Update app.py

app.py

@@ -1,28 +1,39 @@
-# ------------------------------
-# Imports & Dependencies
-# ------------------------------
+# -----------------------------------------------------
+# 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")
+
 from langchain_openai import OpenAIEmbeddings
 from langchain_community.vectorstores import Chroma
-from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
+from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
 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
 
-# ------------------------------
-# Configuration & Constants
-# ------------------------------
+# -----------------------------------------------------
+# State Schema Definition
+# -----------------------------------------------------
+class AgentState(TypedDict):
+    messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]
+    context: Dict[str, Any]
+    metadata: Dict[str, Any]
+
+# -----------------------------------------------------
+# Configuration
+# -----------------------------------------------------
 class ResearchConfig:
     DEEPSEEK_API_KEY = os.environ.get("DEEPSEEK_API_KEY")
     CHROMA_PATH = "chroma_db"
@@ -30,171 +41,139 @@ class ResearchConfig:
     CHUNK_OVERLAP = 64
     MAX_CONCURRENT_REQUESTS = 5
     EMBEDDING_DIMENSIONS = 1536
-    ANALYSIS_TEMPLATE = """**Technical Analysis Request**
+    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:
 {context}
 
 Respond with:
-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:
+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:
     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 = [], [], []
         
-        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
-
+    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\n", "\n• "]
+            separators=["\n\n", "\n", "|||"]
         )
-        
-        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
-# ------------------------------
+        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
+# -----------------------------------------------------
 class ResearchRetriever:
     def __init__(self):
-        self.dm = DocumentManager()
-        self.dm.initialize_collections()
+        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, domain: str) -> List[DocumentSchema]:
+    def retrieve(self, query: str, domain: str) -> List[Any]:
+        """Retrieve documents from the specified domain."""
         try:
-            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 self.retrievers[domain].invoke(query)
+        except KeyError:
+            st.error(f"[ERROR] Retrieval domain '{domain}' not found.")
             return []
 
-# ------------------------------
-# Analysis Processor
-# ------------------------------
-class AnalysisEngine:
+retriever = ResearchRetriever()
+
+# -----------------------------------------------------
+# Cognitive Processing Unit
+# -----------------------------------------------------
+class CognitiveProcessor:
     def __init__(self):
         self.executor = ThreadPoolExecutor(max_workers=ResearchConfig.MAX_CONCURRENT_REQUESTS)
-        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:
+        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."""
         headers = {
             "Authorization": f"Bearer {ResearchConfig.DEEPSEEK_API_KEY}",
-            "X-Session-ID": self.session_hash,
-            "Content-Type": "application/json"
+            "Content-Type": "application/json",
+            "X-Research-Session": self.session_id
         }
         
         try:
@@ -203,252 +182,293 @@ class AnalysisEngine:
                 headers=headers,
                 json={
                     "model": "deepseek-chat",
-                    "messages": [{"role": "user", "content": prompt}],
+                    "messages": [{
+                        "role": "user",
+                        "content": f"Respond as Senior AI Researcher:\n{prompt}"
+                    }],
                     "temperature": 0.7,
-                    "max_tokens": 2000
+                    "max_tokens": 1500,
+                    "top_p": 0.9
                 },
-                timeout=30
+                timeout=45
             )
             response.raise_for_status()
             return response.json()
-        except Exception as e:
-            return {"error": str(e), "status_code": 500}
-
-    def _validate_results(self, results: List[Dict]) -> Dict:
+        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."""
         valid = [r for r in results if "error" not in r]
         if not valid:
-            return {"error": "All analysis attempts failed", "results": results}
-        
-        best = max(valid, key=lambda x: len(x.get('choices', [{}])[0].get('message', {}).get('content', ''))
-        return best
+            return {"error": "All API requests failed"}
+        return max(valid, key=lambda x: len(x.get('choices', [{}])[0].get('message', {}).get('content', '')))
 
-# ------------------------------
-# Workflow Implementation
-# ------------------------------
+# -----------------------------------------------------
+# Research Workflow Engine
+# -----------------------------------------------------
 class ResearchWorkflow:
     def __init__(self):
-        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.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.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_gate,
+            self._quality_check,
             {"valid": "validate", "invalid": "refine"}
         )
         self.workflow.add_edge("validate", END)
         self.workflow.add_edge("refine", "retrieve")
 
-    def _ingest(self, state: ResearchState) -> ResearchState:
+        # 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."""
         try:
-            query = next(msg.content for msg in reversed(state["messages"]) 
-                      if isinstance(msg, HumanMessage))
+            query = state["messages"][-1].content
             return {
-                "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()
-                }
+                "messages": [AIMessage(content="Query ingested successfully")],
+                "context": {"raw_query": query},
+                "metadata": {"timestamp": datetime.now().isoformat()}
             }
         except Exception as e:
-            return self._handle_error(f"Ingest failed: {str(e)}", state)
+            return self._error_state(f"Ingestion Error: {str(e)}")
 
