Update app.py

app.py

@@ -1,39 +1,28 @@
-# -----------------------------------------------------
-# Imports & Initial Configuration
-# -----------------------------------------------------
-import streamlit as st
-
-# IMPORTANT: Must be the first Streamlit command
-st.set_page_config(page_title="NeuroResearch AI", layout="wide", initial_sidebar_state="expanded")
-
+# ------------------------------
+# Imports & Dependencies
+# ------------------------------
 from langchain_openai import OpenAIEmbeddings
 from langchain_community.vectorstores import Chroma
-from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
+from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langgraph.graph import END, StateGraph
-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 os
+import streamlit as st
 import requests
 import hashlib
+import json
 import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime
+from pydantic import BaseModel, ValidationError
+import traceback
 
-# -----------------------------------------------------
-# State Schema Definition
-# -----------------------------------------------------
-class AgentState(TypedDict):
-    messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]
-    context: Dict[str, Any]
-    metadata: Dict[str, Any]
-
-# -----------------------------------------------------
-# Configuration
-# -----------------------------------------------------
+# ------------------------------
+# Configuration & Constants
+# ------------------------------
 class ResearchConfig:
     DEEPSEEK_API_KEY = os.environ.get("DEEPSEEK_API_KEY")
     CHROMA_PATH = "chroma_db"
@@ -41,139 +30,171 @@ class ResearchConfig:
     CHUNK_OVERLAP = 64
     MAX_CONCURRENT_REQUESTS = 5
     EMBEDDING_DIMENSIONS = 1536
