Update app.py

app.py

@@ -3,33 +3,30 @@
 # ------------------------------
 from langchain_openai import OpenAIEmbeddings
 from langchain_community.vectorstores import Chroma
-from langchain_core.messages import HumanMessage, AIMessage
+from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_core.documents import Document
 from langgraph.graph import END, StateGraph
-from langgraph.graph.message import add_messages  # <-- FIXED IMPORT
+from langgraph.prebuilt import ToolNode
+from langgraph.graph.message import add_messages
 from typing_extensions import TypedDict, Annotated
 from typing import Sequence, Dict, List, Optional, Any
 import chromadb
-import numpy as np
+import re
 import os
 import streamlit as st
 import requests
 import hashlib
-import re
+import json
 import time
+from langchain.tools.retriever import create_retriever_tool
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime
-from sklearn.metrics.pairwise import cosine_similarity
 
 # ------------------------------
-# State Schema Definition (Fixed)
+# State Schema Definition
 # ------------------------------
 class AgentState(TypedDict):
-    messages: Annotated[
-        Sequence[AIMessage | HumanMessage], 
-        add_messages  # <-- NOW PROPERLY DEFINED
-    ]
+    messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]
     context: Dict[str, Any]
     metadata: Dict[str, Any]
 
@@ -43,62 +40,37 @@ class ResearchConfig:
     CHUNK_OVERLAP = 64
     MAX_CONCURRENT_REQUESTS = 5
     EMBEDDING_DIMENSIONS = 1536
-    
     DOCUMENT_MAP = {
-        "CV-Transformer Model": {
-            "title": "Hybrid CV-Transformer Architecture",
-            "content": """
-            Combines CNN feature extraction with transformer attention mechanisms.
-            Key equation: $f(x) = \text{Softmax}(\frac{QK^T}{\sqrt{d_k}})V$
-            Achieves 98.2% accuracy on ImageNet-1k with 42ms inference speed
-            """
-        },
-        "Quantum ML": {
-            "title": "Quantum Machine Learning",
-            "content": """
-            Quantum-enhanced optimization techniques for ML models.
-            $\theta_{t+1} = \theta_t - \eta \nabla_\theta \mathcal{L}(\theta_t)$
-            100x speedup on optimization tasks with 58% energy reduction
-            """
-        }
+        "Research Report: Results of a New AI Model Improving Image Recognition Accuracy to 98%": 
+            "CV-Transformer Hybrid Architecture",
+        "Academic Paper Summary: Why Transformers Became the Mainstream Architecture in Natural Language Processing": 
+            "Transformer Architecture Analysis",
+        "Latest Trends in Machine Learning Methods Using Quantum Computing": 
+            "Quantum ML Frontiers"
     }
-    
-    ANALYSIS_TEMPLATE = """Analyze these technical documents:
+    ANALYSIS_TEMPLATE = """Analyze these technical documents with scientific rigor:
 {context}
 
-Respond in MARKDOWN with:
-1. **Key Innovations** (mathematical formulations)
-2. **Methodologies** (algorithms & architectures)
-3. **Empirical Results** (comparative metrics)
-4. **Applications** (industry use cases)
-5. **Limitations** (theoretical boundaries)
+Respond with:
+1. Key Technical Contributions (bullet points)
+2. Novel Methodologies
+3. Empirical Results (with metrics)
+4. Potential Applications
+5. Limitations & Future Directions
 
-Include LaTeX equations where applicable."""
-
-# Check for Chroma migration
-if os.path.exists(ResearchConfig.CHROMA_PATH):
-    st.warning("""
-    **ChromDB Migration Required**
-    Existing Chroma database detected. Run these commands:
-    
-    ```bash
-    pip install chroma-migrate
-    chroma-migrate
-    ```
-    
-    Then restart the application.
-    """)
-    st.stop()
+Format: Markdown with LaTeX mathematical notation where applicable
+"""
 
+# Validation
 if not ResearchConfig.DEEPSEEK_API_KEY:
