Update app.py

app.py

@@ -1,5 +1,5 @@
 # ------------------------------
-# Imports & Dependencies
+# Imports
 # ------------------------------
 from langchain_openai import OpenAIEmbeddings
 from langchain_community.vectorstores import Chroma
@@ -21,7 +21,7 @@ from langchain.tools.retriever import create_retriever_tool
 from datetime import datetime
 
 # ------------------------------
-# Data Definitions
+# Data
 # ------------------------------
 research_texts = [
     "Research Report: Results of a New AI Model Improving Image Recognition Accuracy to 98%",
@@ -40,7 +40,7 @@ development_texts = [
 ]
 
 # ------------------------------
-# Configuration Class
+# Configuration
 # ------------------------------
 class AppConfig:
     def __init__(self):
@@ -52,49 +52,43 @@ class AppConfig:
         self.DOCUMENT_OVERLAP = 50
         self.SEARCH_K = 5
         self.SEARCH_TYPE = "mmr"
-        self.validate_config()
-    
-    def validate_config(self):
+
+    def validate(self):
         if not self.DEEPSEEK_API_KEY:
             st.error("""
-            **Critical Configuration Missing**  
-            🔑 DeepSeek API key not found in environment variables.  
-            Please configure through Hugging Face Space secrets:
-            1. Go to Space Settings → Repository secrets
-            2. Add secret: Name=DEEPSEEK_API_KEY, Value=your_api_key
+            **Configuration Error**  
+            🔑 Missing DeepSeek API key.  
+            Configure through Hugging Face Space secrets:
+            1. Space Settings → Repository secrets
+            2. Add secret: DEEPSEEK_API_KEY=your_key
             3. Rebuild Space
             """)
             st.stop()
 
-config = AppConfig()
-
-# ------------------------------
-# ChromaDB Manager
-# ------------------------------
 class ChromaManager:
-    def __init__(self, research_data: List[str], development_data: List[str]):
+    def __init__(self, config: AppConfig):
         os.makedirs(config.CHROMA_PATH, exist_ok=True)
         self.client = chromadb.PersistentClient(path=config.CHROMA_PATH)
         self.embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
         
-        self.research_collection = self.create_collection(
-            research_data, 
+        self.research_collection = self._create_collection(
+            research_texts, 
             "research_collection",
-            {"category": "research", "version": "1.2"}
+            {"category": "research"}
         )
-        self.dev_collection = self.create_collection(
-            development_data,
+        self.dev_collection = self._create_collection(
+            development_texts,
             "development_collection", 
-            {"category": "development", "version": "1.1"}
+            {"category": "development"}
         )
 
-    def create_collection(self, documents: List[str], name: str, metadata: dict) -> Chroma:
-        text_splitter = RecursiveCharacterTextSplitter(
-            chunk_size=config.DOCUMENT_CHUNK_SIZE,
-            chunk_overlap=config.DOCUMENT_OVERLAP,
-            separators=["\n\n", "\n", "。", " "]
+    def _create_collection(self, documents: List[str], name: str, metadata: dict) -> Chroma:
+        splitter = RecursiveCharacterTextSplitter(
+            chunk_size=300,
+            chunk_overlap=50,
+            separators=["\n\n", "\n", "。"]
         )
-        docs = text_splitter.create_documents(documents)
+        docs = splitter.create_documents(documents)
         return Chroma.from_documents(
             documents=docs,
             embedding=self.embeddings,
@@ -103,383 +97,208 @@ class ChromaManager:
             collection_metadata=metadata
         )
 
-# Initialize Chroma with data
-chroma_manager = ChromaManager(research_texts, development_texts)
-
 # ------------------------------
 # Document Processing
 # ------------------------------
 class DocumentProcessor:
     @staticmethod
-    def deduplicate_documents(docs: List[Any]) -> List[Any]:
+    def deduplicate(docs: List[Any]) -> List[Any]:
         seen = set()
-        unique_docs = []
-        for doc in docs:
-            content_hash = hashlib.md5(doc.page_content.encode()).hexdigest()
-            if content_hash not in seen:
-                unique_docs.append(doc)
-                seen.add(content_hash)
-        return unique_docs
+        return [doc for doc in docs 
+                if not (hashlib.md5(doc.page_content.encode()).hexdigest() in seen 
+                        or seen.add(hashlib.md5(doc.page_content.encode()).hexdigest()))]
 
     @staticmethod
-    def extract_key_points(docs: List[Any]) -> str:
-        key_points = []
+    def extract_keypoints(docs: List[Any]) -> str:
         categories = {
-            "quantum": ["quantum", "qpu", "qubit"],
-            "vision": ["image", "recognition", "vision"],
-            "nlp": ["transformer", "language", "llm"]
+            "quantum": ["quantum", "qubit"],
+            "vision": ["image", "recognition"],
+            "nlp": ["transformer", "language"]
         }
-        
-        for doc in docs:
-            content = doc.page_content.lower()
-            if any(kw in content for kw in categories["quantum"]):
-                key_points.append("- Quantum computing integration showing promising results")
-            if any(kw in content for kw in categories["vision"]):
-                key_points.append("- Computer vision models achieving state-of-the-art accuracy")
-            if any(kw in content for kw in categories["nlp"]):
-                key_points.append("- NLP architectures evolving with memory-augmented transformers")
-        
-        return "\n".join(list(set(key_points)))
+        return "\n".join(sorted({
+            "- " + {
+                "quantum": "Quantum computing breakthroughs",
+                "vision": "Computer vision advancements",
+                "nlp": "NLP architecture innovations"
+            }[cat] 
+            for doc in docs 
+            for cat, kw in categories.items() 
+            if any(k in doc.page_content.lower() for k in kw)
+        }))
 
 # ------------------------------
-# Enhanced Agent Components
+# Workflow
 # ------------------------------
-class EnhancedAgent:
-    def __init__(self):
-        self.session_stats = {
-            "processing_times": [],
-            "doc_counts": [],
-            "error_count": 0
-        }
-    
-    def api_request_with_retry(self, endpoint: str, payload: Dict) -> Dict:
-        headers = {
-            "Authorization": f"Bearer {config.DEEPSEEK_API_KEY}",
-            "Content-Type": "application/json"
-        }
+class AgentWorkflow:
+    def __init__(self, chroma: ChromaManager):
+        self.chroma = chroma
+        self.workflow = StateGraph(AgentState)
         
-        for attempt in range(config.MAX_RETRIES):
-            try:
-                response = requests.post(
-                    endpoint,
-                    headers=headers,
-                    json=payload,
-                    timeout=30,
-                    verify=False
-                )
-                response.raise_for_status()
-                return response.json()
-            except requests.exceptions.HTTPError as e:
-                if e.response.status_code == 429:
-                    delay = config.RETRY_DELAY ** (attempt + 1)
-                    time.sleep(delay)
-                    continue
-                raise
-        raise Exception(f"API request failed after {config.MAX_RETRIES} attempts")
+        # Define nodes
+        self.workflow.add_node("agent", self.agent)
+        self.workflow.add_node("retrieve", ToolNode([
+            create_retriever_tool(
+                chroma.research_collection.as_retriever(),
+                "research_tool",
+                "Search research documents"
+            ),
+            create_retriever_tool(
+                chroma.dev_collection.as_retriever(),
+                "dev_tool",
+                "Search development updates"
+            )
+        ]))
+        self.workflow.add_node("generate", self.generate)
+        self.workflow.add_node("rewrite", self.rewrite)
 
-# ------------------------------
-# Workflow Configuration
-# ------------------------------
-class AgentState(TypedDict):
-    messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]
-
-def agent(state: AgentState):
-    print("---CALL AGENT---")
-    messages = state["messages"]
-    user_message = messages[0].content if not isinstance(messages[0], tuple) else messages[0][1]
-
-    prompt = f"""Given this user question: "{user_message}"
-If about research/academic topics, respond EXACTLY:
-SEARCH_RESEARCH: <search terms>
-If about development status, respond EXACTLY:
-SEARCH_DEV: <search terms>
-Otherwise, answer directly."""
-    
-    headers = {
-        "Accept": "application/json",
-        "Authorization": f"Bearer {config.DEEPSEEK_API_KEY}",
-        "Content-Type": "application/json"
-    }
-    
-    data = {
-        "model": "deepseek-chat",
-        "messages": [{"role": "user", "content": prompt}],
-        "temperature": 0.7,
-        "max_tokens": 1024
-    }
-    
-    try:
-        response = requests.post(
-            "https://api.deepseek.com/v1/chat/completions",
-            headers=headers,
-            json=data,
-            verify=False,
-            timeout=30
+        # Define edges
+        self.workflow.set_entry_point("agent")
+        self.workflow.add_conditional_edges(
+            "agent",
+            self._tools_condition,
+            {"retrieve": "retrieve", "end": END}
         )
-        response.raise_for_status()
-        response_text = response.json()['choices'][0]['message']['content']
-        
-        if "SEARCH_RESEARCH:" in response_text:
-            query = response_text.split("SEARCH_RESEARCH:")[1].strip()
-            results = chroma_manager.research_collection.as_retriever().invoke(query)
-            return {"messages": [AIMessage(content=f'Action: research_db_tool\n{{"query": "{query}"}}\n\nResults: {str(results)}')]}
-        
-        elif "SEARCH_DEV:" in response_text:
-            query = response_text.split("SEARCH_DEV:")[1].strip()
-            results = chroma_manager.dev_collection.as_retriever().invoke(query)
-            return {"messages": [AIMessage(content=f'Action: development_db_tool\n{{"query": "{query}"}}\n\nResults: {str(results)}')]}
+        self.workflow.add_conditional_edges(
+            "retrieve",
+            self._grade_documents,
+            {"generate": "generate", "rewrite": "rewrite"}
+        )
+        self.workflow.add_edge("generate", END)
+        self.workflow.add_edge("rewrite", "agent")
         
-        return {"messages": [AIMessage(content=response_text)]}
-    except Exception as e:
-        error_msg = f"API Error: {str(e)}"
-        if "Insufficient Balance" in str(e):
-            error_msg += "\n\nPlease check your DeepSeek API account balance."
-        return {"messages": [AIMessage(content=error_msg)]}
+        self.app = self.workflow.compile()
 
-def simple_grade_documents(state: AgentState):
-    messages = state["messages"]
-    last_message = messages[-1]
-    return "generate" if "Results: [Document" in last_message.content else "rewrite"
+    class AgentState(TypedDict):
+        messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]
 
-def generate(state: AgentState):
-    messages = state["messages"]
-    question = messages[0].content
-    last_message = messages[-1]
+    def agent(self, state: AgentState):
+        try:
+            messages = state["messages"]
+            query = messages[-1].content if isinstance(messages[-1], HumanMessage) else messages[-1]['content']
+            
+            response = requests.post(
+                "https://api.deepseek.com/v1/chat/completions",
+                headers={"Authorization": f"Bearer {config.DEEPSEEK_API_KEY}"},
+                json={
+                    "model": "deepseek-chat",
+                    "messages": [{
+                        "role": "user",
+                        "content": f"""Analyze this query: "{query}"
+                        Respond EXACTLY as:
+                        - SEARCH_RESEARCH: <terms> (for research topics)
+                        - SEARCH_DEV: <terms> (for development updates)
+                        - DIRECT: <answer> (otherwise)"""
+                    }]
+                }
+            ).json()
+            
+            content = response['choices'][0]['message']['content']
+            if "SEARCH_RESEARCH:" in content:
+                terms = content.split("SEARCH_RESEARCH:")[1].strip()
+                results = self.chroma.research_collection.similarity_search(terms)
+                return {"messages": [AIMessage(content=f"Research Results: {str(results)}")]}
+            elif "SEARCH_DEV:" in content:
+                terms = content.split("SEARCH_DEV:")[1].strip()
+                results = self.chroma.dev_collection.similarity_search(terms)
+                return {"messages": [AIMessage(content=f"Development Results: {str(results)}")]}
+            return {"messages": [AIMessage(content=content)]}
+            
+        except Exception as e:
+            return {"messages": [AIMessage(content=f"Error: {str(e)}")]}
 
-    docs_content = []
-    if "Results: [" in last_message.content:
-        docs_str = last_message.content.split("Results: ")[1]
-        docs_content = eval(docs_str)
-    
-    processed_info = DocumentProcessor.extract_key_points(
-        DocumentProcessor.deduplicate_documents(docs_content)
-    )
-    
-    prompt = f"""Generate structured research summary:
-    Key Information:
-    {processed_info}
-    Include:
-    1. Section headings
-    2. Bullet points
-    3. Significance
-    4. Applications"""
-    
-    try:
+    def generate(self, state: AgentState):
+        docs = eval(state["messages"][-1].content.split("Results: ")[1])
+        processed = "\n".join([d.page_content[:200] for d in DocumentProcessor.deduplicate(docs)])
+        
         response = requests.post(
             "https://api.deepseek.com/v1/chat/completions",
-            headers={
-                "Authorization": f"Bearer {config.DEEPSEEK_API_KEY}",
-                "Content-Type": "application/json"
-            },
+            headers={"Authorization": f"Bearer {config.DEEPSEEK_API_KEY}"},
             json={
                 "model": "deepseek-chat",
-                "messages": [{"role": "user", "content": prompt}],
-                "temperature": 0.7,
-                "max_tokens": 1024
-            },
-            timeout=30
-        )
-        response.raise_for_status()
-        return {"messages": [AIMessage(content=response.json()['choices'][0]['message']['content'])]}
-    except Exception as e:
-        return {"messages": [AIMessage(content=f"Generation Error: {str(e)}")]}
+                "messages": [{
+                    "role": "user",
+                    "content": f"Summarize these findings:\n{processed}"
+                }]
+            }
+        ).json()
+        
+        return {"messages": [AIMessage(content=response['choices'][0]['message']['content'])]}
 
-def rewrite(state: AgentState):
-    messages = state["messages"]
-    original_question = messages[0].content
-    
-    try:
+    def rewrite(self, state: AgentState):
+        original = state["messages"][0].content
         response = requests.post(
             "https://api.deepseek.com/v1/chat/completions",
-            headers={
-                "Authorization": f"Bearer {config.DEEPSEEK_API_KEY}",
-                "Content-Type": "application/json"
-            },
+            headers={"Authorization": f"Bearer {config.DEEPSEEK_API_KEY}"},
             json={
                 "model": "deepseek-chat",
                 "messages": [{
-                    "role": "user", 
-                    "content": f"Rewrite for clarity: {original_question}"
-                }],
-                "temperature": 0.7,
-                "max_tokens": 1024
-            },
-            timeout=30
-        )
-        response.raise_for_status()
-        return {"messages": [AIMessage(content=response.json()['choices'][0]['message']['content'])]}
-    except Exception as e:
-        return {"messages": [AIMessage(content=f"Rewrite Error: {str(e)}")]}
+                    "role": "user",
+                    "content": f"Rephrase this query: {original}"
+                }]
+            }
+        ).json()
+        return {"messages": [AIMessage(content=response['choices'][0]['message']['content'])]}
 