-    DOCUMENT_MAP = {
-        "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 with scientific rigor:
+    ANALYSIS_TEMPLATE = """**Technical Analysis Request**
 {context}
 
 Respond with:
-1. Key Technical Contributions (bullet points)
-2. Novel Methodologies
-3. Empirical Results (with metrics)
-4. Potential Applications
-5. Limitations & Future Directions
-
-Format: Markdown with LaTeX mathematical notation where applicable
-"""
-
-# Validate API key configuration
-if not ResearchConfig.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()
-
-# -----------------------------------------------------
-# Quantum Document Processing
-# -----------------------------------------------------
-class QuantumDocumentManager:
+1. Key Technical Innovations (markdown table)
+2. Methodological Breakdown (bullet points)
+3. Quantitative Results (LaTeX equations)
+4. Critical Evaluation
+5. Research Impact Assessment
+
+Include proper academic citations where applicable."""
+
+# ------------------------------
+# Document Schema & Content
+# ------------------------------
+DOCUMENT_CONTENT = {
+    "CV-Transformer Hybrid": {
+        "content": """## Hybrid Architecture for Computer Vision
+**Authors**: DeepVision Research Team  
+**Abstract**: Novel combination of convolutional layers with transformer attention mechanisms.
+
+### Key Innovations:
+- Cross-attention feature fusion
+- Adaptive spatial pooling
+- Multi-scale gradient propagation
+
+$$\\mathcal{L}_{total} = \\alpha\\mathcal{L}_{CE} + \\beta\\mathcal{L}_{SSIM}$$""",
+        "metadata": {
+            "year": 2024,
+            "domain": "computer_vision",
+            "citations": 142
+        }
+    },
+    "Quantum ML Advances": {
+        "content": """## Quantum Machine Learning Breakthroughs
+**Authors**: Quantum AI Lab
+
+### Achievements:
+- Quantum-enhanced SGD (40% faster convergence)
+- 5-qubit QNN achieving 98% accuracy
+- Hybrid quantum-classical GANs
+
+$$\\mathcal{H} = -\\sum_{i<j} J_{ij}\\sigma_i^z\\sigma_j^z - \\Gamma\\sum_i\\sigma_i^x$$""",
+        "metadata": {
+            "year": 2023,
+            "domain": "quantum_ml",
+            "citations": 89
+        }
+    }
+}
+
+class DocumentSchema(BaseModel):
+    content: str
+    metadata: dict
+    doc_id: str
+
+# ------------------------------
+# State Management
+# ------------------------------
+class ResearchState(TypedDict):
+    messages: Annotated[List[BaseMessage], add_messages]
+    context: Annotated[Dict[str, Any], "research_context"]
+    metadata: Annotated[Dict[str, str], "system_metadata"]
+
+# ------------------------------
+# Document Processing
+# ------------------------------
+class DocumentManager:
     def __init__(self):
         self.client = chromadb.PersistentClient(path=ResearchConfig.CHROMA_PATH)
         self.embeddings = OpenAIEmbeddings(
             model="text-embedding-3-large",
             dimensions=ResearchConfig.EMBEDDING_DIMENSIONS
         )
+    
+    def initialize_collections(self):
+        try:
+            self.research_col = self._create_collection("research")
+            self.dev_col = self._create_collection("development")
+        except Exception as e:
+            st.error(f"Collection initialization failed: {str(e)}")
+            traceback.print_exc()
+
+    def _create_collection(self, name: str) -> Chroma:
+        documents, metadatas, ids = [], [], []
         
-    def create_collection(self, documents: List[str], collection_name: str) -> Chroma:
+        for title, data in DOCUMENT_CONTENT.items():
+            try:
+                doc = DocumentSchema(
+                    content=data["content"],
+                    metadata=data["metadata"],
+                    doc_id=hashlib.sha256(title.encode()).hexdigest()[:16]
+                )
+                documents.append(doc.content)
+                metadatas.append(doc.metadata)
+                ids.append(doc.doc_id)
+            except ValidationError as e:
+                st.error(f"Invalid document format: {title} - {str(e)}")
+                continue
+
         splitter = RecursiveCharacterTextSplitter(
             chunk_size=ResearchConfig.CHUNK_SIZE,
             chunk_overlap=ResearchConfig.CHUNK_OVERLAP,
-            separators=["\n\n", "\n", "|||"]
+            separators=["\n## ", "\n### ", "\n\n", "\n• "]
         )
-        docs = splitter.create_documents(documents)
-        # Debug lines about chunk creation removed
-        return Chroma.from_documents(
-            documents=docs,
-            embedding=self.embeddings,
-            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:
-        """Create a unique ID for each document chunk."""
-        return f"{hashlib.sha256(content.encode()).hexdigest()[:16]}-{int(time.time())}"
-
-# Initialize document collections
-qdm = QuantumDocumentManager()
-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")
-
-# -----------------------------------------------------
-# Advanced Retrieval System
-# -----------------------------------------------------
+        
+        try:
+            docs = splitter.create_documents(documents, metadatas=metadatas)
+            return Chroma.from_documents(
+                docs,
+                self.embeddings,
+                client=self.client,
+                collection_name=name,
+                ids=ids
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failed creating {name} collection: {str(e)}")
+
+# ------------------------------
+# Retrieval System
+# ------------------------------
 class ResearchRetriever:
     def __init__(self):
-        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}
-            )
-        }
+        self.dm = DocumentManager()
+        self.dm.initialize_collections()
     
-    def retrieve(self, query: str, domain: str) -> List[Any]:
-        """Retrieve documents from the specified domain."""
+    def retrieve(self, query: str, domain: str) -> List[DocumentSchema]:
         try:
-            return self.retrievers[domain].invoke(query)
-        except KeyError:
-            st.error(f"[ERROR] Retrieval domain '{domain}' not found.")
+            collection = self.dm.research_col if domain == "research" else self.dm.dev_col
+            if not collection:
+                return []
+            
+            results = collection.as_retriever(
+                search_type="mmr",
+                search_kwargs={'k': 4, 'fetch_k': 20}
+            ).invoke(query)
+            
+            return [DocumentSchema(
+                content=doc.page_content,
+                metadata=doc.metadata,
+                doc_id=doc.metadata.get("doc_id", "")
+            ) for doc in results if doc.page_content]
+        
+        except Exception as e:
+            st.error(f"Retrieval failure: {str(e)}")
+            traceback.print_exc()
             return []
 
-retriever = ResearchRetriever()
-
-# -----------------------------------------------------
-# Cognitive Processing Unit
-# -----------------------------------------------------
-class CognitiveProcessor:
+# ------------------------------
+# Analysis Processor
+# ------------------------------
+class AnalysisEngine:
     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:
-        """Send the prompt to the DeepSeek API using triple redundancy for robustness."""
-        futures = []
-        for _ in range(3):
-            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:
-        """Make a single request to the DeepSeek API."""
+        self.session_hash = hashlib.sha256(str(time.time()).encode()).hexdigest()[:12]
+
+    def analyze(self, prompt: str) -> Dict:
+        futures = [self.executor.submit(self._api_request, prompt) for _ in range(3)]
+        return self._validate_results([f.result() for f in as_completed(futures)])
+
+    def _api_request(self, prompt: str) -> Dict:
         headers = {
             "Authorization": f"Bearer {ResearchConfig.DEEPSEEK_API_KEY}",
-            "Content-Type": "application/json",
-            "X-Research-Session": self.session_id
+            "X-Session-ID": self.session_hash,
+            "Content-Type": "application/json"
         }
         
         try:
@@ -182,293 +203,251 @@ class CognitiveProcessor:
                 headers=headers,
                 json={
                     "model": "deepseek-chat",
-                    "messages": [{
-                        "role": "user",
-                        "content": f"Respond as Senior AI Researcher:\n{prompt}"
-                    }],
+                    "messages": [{"role": "user", "content": prompt}],
                     "temperature": 0.7,
-                    "max_tokens": 1500,
-                    "top_p": 0.9
+                    "max_tokens": 2000
                 },
-                timeout=45
+                timeout=30
             )
             response.raise_for_status()
             return response.json()
-        except requests.exceptions.RequestException as e:
-            return {"error": str(e)}
-    
-    def _consensus_check(self, results: List[Dict]) -> Dict:
-        """Pick the best result by comparing content length among successful responses."""
+        except Exception as e:
+            return {"error": str(e), "status_code": 500}
+
+    def _validate_results(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"}
-        return max(valid, key=lambda x: len(x.get('choices', [{}])[0].get('message', {}).get('content', '')))
+            return {"error": "All analysis attempts failed", "results": results}
+        
+        # Corrected line with proper parenthesis closure
+        best = max(valid, key=lambda x: len(x.get('choices', [{}])[0].get('message', {}).get('content', '')))
+        return best
 
-# -----------------------------------------------------
-# Research Workflow Engine
-# -----------------------------------------------------
+# ------------------------------
+# Workflow Implementation
+# ------------------------------
 class ResearchWorkflow:
     def __init__(self):
-        self.processor = CognitiveProcessor()
-        self.workflow = StateGraph(AgentState)
-        self._build_workflow()
-
-    def _build_workflow(self):
-        # Register nodes in the state graph
-        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)
-
-        # Define workflow transitions
+        self.retriever = ResearchRetriever()
+        self.engine = AnalysisEngine()
+        self.workflow = StateGraph(ResearchState)
+        self._build_graph()
+
+    def _build_graph(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.set_entry_point("ingest")
         self.workflow.add_edge("ingest", "retrieve")
         self.workflow.add_edge("retrieve", "analyze")
         self.workflow.add_conditional_edges(
             "analyze",
-            self._quality_check,
+            self._quality_gate,
             {"valid": "validate", "invalid": "refine"}
         )
         self.workflow.add_edge("validate", END)
         self.workflow.add_edge("refine", "retrieve")
 
-        # Compile the final state machine
-        self.app = self.workflow.compile()
-
-    def ingest_query(self, state: AgentState) -> Dict:
-        """Extract the user query and store it in the state."""
+    def _ingest(self, state: ResearchState) -> ResearchState:
         try:
-            query = state["messages"][-1].content
+            query = next(msg.content for msg in reversed(state["messages"]) 
+                      if isinstance(msg, HumanMessage))
             return {
-                "messages": [AIMessage(content="Query ingested successfully")],
-                "context": {"raw_query": query},
-                "metadata": {"timestamp": datetime.now().isoformat()}
+                "messages": [AIMessage(content="Query ingested")],
+                "context": {
+                    "query": query,
+                    "documents": [],
+                    "errors": []
+                },
+                "metadata": {
+                    "session_id": hashlib.sha256(str(time.time()).encode()).hexdigest()[:8],
+                    "timestamp": datetime.now().isoformat()
+                }
             }
         except Exception as e:
-            return self._error_state(f"Ingestion Error: {str(e)}")
+            return self._handle_error(f"Ingest failed: {str(e)}", state)
 
-    def retrieve_documents(self, state: AgentState) -> Dict:
-        """Retrieve relevant documents from the 'research' domain."""
+    def _retrieve(self, state: ResearchState) -> ResearchState:
         try:
-            # Fallback check for 'raw_query'
-            if "raw_query" not in state["context"]:
-                return self._error_state("No 'raw_query' found in context. Make sure the ingest step has run.")
-
-            query = state["context"]["raw_query"]
-            docs = retriever.retrieve(query, "research")
+            docs = self.retriever.retrieve(state["context"]["query"], "research")
             return {
                 "messages": [AIMessage(content=f"Retrieved {len(docs)} documents")],
                 "context": {
+                    **state["context"],
                     "documents": docs,
                     "retrieval_time": time.time()
-                }
+                },
+                "metadata": state["metadata"]
             }
         except Exception as e:
-            return self._error_state(f"Retrieval Error: {str(e)}")
+            return self._handle_error(f"Retrieval error: {str(e)}", state)
 
-    def analyze_content(self, state: AgentState) -> Dict:
-        """Concatenate document contents and analyze them using the CognitiveProcessor."""
+    def _analyze(self, state: ResearchState) -> ResearchState:
+        docs = state["context"].get("documents", [])
+        if not docs:
+            return self._handle_error("No documents for analysis", state)
+        
         try:
-            if "documents" not in state["context"] or not state["context"]["documents"]:
-                return self._error_state("No documents retrieved; please check your query or retrieval process.")
+            context = "\n\n".join([d.content for d in docs])
+            prompt = ResearchConfig.ANALYSIS_TEMPLATE.format(context=context)
+            result = self.engine.analyze(prompt)
+            
+            if "error" in result:
+                raise RuntimeError(result["error"])
             
-            docs = "\n\n".join([
-                d.page_content for d in state["context"]["documents"]
-                if hasattr(d, "page_content") and d.page_content
-            ])
-            prompt = ResearchConfig.ANALYSIS_TEMPLATE.format(context=docs)
-            response = self.processor.process_query(prompt)
+            content = result['choices'][0]['message']['content']
             
-            if "error" in response:
-                return self._error_state(response["error"])
+            if len(content) < 200 or not any(c.isalpha() for c in content):
+                raise ValueError("Insufficient analysis content")
             
             return {
-                "messages": [AIMessage(content=response['choices'][0]['message']['content'])],
-                "context": {"analysis": response}
+                "messages": [AIMessage(content=content)],
+                "context": state["context"],
+                "metadata": state["metadata"]
             }
         except Exception as e:
-            return self._error_state(f"Analysis Error: {str(e)}")
-
-    def validate_output(self, state: AgentState) -> Dict:
-        """Validate the technical correctness of the analysis output."""
-        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)
-        return {
-            "messages": [AIMessage(content=analysis + f"\n\nValidation: {response.get('choices', [{}])[0].get('message', {}).get('content', '')}")]
-        }
+            return self._handle_error(f"Analysis failed: {str(e)}", state)
 
-    def refine_results(self, state: AgentState) -> Dict:
-        """Refine the analysis based on the validation feedback."""
-        refinement_prompt = f"""Refine this analysis:
-{state["messages"][-1].content}
+    def _validate(self, state: ResearchState) -> ResearchState:
+        return state
 
-Improve:
-1. Technical precision
-2. Empirical grounding
-3. Theoretical coherence"""
-        
-        response = self.processor.process_query(refinement_prompt)
-        return {
-            "messages": [AIMessage(content=response.get('choices', [{}])[0].get('message', {}).get('content', ''))],
-            "context": state["context"]
-        }
+    def _refine(self, state: ResearchState) -> ResearchState:
+        return state
 
-    def _quality_check(self, state: AgentState) -> str:
-        """Check if the validation step indicates a 'VALID' or 'INVALID' output."""
-        content = state["messages"][-1].content
-        return "valid" if "VALID" in content else "invalid"
+    def _quality_gate(self, state: ResearchState) -> str:
+        content = state["messages"][-1].content if state["messages"] else ""
+        required = ["Innovations", "Results", "Evaluation"]
+        return "valid" if all(kw in content for kw in required) else "invalid"
 
-    def _error_state(self, message: str) -> Dict:
-        """Return an error message and mark the state as erroneous."""
-        st.error(f"[ERROR] {message}")
+    def _handle_error(self, message: str, state: ResearchState) -> ResearchState:
         return {
-            "messages": [AIMessage(content=f"❌ {message}")],
-            "context": {"error": True},
-            "metadata": {"status": "error"}
+            "messages": [AIMessage(content=f"🚨 Error: {message}")],
+            "context": {
+                **state["context"],
+                "errors": state["context"]["errors"] + [message]
+            },
+            "metadata": state["metadata"]
         }
 
-# -----------------------------------------------------
-# Research Interface
-# -----------------------------------------------------
+# ------------------------------
+# User Interface
+# ------------------------------
 class ResearchInterface:
     def __init__(self):
-        self.workflow = ResearchWorkflow()
-        # Page config already set at the top.
-        self._inject_styles()
+        self.workflow = ResearchWorkflow().workflow.compile()
+        self._setup_interface()
+
+    def _setup_interface(self):
+        st.set_page_config(
+            page_title="Research Assistant",
+            layout="wide",
+            initial_sidebar_state="expanded"
+        )
+        self._apply_styles()
         self._build_sidebar()
-        self._build_main_interface()
+        self._build_main()
 
-    def _inject_styles(self):
-        """Inject custom CSS for a sleek interface."""
+    def _apply_styles(self):
         st.markdown("""
         <style>
-        :root {
-            --primary: #2ecc71;
-            --secondary: #3498db;
-            --background: #0a0a0a;
-            --text: #ecf0f1;
-        }
-        
         .stApp {
-            background: var(--background);
-            color: var(--text);
-            font-family: 'Roboto', sans-serif;
+            background: #0a192f;
+            color: #64ffda;
         }
-        
         .stTextArea textarea {
-            background: #1a1a1a !important;
-            color: var(--text) !important;
-            border: 2px solid var(--secondary);
-            border-radius: 8px;
-            padding: 1rem;
+            background: #172a45 !important;
+            color: #a8b2d1 !important;
         }
-        
         .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);