-    def _retrieve(self, state: ResearchState) -> ResearchState:
+    def retrieve_documents(self, state: AgentState) -> Dict:
+        """Retrieve relevant documents from the 'research' domain."""
         try:
-            docs = self.retriever.retrieve(state["context"]["query"], "research")
+            # 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")
             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._handle_error(f"Retrieval error: {str(e)}", state)
+            return self._error_state(f"Retrieval Error: {str(e)}")
 
-    def _analyze(self, state: ResearchState) -> ResearchState:
-        docs = state["context"].get("documents", [])
-        if not docs:
-            return self._handle_error("No documents for analysis", state)
-        
+    def analyze_content(self, state: AgentState) -> Dict:
+        """Concatenate document contents and analyze them using the CognitiveProcessor."""
         try:
-            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"])
+            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.")
             
-            content = result['choices'][0]['message']['content']
+            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)
             
-            if len(content) < 200 or not any(c.isalpha() for c in content):
-                raise ValueError("Insufficient analysis content")
+            if "error" in response:
+                return self._error_state(response["error"])
             
             return {
-                "messages": [AIMessage(content=content)],
-                "context": state["context"],
-                "metadata": state["metadata"]
+                "messages": [AIMessage(content=response['choices'][0]['message']['content'])],
+                "context": {"analysis": response}
             }
         except Exception as e:
-            return self._handle_error(f"Analysis failed: {str(e)}", state)
+            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', '')}")]
+        }
 
-    def _validate(self, state: ResearchState) -> ResearchState:
-        # Implementation of validation logic
-        return 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 _refine(self, state: ResearchState) -> ResearchState:
-        # Implementation of refinement logic
-        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 _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 _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 _handle_error(self, message: str, state: ResearchState) -> ResearchState:
+    def _error_state(self, message: str) -> Dict:
+        """Return an error message and mark the state as erroneous."""
+        st.error(f"[ERROR] {message}")
         return {
-            "messages": [AIMessage(content=f"🚨 Error: {message}")],
-            "context": {
-                **state["context"],
-                "errors": state["context"]["errors"] + [message]
-            },
-            "metadata": state["metadata"]
+            "messages": [AIMessage(content=f"❌ {message}")],
+            "context": {"error": True},
+            "metadata": {"status": "error"}
         }
 
-# ------------------------------
-# User Interface
-# ------------------------------
+# -----------------------------------------------------
+# Research Interface
+# -----------------------------------------------------
 class ResearchInterface:
     def __init__(self):
-        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.workflow = ResearchWorkflow()
+        # Page config already set at the top.
+        self._inject_styles()
         self._build_sidebar()
-        self._build_main()
+        self._build_main_interface()
 
-    def _apply_styles(self):
+    def _inject_styles(self):
+        """Inject custom CSS for a sleek interface."""
         st.markdown("""
         <style>
+        :root {
+            --primary: #2ecc71;
+            --secondary: #3498db;
+            --background: #0a0a0a;
+            --text: #ecf0f1;
+        }
+        
         .stApp {
-            background: #0a192f;
-            color: #64ffda;
+            background: var(--background);
+            color: var(--text);
+            font-family: 'Roboto', sans-serif;
         }
+        
         .stTextArea textarea {
-            background: #172a45 !important;
-            color: #a8b2d1 !important;
+            background: #1a1a1a !important;
+            color: var(--text) !important;
+            border: 2px solid var(--secondary);
+            border-radius: 8px;
+            padding: 1rem;
         }
+        
         .stButton>button {
-            background: #233554;
-            border: 1px solid #64ffda;
+            background: linear-gradient(135deg, var(--primary), var(--secondary));
+            border: none;
+            border-radius: 8px;
+            padding: 1rem 2rem;
+            transition: all 0.3s;
         }
-        .error-box {
-            border: 1px solid #ff4444;
-            border-radius: 5px;
-            padding: 1rem;
+        
+        .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):
+        """Construct the left sidebar with document info and metrics."""
         with st.sidebar:
-            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)
+            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...")
         
-        if st.button("Start Analysis", type="primary"):
-            self._run_analysis(query)
+        if st.button("Execute Analysis", type="primary"):
+            self._execute_analysis(query)
 
-    def _run_analysis(self, query: str):
+    def _execute_analysis(self, query: str):
+        """Execute the entire research workflow and render the results."""
         try:
-            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)
-                
+            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")
         except Exception as e:
             st.error(f"""**Analysis Failed**
-            {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")
+{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])
                 else:
-                    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)
+                    st.warning("⚠️ Validation Issues Detected")
+                    with st.expander("View Validation Details", expanded=True):
+                        st.markdown(content)
 
+# -----------------------------------------------------
+# Main Execution
+# -----------------------------------------------------
 if __name__ == "__main__":
-    ResearchInterface()
+    ResearchInterface()