-tools_pattern = re.compile(r"Action: .*")
+    def _tools_condition(self, state: AgentState):
+        return "retrieve" if "Results:" in state["messages"][-1].content else "end"
 
-def custom_tools_condition(state: AgentState):
-    content = state["messages"][-1].content
-    return "tools" if tools_pattern.match(content) else END
+    def _grade_documents(self, state: AgentState):
+        return "generate" if len(eval(state["messages"][-1].content.split("Results: ")[1])) > 0 else "rewrite"
 
 # ------------------------------
-# Workflow Graph Setup
+# Streamlit App
 # ------------------------------
-workflow = StateGraph(AgentState)
-workflow.add_node("agent", agent)
-workflow.add_node("retrieve", ToolNode([
-    create_retriever_tool(
-        chroma_manager.research_collection.as_retriever(),
-        "research_db_tool",
-        "Search research database"
-    ),
-    create_retriever_tool(
-        chroma_manager.dev_collection.as_retriever(),
-        "development_db_tool",
-        "Search development database"
-    )
-]))
-workflow.add_node("rewrite", rewrite)
-workflow.add_node("generate", generate)
-
-workflow.set_entry_point("agent")
-workflow.add_conditional_edges("agent", custom_tools_condition, {"tools": "retrieve", END: END})
-workflow.add_conditional_edges("retrieve", simple_grade_documents, {"generate": "generate", "rewrite": "rewrite"})
-workflow.add_edge("generate", END)
-workflow.add_edge("rewrite", "agent")
-
-app = workflow.compile()
-
-# ------------------------------
-# Streamlit UI
-# ------------------------------
-class UITheme:
-    primary_color = "#2E86C1"
-    secondary_color = "#28B463"
-    background_color = "#1A1A1A"
-    text_color = "#EAECEE"
-    
-    @classmethod
-    def apply(cls):
-        st.markdown(f"""
-        <style>
-        .stApp {{ background-color: {cls.background_color}; color: {cls.text_color}; }}
-        .stTextArea textarea {{ 
-            background-color: #2D2D2D !important; 
-            color: {cls.text_color} !important;
-            border: 1px solid {cls.primary_color};
-        }}
-        .stButton > button {{
-            background-color: {cls.primary_color};
-            color: white;
-            border: none;
-            padding: 12px 28px;
-            border-radius: 6px;
-            transition: all 0.3s ease;
-            font-weight: 500;
-        }}
-        .stButton > button:hover {{
-            background-color: {cls.secondary_color};
-            transform: translateY(-1px);
-            box-shadow: 0 4px 12px rgba(0,0,0,0.2);
-        }}
-        .data-box {{
-            background-color: #2D2D2D;
-            border-left: 4px solid {cls.primary_color};
-            padding: 18px;
-            margin: 14px 0;
-            border-radius: 8px;
-            box-shadow: 0 2px 8px rgba(0,0,0,0.15);
-        }}
-        </style>
-        """, unsafe_allow_html=True)
+def apply_theme():
+    st.markdown("""
+    <style>
+    .stApp { background: #1a1a1a; color: white; }
+    .stTextArea textarea { background: #2d2d2d !important; color: white !important; }
+    .stButton>button { background: #2E86C1; transition: 0.3s; }
+    .stButton>button:hover { background: #1B4F72; transform: scale(1.02); }
+    .data-box { background: #2d2d2d; border-left: 4px solid #2E86C1; padding: 15px; margin: 10px 0; }
+    </style>
+    """, unsafe_allow_html=True)
 
 def main():