+            background: #233554;
+            border: 1px solid #64ffda;
         }
-        
-        .stExpander {
-            background: #1a1a1a;
-            border: 1px solid #2a2a2a;
-            border-radius: 8px;
+        .error-box {
+            border: 1px solid #ff4444;
+            border-radius: 5px;
+            padding: 1rem;
             margin: 1rem 0;
         }
         </style>
         """, unsafe_allow_html=True)
 
     def _build_sidebar(self):
-        """Construct the left sidebar with document info and metrics."""
         with st.sidebar:
-            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):
-        """Construct the main interface for query input and result display."""
-        st.title("🧠 NeuroResearch AI")
-        query = st.text_area("Research Query:", height=200,
-                             placeholder="Enter technical research question...")
+            st.title("🔍 Document Database")
+            for title, data in DOCUMENT_CONTENT.items():
+                with st.expander(title[:25]+"..."):
+                    st.markdown(f"```\n{data['content'][:300]}...\n```")
+
+    def _build_main(self):
+        st.title("🧠 Research Analysis System")
+        query = st.text_area("Enter your research query:", height=150)
         
-        if st.button("Execute Analysis", type="primary"):
-            self._execute_analysis(query)
+        if st.button("Start Analysis", type="primary"):
+            self._run_analysis(query)
 
-    def _execute_analysis(self, query: str):
-        """Execute the entire research workflow and render the results."""
+    def _run_analysis(self, query: str):
         try:
-            with st.spinner("Initializing Quantum Analysis..."):
-                results = self.workflow.app.stream(
-                    {"messages": [HumanMessage(content=query)], "context": {}, "metadata": {}}
-                )
-                for event in results:
-                    self._render_event(event)
-                st.success("✅ Analysis Completed Successfully")
+            with st.spinner("🔍 Analyzing documents..."):
+                state = {
+                    "messages": [HumanMessage(content=query)],
+                    "context": {
+                        "query": "",
+                        "documents": [],
+                        "errors": []
+                    },
+                    "metadata": {}
+                }
+                
+                for event in self.workflow.stream(state):
+                    self._display_progress(event)
+                
+                final_state = self.workflow.invoke(state)
+                self._show_results(final_state)
+                
         except Exception as e:
             st.error(f"""**Analysis Failed**
-{str(e)}
-Potential issues:
-- Complex query structure
-- Document correlation failure
-- Temporal processing constraints""")
-
-    def _render_event(self, event: Dict):
-        """Render each node's output in the UI as it streams through the workflow."""
-        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])
+            {str(e)}
+            Common solutions:
+            - Simplify your query
+            - Check document database status
+            - Verify API connectivity""")
+
+    def _display_progress(self, event):
+        current_state = next(iter(event.values()))
+        with st.container():
+            st.markdown("---")
+            cols = st.columns([1,2,1])
+            
+            with cols[0]:
+                st.subheader("Processing Stage")
+                stage = list(event.keys())[0].title()
+                st.code(stage)
+                
+            with cols[1]:
+                st.subheader("Documents")
+                docs = current_state["context"].get("documents", [])
+                st.metric("Retrieved", len(docs))
+                
+            with cols[2]:
+                st.subheader("Status")
+                if current_state["context"].get("errors"):
+                    st.error("Errors detected")
                 else:
-                    st.warning("⚠️ Validation Issues Detected")
-                    with st.expander("View Validation Details", expanded=True):
-                        st.markdown(content)
+                    st.success("Normal operation")
+
+    def _show_results(self, state: ResearchState):
+        if state["context"].get("errors"):
+            st.error("Analysis completed with errors")
+            with st.expander("Error Details"):
+                for error in state["context"]["errors"]:
+                    st.markdown(f"- {error}")
+        else:
+            st.success("Analysis completed successfully ✅")
+            with st.expander("Full Report"):
+                st.markdown(state["messages"][-1].content)
 
-# -----------------------------------------------------
-# Main Execution
-# -----------------------------------------------------
 if __name__ == "__main__":
-    ResearchInterface()
+    ResearchInterface()