-    st.error("""**Configuration Required**  
-    1. Get DeepSeek API key: [platform.deepseek.com](https://platform.deepseek.com/)  
-    2. Set secret: `DEEPSEEK_API_KEY`  
+    st.error("""**Research Portal Configuration Required**  
+    1. Obtain DeepSeek API key: [platform.deepseek.com](https://platform.deepseek.com/)  
+    2. Configure secret: `DEEPSEEK_API_KEY` in Space settings  
     3. Rebuild deployment""")
     st.stop()
 
 # ------------------------------
-# ChromaDB Document Manager
+# Quantum Document Processing
 # ------------------------------
 class QuantumDocumentManager:
     def __init__(self):
@@ -107,82 +79,97 @@ class QuantumDocumentManager:
             model="text-embedding-3-large",
             dimensions=ResearchConfig.EMBEDDING_DIMENSIONS
         )
-
-    def create_collection(self, document_map: Dict[str, Dict[str, str]], collection_name: str) -> Chroma:
+        
+    def create_collection(self, documents: List[str], collection_name: str) -> Chroma:
         splitter = RecursiveCharacterTextSplitter(
             chunk_size=ResearchConfig.CHUNK_SIZE,
             chunk_overlap=ResearchConfig.CHUNK_OVERLAP,
             separators=["\n\n", "\n", "|||"]
         )
-        
-        docs = []
-        for key, data in document_map.items():
-            chunks = splitter.split_text(data["content"])
-            for chunk in chunks:
-                docs.append(Document(
-                    page_content=chunk,
-                    metadata={
-                        "title": data["title"],
-                        "source": collection_name,
-                        "hash": hashlib.sha256(chunk.encode()).hexdigest()[:16]
-                    }
-                ))
-        
+        docs = splitter.create_documents(documents)
         return Chroma.from_documents(
             documents=docs,
             embedding=self.embeddings,
-            collection_name=collection_name,
             client=self.client,
+            collection_name=collection_name,
             ids=[self._document_id(doc.page_content) for doc in docs]
         )
     
     def _document_id(self, content: str) -> str:
         return f"{hashlib.sha256(content.encode()).hexdigest()[:16]}-{int(time.time())}"
 
-# Initialize document system
+# Initialize document collections
 qdm = QuantumDocumentManager()
-research_docs = qdm.create_collection(ResearchConfig.DOCUMENT_MAP, "research_papers")
-
+research_docs = qdm.create_collection([
+    "Research Report: Results of a New AI Model Improving Image Recognition Accuracy to 98%",
+    "Academic Paper Summary: Why Transformers Became the Mainstream Architecture in Natural Language Processing",
+    "Latest Trends in Machine Learning Methods Using Quantum Computing"
+], "research")
+
+development_docs = qdm.create_collection([
+    "Project A: UI Design Completed, API Integration in Progress",
+    "Project B: Testing New Feature X, Bug Fixes Needed",
+    "Product Y: In the Performance Optimization Stage Before Release"
+], "development")
 
 # ------------------------------
-# Retrieval System
+# Advanced Retrieval System
 # ------------------------------
 class ResearchRetriever:
     def __init__(self):
-        self.retriever = research_docs.as_retriever(
-            search_type="mmr",
-            search_kwargs={
-                'k': 4,
-                'fetch_k': 20,
-                'lambda_mult': 0.85
-            }
-        )
+        self.retrievers = {
+            "research": research_docs.as_retriever(
+                search_type="mmr",
+                search_kwargs={
+                    'k': 4,
+                    'fetch_k': 20,
+                    'lambda_mult': 0.85
+                }
+            ),
+            "development": development_docs.as_retriever(
+                search_type="similarity",
+                search_kwargs={'k': 3}
+            )
+        }
     
-    def retrieve(self, query: str) -> List[Document]:
+    def retrieve(self, query: str, domain: str) -> List[Any]:
         try:
-            docs = self.retriever.invoke(query)
-            if len(docs) < 1:
-                raise ValueError("No relevant documents found")
-            return docs
-        except Exception as e:
-            st.error(f"Retrieval Error: {str(e)}")
+            return self.retrievers[domain].invoke(query)
+        except KeyError:
             return []
 
+retriever = ResearchRetriever()
+
 # ------------------------------
-# Analysis Processor
+# Cognitive Processing Unit
 # ------------------------------
 class CognitiveProcessor:
     def __init__(self):
         self.executor = ThreadPoolExecutor(max_workers=ResearchConfig.MAX_CONCURRENT_REQUESTS)
+        self.session_id = hashlib.sha256(datetime.now().isoformat().encode()).hexdigest()[:12]
     
     def process_query(self, prompt: str) -> Dict:
-        futures = [self.executor.submit(self._api_request, prompt) for _ in range(3)]
-        return self._best_result([f.result() for f in as_completed(futures)])
-
-    def _api_request(self, prompt: str) -> Dict:
+        futures = []
+        for _ in range(3):  # Triple redundancy
+            futures.append(self.executor.submit(
+                self._execute_api_request,
+                prompt
+            ))
+        
+        results = []
+        for future in as_completed(futures):
+            try:
+                results.append(future.result())
+            except Exception as e:
+                st.error(f"Processing Error: {str(e)}")
+        
+        return self._consensus_check(results)
+    
+    def _execute_api_request(self, prompt: str) -> Dict:
         headers = {
             "Authorization": f"Bearer {ResearchConfig.DEEPSEEK_API_KEY}",
-            "Content-Type": "application/json"
+            "Content-Type": "application/json",
+            "X-Research-Session": self.session_id
         }
         
         try:
@@ -203,35 +190,30 @@ class CognitiveProcessor:
             )
             response.raise_for_status()
             return response.json()
-        except Exception as e:
+        except requests.exceptions.RequestException as e:
             return {"error": str(e)}
     
-    def _best_result(self, results: List[Dict]) -> Dict:
+    def _consensus_check(self, results: List[Dict]) -> Dict:
         valid = [r for r in results if "error" not in r]
         if not valid:
             return {"error": "All API requests failed"}
-        
-        # Select response with most technical content
-        contents = [r.get('choices', [{}])[0].get('message', {}).get('content', '') for r in valid]
-        tech_scores = [len(re.findall(r"\$.*?\$", c)) for c in contents]
-        return valid[np.argmax(tech_scores)]
+        return max(valid, key=lambda x: len(x.get('choices', [{}])[0].get('message', {}).get('content', '')))
 
 # ------------------------------
-# Workflow Engine
+# Research Workflow Engine
 # ------------------------------
 class ResearchWorkflow:
     def __init__(self):
-        self.retriever = ResearchRetriever()
         self.processor = CognitiveProcessor()
         self.workflow = StateGraph(AgentState)
         self._build_workflow()
 
     def _build_workflow(self):
-        self.workflow.add_node("ingest", self.ingest)
-        self.workflow.add_node("retrieve", self.retrieve)
-        self.workflow.add_node("analyze", self.analyze)
-        self.workflow.add_node("validate", self.validate)
-        self.workflow.add_node("refine", self.refine)
+        self.workflow.add_node("ingest", self.ingest_query)
+        self.workflow.add_node("retrieve", self.retrieve_documents)
+        self.workflow.add_node("analyze", self.analyze_content)
+        self.workflow.add_node("validate", self.validate_output)
+        self.workflow.add_node("refine", self.refine_results)
 
         self.workflow.set_entry_point("ingest")
         self.workflow.add_edge("ingest", "retrieve")
@@ -246,93 +228,83 @@ class ResearchWorkflow:
 
         self.app = self.workflow.compile()
 
-    def ingest(self, state: AgentState) -> Dict:
+    def ingest_query(self, state: AgentState) -> Dict:
         try:
             query = state["messages"][-1].content
             return {
-                "messages": [AIMessage(content="Query ingested")],
-                "context": {"query": query},
+                "messages": [AIMessage(content="Query ingested successfully")],
+                "context": {"raw_query": query},
                 "metadata": {"timestamp": datetime.now().isoformat()}
             }
         except Exception as e:
             return self._error_state(f"Ingestion Error: {str(e)}")
 
-    def retrieve(self, state: AgentState) -> Dict:
+    def retrieve_documents(self, state: AgentState) -> Dict:
         try:
-            docs = self.retriever.retrieve(state["context"]["query"])
+            query = state["context"]["raw_query"]
+            docs = retriever.retrieve(query, "research")
             return {
-                "messages": [AIMessage(content=f"Found {len(docs)} relevant papers")],
-                "context": {"docs": docs}
+                "messages": [AIMessage(content=f"Retrieved {len(docs)} documents")],
+                "context": {
+                    "documents": docs,
+                    "retrieval_time": time.time()
+                }
             }
         except Exception as e:
             return self._error_state(f"Retrieval Error: {str(e)}")
 
-    def analyze(self, state: AgentState) -> Dict:
+    def analyze_content(self, state: AgentState) -> Dict:
         try:
-            context = "\n\n".join([
-                f"### {doc.metadata['title']}\n{doc.page_content}" 
-                for doc in state["context"]["docs"]
-            ])
-            prompt = ResearchConfig.ANALYSIS_TEMPLATE.format(context=context)
+            docs = "\n\n".join([d.page_content for d in state["context"]["documents"]])
+            prompt = ResearchConfig.ANALYSIS_TEMPLATE.format(context=docs)
             response = self.processor.process_query(prompt)
             
             if "error" in response:
-                raise RuntimeError(response["error"])
-                
-            content = response['choices'][0]['message']['content']
-            self._validate_analysis(content)
+                return self._error_state(response["error"])
             
-            return {"messages": [AIMessage(content=content)]}
+            return {
+                "messages": [AIMessage(content=response['choices'][0]['message']['content'])],
+                "context": {"analysis": response}
+            }
         except Exception as e:
             return self._error_state(f"Analysis Error: {str(e)}")
 
-    def validate(self, state: AgentState) -> Dict:
-        validation_prompt = f"""Validate this technical analysis:
-{state["messages"][-1].content}
-
-Check for:
-1. Mathematical accuracy
-2. Technical depth
-3. Logical consistency
-
-Respond with 'VALID' or 'INVALID'"""
+    def validate_output(self, state: AgentState) -> Dict:
+        analysis = state["messages"][-1].content
+        validation_prompt = f"""Validate research analysis:
+        {analysis}
+        
+        Check for:
+        1. Technical accuracy
+        2. Citation support
+        3. Logical consistency
+        4. Methodological soundness
+        
+        Respond with 'VALID' or 'INVALID'"""
         
         response = self.processor.process_query(validation_prompt)
-        valid = "VALID" in response.get('choices', [{}])[0].get('message', {}).get('content', '')
         return {
-            "messages": [AIMessage(content=f"{state['messages'][-1].content}\n\nValidation: {'✅ Valid' if valid else '❌ Invalid'}")],
-            "context": {"valid": valid}
+            "messages": [AIMessage(content=analysis + f"\n\nValidation: {response.get('choices', [{}])[0].get('message', {}).get('content', '')}")]
         }
 
-    def refine(self, state: AgentState) -> Dict:
-        refinement_prompt = f"""Improve this analysis:
-{state["messages"][-1].content}
-
-Focus on:
-1. Mathematical precision
-2. Technical terminology
-3. Empirical references"""
+    def refine_results(self, state: AgentState) -> Dict:
+        refinement_prompt = f"""Refine this analysis:
+        {state["messages"][-1].content}
+        
+        Improve:
+        1. Technical precision
+        2. Empirical grounding
+        3. Theoretical coherence"""
         
         response = self.processor.process_query(refinement_prompt)
-        return {"messages": [AIMessage(content=response['choices'][0]['message']['content'])]}
+        return {
+            "messages": [AIMessage(content=response.get('choices', [{}])[0].get('message', {}).get('content', ''))],
+            "context": state["context"]
+        }
 
     def _quality_check(self, state: AgentState) -> str:
-        return "valid" if state.get("context", {}).get("valid", False) else "invalid"
-
-    def _validate_analysis(self, content: str):
-        required_sections = [
-            "Key Innovations",
-            "Methodologies", 
-            "Empirical Results",
-            "Applications",
-            "Limitations"
-        ]
-        missing = [s for s in required_sections if f"## {s}" not in content]
-        if missing:
-            raise ValueError(f"Missing sections: {', '.join(missing)}")
-            
-        if not re.search(r"\$.*?\$", content):
-            raise ValueError("Analysis lacks mathematical notation")
+        content = state["messages"][-1].content
+        return "valid" if "VALID" in content else "invalid"
 
     def _error_state(self, message: str) -> Dict:
         return {
@@ -342,22 +314,22 @@ Focus on:
         }
 
 # ------------------------------
-# Streamlit Interface
+# Research Interface
 # ------------------------------
 class ResearchInterface:
     def __init__(self):
         self.workflow = ResearchWorkflow()
-        self._initialize()
+        self._initialize_interface()
 
-    def _initialize(self):
+    def _initialize_interface(self):
         st.set_page_config(
-            page_title="AI Research Assistant",
+            page_title="NeuroResearch AI",
             layout="wide",
             initial_sidebar_state="expanded"
         )
         self._inject_styles()
         self._build_sidebar()
-        self._build_main()
+        self._build_main_interface()
 
     def _inject_styles(self):
         st.markdown("""
@@ -366,72 +338,115 @@ class ResearchInterface:
             --primary: #2ecc71;
             --secondary: #3498db;
             --background: #0a0a0a;
+            --text: #ecf0f1;
         }
+        
         .stApp {
             background: var(--background);
-            color: white;
+            color: var(--text);
+            font-family: 'Roboto', sans-serif;
         }
+        
         .stTextArea textarea {
             background: #1a1a1a !important;
-            border: 2px solid var(--secondary) !important;
+            color: var(--text) !important;
+            border: 2px solid var(--secondary);
+            border-radius: 8px;
+            padding: 1rem;
         }
-        code {
-            color: var(--primary);
-            background: #002200;
-            padding: 2px 4px;
+        
+        .stButton>button {
+            background: linear-gradient(135deg, var(--primary), var(--secondary));
+            border: none;
+            border-radius: 8px;
+            padding: 1rem 2rem;
+            transition: all 0.3s;
+        }
+        
+        .stButton>button:hover {
+            transform: translateY(-2px);
+            box-shadow: 0 4px 12px rgba(46, 204, 113, 0.3);
+        }
+        
+        .stExpander {
+            background: #1a1a1a;
+            border: 1px solid #2a2a2a;
+            border-radius: 8px;
+            margin: 1rem 0;
         }
         </style>
         """, unsafe_allow_html=True)
 
     def _build_sidebar(self):
         with st.sidebar:
-            st.title("🔬 Research Corpus")
-            for key, data in ResearchConfig.DOCUMENT_MAP.items():
-                with st.expander(data["title"]):
-                    st.markdown(f"```latex\n{data['content']}\n```")
-            st.metric("Vector DB Size", len(research_docs.get()['ids']))
-
-    def _build_main(self):
-        st.title("🧠 AI Research Analyst")
-        query = st.text_area("Research Query:", height=150,
-                           placeholder="Enter technical question...")
+            st.title("🔍 Research Database")
+            st.subheader("Technical Papers")
+            for title, short in ResearchConfig.DOCUMENT_MAP.items():
+                with st.expander(short):
+                    st.markdown(f"```\n{title}\n```")
+            
+            st.subheader("Analysis Metrics")
+            st.metric("Vector Collections", 2)
+            st.metric("Embedding Dimensions", ResearchConfig.EMBEDDING_DIMENSIONS)
+
+    def _build_main_interface(self):
+        st.title("🧠 NeuroResearch AI")
+        query = st.text_area("Research Query:", height=200,
+                           placeholder="Enter technical research question...")
         
-        if st.button("Analyze", type="primary"):
+        if st.button("Execute Analysis", type="primary"):
             self._execute_analysis(query)
 
     def _execute_analysis(self, query: str):
         try:
-            with st.spinner("Analyzing research corpus..."):
-                result = self.workflow.app.invoke(
-                    {"messages": [HumanMessage(content=query)]}
+            with st.spinner("Initializing Quantum Analysis..."):
+                results = self.workflow.app.stream(
+                    {"messages": [HumanMessage(content=query)], "context": {}, "metadata": {}}
                 )
                 
-                if result.get("context", {}).get("error"):
-                    self._show_error(result["context"]["error"])
-                else:
-                    self._display_result(result)
+                for event in results:
+                    self._render_event(event)
+                
+                st.success("✅ Analysis Completed Successfully")
         except Exception as e:
-            self._show_error(str(e))
-
-    def _display_result(self, result):
-        with st.expander("Technical Report", expanded=True):
-            st.markdown(result["messages"][-1].content)
-            
-        with st.expander("Source Excerpts", expanded=False):
-            for doc in result["context"].get("docs", []):
-                st.markdown(f"**{doc.metadata['title']}**")
-                st.code(doc.page_content, language='latex')
-
-    def _show_error(self, message):
-        st.error(f"""
-        ⚠️ Analysis Failed
-        {message}
-        
-        Mitigation Steps:
-        1. Simplify query complexity
-        2. Check document connections
-        3. Verify technical terms
-        """)
+            st.error(f"""**Analysis Failed**
+            {str(e)}
+            Potential issues:
+            - Complex query structure
+            - Document correlation failure
+            - Temporal processing constraints""")
+
+    def _render_event(self, event: Dict):
+        if 'ingest' in event:
+            with st.container():
+                st.success("✅ Query Ingested")
+                
+        elif 'retrieve' in event:
+            with st.container():
+                docs = event['retrieve']['context']['documents']
+                st.info(f"📚 Retrieved {len(docs)} documents")
+                with st.expander("View Retrieved Documents", expanded=False):
+                    for i, doc in enumerate(docs, 1):
+                        st.markdown(f"**Document {i}**")
+                        st.code(doc.page_content, language='text')
+                        
+        elif 'analyze' in event:
+            with st.container():
+                content = event['analyze']['messages'][0].content
+                with st.expander("Technical Analysis Report", expanded=True):
+                    st.markdown(content)
+                    
+        elif 'validate' in event:
+            with st.container():
+                content = event['validate']['messages'][0].content
+                if "VALID" in content:
+                    st.success("✅ Validation Passed")
+                    with st.expander("View Validated Analysis", expanded=True):
+                        st.markdown(content.split("Validation:")[0])
+                else:
+                    st.warning("⚠️ Validation Issues Detected")
+                    with st.expander("View Validation Details", expanded=True):
+                        st.markdown(content)
 
 if __name__ == "__main__":
     ResearchInterface()