-    UITheme.apply()
-    
-    st.set_page_config(
-        page_title="AI Research Assistant Pro",
-        layout="wide",
-        initial_sidebar_state="expanded",
-        menu_items={
-            'Get Help': 'https://example.com/docs',
-            'Report a bug': 'https://example.com/issues',
-            'About': "v2.1 | Enhanced Research Assistant"
-        }
-    )
+    apply_theme()
     
     with st.sidebar:
-        st.header("📂 Knowledge Bases")
-        with st.expander("Research Database", expanded=True):
+        st.header("📚 Databases")
+        with st.expander("Research", expanded=True):
             for text in research_texts:
-                st.markdown(f'<div class="data-box research-box">{text}</div>', unsafe_allow_html=True)
-        
-        with st.expander("Development Database"):
+                st.markdown(f'<div class="data-box">{text}</div>', unsafe_allow_html=True)
+        with st.expander("Development"):
             for text in development_texts:
-                st.markdown(f'<div class="data-box dev-box">{text}</div>', unsafe_allow_html=True)
-    
-    st.title("🔬 AI Research Assistant Pro")
-    st.markdown("---")
-    
-    query = st.text_area(
-        "Research Query Input",
-        height=120,
-        placeholder="Enter your research question...",
-        help="Be specific about domains for better results"
-    )
+                st.markdown(f'<div class="data-box">{text}</div>', unsafe_allow_html=True)
+
+    st.title("🔍 AI Research Assistant")
+    query = st.text_area("Enter your query:", height=100)
     
-    col1, col2 = st.columns([1, 2])
-    with col1:
-        if st.button("🚀 Analyze Documents", use_container_width=True):
-            if not query:
-                st.warning("⚠️ Please enter a research question")
-                return
+    if st.button("Analyze"):
+        with st.spinner("Processing..."):
+            workflow = AgentWorkflow(chroma_manager)
+            results = workflow.app.invoke({"messages": [HumanMessage(content=query)]})
+            
+            with st.expander("Processing Details", expanded=True):
+                st.write("### Raw Results", results)
                 
-            with st.status("Processing Workflow...", expanded=True) as status:
-                try:
-                    start_time = time.time()
-                    events = process_question(query, app, {"configurable": {"thread_id": "1"}})
-                    
-                    processed_data = []
-                    for event in events:
-                        if 'agent' in event:
-                            content = event['agent']['messages'][0].content
-                            if "Results:" in content:
-                                docs = eval(content.split("Results: ")[1])
-                                unique_docs = DocumentProcessor.deduplicate_documents(docs)
-                                key_points = DocumentProcessor.extract_key_points(unique_docs)
-                                processed_data.append(key_points)
-                                
-                                with st.expander("📄 Retrieved Documents", expanded=False):
-                                    st.info(f"Found {len(unique_docs)} unique documents")
-                                    st.write(docs)
-                        
-                        elif 'generate' in event:
-                            final_answer = event['generate']['messages'][0].content
-                            status.update(label="✅ Analysis Complete", state="complete")
-                            st.markdown("## 📝 Research Summary")
-                            st.markdown(final_answer)
-                    
-                    st.caption(f"⏱️ Processed in {time.time()-start_time:.2f}s | {len(processed_data)} clusters")
-                
-                except Exception as e:
-                    status.update(label="❌ Processing Failed", state="error")
-                    st.error(f"**Error:** {str(e)}\n\nCheck API key and network connection")
-                    with open("error_log.txt", "a") as f:
-                        f.write(f"{datetime.now()} | {str(e)}\n")
-
-    with col2:
-        st.markdown("""
-        ## 📘 Usage Guide
-        **1. Query Formulation**  
-        - Specify domains (e.g., "quantum NLP")  
-        - Include timeframes for recent advances  
-        
-        **2. Results Interpretation**  
-        - Expand sections for source documents  
-        - Key points show technical breakthroughs  
-        - Summary includes commercial implications  
-        
-        **3. Advanced Features**  
-        - Use keyboard shortcuts for efficiency  
-        - Click documents for raw context  
-        - Export via screenshot/PDF  
-        """)
+            st.success("### Final Answer")
+            st.markdown(results['messages'][-1].content)
 
+# ------------------------------
+# Initialization
+# ------------------------------
 if __name__ == "__main__":
-    main()
+    st.set_page_config(
+        page_title="AI Research Assistant",
+        layout="wide",
+        initial_sidebar_state="expanded"
+    )
+    
+    try:
+        config = AppConfig()
+        config.validate()
+        chroma_manager = ChromaManager(config)
+        main()
+    except Exception as e:
+        st.error(f"Initialization failed: {str(e)}")