Update app.py

app.py

@@ -9,27 +9,30 @@ 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
+from typing import Sequence, List, Dict
 import chromadb
 import re
 import os
 import streamlit as st
 import requests
+import hashlib
 from langchain.tools.retriever import create_retriever_tool
+from langchain.schema import Document
 
 # ------------------------------
 # Configuration
 # ------------------------------
-# Get DeepSeek API key from Hugging Face Space secrets
+# Get DeepSeek API key from environment variables
 DEEPSEEK_API_KEY = os.environ.get("DEEPSEEK_API_KEY")
 
+# Validate API key configuration
 if not DEEPSEEK_API_KEY:
     st.error("""
-    **Missing API Configuration**  
-    Please configure your DeepSeek API key in Hugging Face Space secrets:
-    1. Go to your Space's Settings
-    2. Click on 'Repository secrets'
-    3. Add a secret named DEEPSEEK_API_KEY
+    **Critical Configuration Missing**  
+    DeepSeek API key not found. Please ensure you have:
+    1. Created a Hugging Face Space secret named DEEPSEEK_API_KEY
+    2. Added your valid API key to the Space secrets
+    3. Restarted the Space after configuration
     """)
     st.stop()
 
@@ -39,10 +42,26 @@ os.makedirs("chroma_db", exist_ok=True)
 # ------------------------------
 # ChromaDB Client Configuration
 # ------------------------------
-chroma_client = chromadb.PersistentClient(path="chroma_db")
+chroma_client = chromadb.PersistentClient(
+    path="chroma_db",
+    settings=chromadb.config.Settings(anonymized_telemetry=False)
+    
+# ------------------------------
+# Document Processing Utilities
+# ------------------------------
+def deduplicate_docs(docs: List[Document]) -> List[Document]:
+    """Remove duplicate documents using content hashing"""
+    seen = set()
+    unique_docs = []
+    for doc in docs:
+        content_hash = hashlib.sha256(doc.page_content.encode()).hexdigest()
+        if content_hash not in seen:
+            seen.add(content_hash)
+            unique_docs.append(doc)
+    return unique_docs
 
 # ------------------------------
-# Dummy Data: Research & Development Texts
+# Data Preparation
 # ------------------------------
 research_texts = [
     "Research Report: Results of a New AI Model Improving Image Recognition Accuracy to 98%",
@@ -56,407 +75,461 @@ development_texts = [
     "Product Y: In the Performance Optimization Stage Before Release"
 ]
 
-# ------------------------------
-# Text Splitting & Document Creation
-# ------------------------------
-splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=10)
-research_docs = splitter.create_documents(research_texts)
-development_docs = splitter.create_documents(development_texts)
+# Create documents with metadata
+splitter = RecursiveCharacterTextSplitter(
+    chunk_size=150,
+    chunk_overlap=20,
+    length_function=len,
+    add_start_index=True
+)
+
+research_docs = splitter.create_documents(
+    research_texts,
+    metadatas=[{"source": "research", "doc_id": f"res_{i}"} for i in range(len(research_texts))]
+)
+
+development_docs = splitter.create_documents(
+    development_texts,
+    metadatas=[{"source": "development", "doc_id": f"dev_{i}"} for i in range(len(development_texts))]
+)
 
 # ------------------------------
-# Creating Vector Stores with Embeddings
+# Vector Store Initialization
 # ------------------------------
 embeddings = OpenAIEmbeddings(
     model="text-embedding-3-large",
-    # dimensions=1024  # Uncomment if needed
+    model_kwargs={"dimensions": 1024}
 )
 
 research_vectorstore = Chroma.from_documents(
     documents=research_docs,
     embedding=embeddings,
     client=chroma_client,
-    collection_name="research_collection"
+    collection_name="research_collection",
+    collection_metadata={"hnsw:space": "cosine"}
 )
 
 development_vectorstore = Chroma.from_documents(
     documents=development_docs,
     embedding=embeddings,
     client=chroma_client,
-    collection_name="development_collection"
+    collection_name="development_collection",
+    collection_metadata={"hnsw:space": "cosine"}
 )
 
-research_retriever = research_vectorstore.as_retriever()
-development_retriever = development_vectorstore.as_retriever()
-
 # ------------------------------
-# Creating Retriever Tools
+# Retriever Tools Configuration
 # ------------------------------
-research_tool = create_retriever_tool(
-    research_retriever,
-    "research_db_tool",
-    "Search information from the research database."
+research_retriever = research_vectorstore.as_retriever(
+    search_type="mmr",
+    search_kwargs={"k": 5, "fetch_k": 10}
 )
 
-development_tool = create_retriever_tool(
-    development_retriever,
-    "development_db_tool",
-    "Search information from the development database."
+development_retriever = development_vectorstore.as_retriever(
+    search_type="similarity",
+    search_kwargs={"k": 5}
 )
 
-tools = [research_tool, development_tool]
+tools = [
+    create_retriever_tool(
+        research_retriever,
+        "research_database",
+        "Searches through academic papers and research reports for technical AI advancements"
+    ),
+    create_retriever_tool(
+        development_retriever,
+        "development_database",
+        "Accesses current project statuses and development timelines"
+    )
+]
 
 # ------------------------------
-# Agent Function & Workflow Functions
+# Agent State Definition
 # ------------------------------
 class AgentState(TypedDict):
     messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]
 
+# ------------------------------
+# Core Agent Function
+# ------------------------------
 def agent(state: AgentState):
-    print("---CALL AGENT---")
+    """Main decision-making agent handling user queries"""
+    print("\n--- AGENT EXECUTION START ---")
     messages = state["messages"]
-
-    if isinstance(messages[0], tuple):
-        user_message = messages[0][1]
-    else:
-        user_message = messages[0].content
-
-    prompt = f"""Given this user question: "{user_message}"
-If it's about research or academic topics, respond EXACTLY in this format:
-SEARCH_RESEARCH: <search terms>
-
-If it's about development status, respond EXACTLY in this format:
-SEARCH_DEV: <search terms>
-
-Otherwise, just answer directly.
-"""
-
-    headers = {
-        "Accept": "application/json",
-        "Authorization": f"Bearer {DEEPSEEK_API_KEY}",
-        "Content-Type": "application/json"
-    }
-    
-    data = {
-        "model": "deepseek-chat",
-        "messages": [{"role": "user", "content": prompt}],
-        "temperature": 0.7,
-        "max_tokens": 1024
-    }
     
     try:
+        # Extract user message content
+        user_message = messages[-1].content if isinstance(messages[-1], HumanMessage) else ""
+        
+        # Construct analysis prompt
+        prompt = f"""Analyze this user query and determine the appropriate action:
+        
+        Query: {user_message}
+        
+        Response Format:
+        - If research-related (technical details, academic concepts), respond:
+          SEARCH_RESEARCH: [keywords]
+          
+        - If development-related (project status, timelines), respond:
+          SEARCH_DEV: [keywords]
+          
+        - If general question, answer directly
+        - If unclear, request clarification
+        """
+
+        # API request configuration
+        headers = {
+            "Accept": "application/json",
+            "Authorization": f"Bearer {DEEPSEEK_API_KEY}",
+            "Content-Type": "application/json"
+        }
+        
+        data = {
+            "model": "deepseek-chat",
+            "messages": [{"role": "user", "content": prompt}],
+            "temperature": 0.5,
+            "max_tokens": 256
+        }
+
+        # Execute API call
         response = requests.post(
             "https://api.deepseek.com/v1/chat/completions",
             headers=headers,
             json=data,
-            verify=False,
             timeout=30
         )
         response.raise_for_status()
         
+        # Process response
         response_text = response.json()['choices'][0]['message']['content']
-        print("Raw response:", response_text)
-        
+        print(f"Agent Decision: {response_text}")
+
+        # Handle different response types
         if "SEARCH_RESEARCH:" in response_text:
             query = response_text.split("SEARCH_RESEARCH:")[1].strip()
             results = research_retriever.invoke(query)
-            return {"messages": [AIMessage(content=f'Action: research_db_tool\n{{"query": "{query}"}}\n\nResults: {str(results)}')]}
+            unique_results = deduplicate_docs(results)
+            return {
+                "messages": [
+                    AIMessage(
+                        content=f'Action: research_database\nQuery: "{query}"\nResults: {len(unique_results)} relevant documents',
+                        additional_kwargs={"documents": unique_results}
+                    )
+                ]
+            }
 
         elif "SEARCH_DEV:" in response_text:
             query = response_text.split("SEARCH_DEV:")[1].strip()
             results = development_retriever.invoke(query)
-            return {"messages": [AIMessage(content=f'Action: development_db_tool\n{{"query": "{query}"}}\n\nResults: {str(results)}')]}
+            unique_results = deduplicate_docs(results)
+            return {
+                "messages": [
+                    AIMessage(
+                        content=f'Action: development_database\nQuery: "{query}"\nResults: {len(unique_results)} relevant documents',
+                        additional_kwargs={"documents": unique_results}
+                    )
+                ]
+            }
 
         else:
             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."
+    except requests.exceptions.HTTPError as e:
+        error_msg = f"API Error: {e.response.status_code} - {e.response.text}"
+        if "insufficient balance" in e.response.text.lower():
+            error_msg += "\n\nPlease check your DeepSeek account balance."
         return {"messages": [AIMessage(content=error_msg)]}
+    except Exception as e:
+        return {"messages": [AIMessage(content=f"Processing Error: {str(e)}")]}
 
+# ------------------------------
+# Document Evaluation Functions
+# ------------------------------
 def simple_grade_documents(state: AgentState):
+    """Evaluate retrieved document relevance"""
     messages = state["messages"]
     last_message = messages[-1]
-    print("Evaluating message:", last_message.content)
     
-    if "Results: [Document" in last_message.content:
-        print("---DOCS FOUND, GO TO GENERATE---")
+    if last_message.additional_kwargs.get("documents"):
+        print("--- Relevant Documents Found ---")
         return "generate"
     else:
-        print("---NO DOCS FOUND, TRY REWRITE---")
+        print("--- No Valid Documents Found ---")
         return "rewrite"
 
 def generate(state: AgentState):
-    print("---GENERATE FINAL ANSWER---")
+    """Generate final answer from documents"""
+    print("\n--- GENERATING FINAL ANSWER ---")
     messages = state["messages"]
-    question = messages[0].content if isinstance(messages[0], tuple) else messages[0].content
-    last_message = messages[-1]
-
-    docs = ""
-    if "Results: [" in last_message.content:
-        results_start = last_message.content.find("Results: [")
-        docs = last_message.content[results_start:]
-    print("Documents found:", docs)
-
-    headers = {
-        "Accept": "application/json",
-        "Authorization": f"Bearer {DEEPSEEK_API_KEY}",
-        "Content-Type": "application/json"
-    }
-    
-    prompt = f"""Based on these research documents, summarize the latest advancements in AI:
-Question: {question}
-Documents: {docs}
-Focus on extracting and synthesizing the key findings from the research papers.
-"""
-    
-    data = {
-        "model": "deepseek-chat",
-        "messages": [{
-            "role": "user",
-            "content": prompt
-        }],
-        "temperature": 0.7,
-        "max_tokens": 1024
-    }
     
     try:
-        print("Sending generate request to API...")
+        # Extract context
+        user_question = next(msg.content for msg in messages if isinstance(msg, HumanMessage))
+        documents = messages[-1].additional_kwargs.get("documents", [])
+        
+        # Format document sources
+        sources = list(set(
+            doc.metadata.get('source', 'unknown') 
+            for doc in documents
+        ))
+        
+        # Create analysis prompt
+        prompt = f"""Synthesize a technical answer using these documents:
+        
+        Question: {user_question}
+        
+        Documents:
+        {[doc.page_content for doc in documents]}
+        
+        Requirements:
+        1. Highlight quantitative metrics
+        2. Cite document sources (research/development)
+        3. Note temporal context
+        4. List potential applications
+        5. Mention limitations/gaps
+        """
+
+        # API request configuration
+        headers = {
+            "Accept": "application/json",
+            "Authorization": f"Bearer {DEEPSEEK_API_KEY}",
+            "Content-Type": "application/json"
+        }
+        
+        data = {
+            "model": "deepseek-chat",
+            "messages": [{"role": "user", "content": prompt}],
+            "temperature": 0.3,
+            "max_tokens": 1024
+        }
+
+        # Execute API call
         response = requests.post(
             "https://api.deepseek.com/v1/chat/completions",
             headers=headers,
             json=data,
-            verify=False,
-            timeout=30
+            timeout=45
         )
         response.raise_for_status()
         
+        # Format final answer
         response_text = response.json()['choices'][0]['message']['content']
-        print("Final Answer:", response_text)
-        return {"messages": [AIMessage(content=response_text)]}
+        formatted_answer = f"{response_text}\n\nSources: {', '.join(sources)}"
+        
+        return {"messages": [AIMessage(content=formatted_answer)]}
+
     except Exception as e:
-        error_msg = f"Generation Error: {str(e)}"
-        return {"messages": [AIMessage(content=error_msg)]}
+        return {"messages": [AIMessage(content=f"Generation Error: {str(e)}")]}
 
 def rewrite(state: AgentState):
-    print("---REWRITE QUESTION---")
+    """Rewrite unclear queries"""
+    print("\n--- REWRITING QUERY ---")
     messages = state["messages"]
-    original_question = messages[0].content if len(messages) > 0 else "N/A"
-
-    headers = {
-        "Accept": "application/json",
-        "Authorization": f"Bearer {DEEPSEEK_API_KEY}",
-        "Content-Type": "application/json"
-    }
-    
-    data = {
-        "model": "deepseek-chat",
-        "messages": [{
-            "role": "user",
-            "content": f"Rewrite this question to be more specific and clearer: {original_question}"
-        }],
-        "temperature": 0.7,
-        "max_tokens": 1024
-    }
     
     try:
-        print("Sending rewrite request...")
+        original_query = next(msg.content for msg in messages if isinstance(msg, HumanMessage))
+        
+        headers = {
+            "Accept": "application/json",
+            "Authorization": f"Bearer {DEEPSEEK_API_KEY}",
+            "Content-Type": "application/json"
+        }
+        
+        data = {
+            "model": "deepseek-chat",
+            "messages": [{
+                "role": "user",
+                "content": f"Clarify this query while preserving technical intent: {original_query}"
+            }],
+            "temperature": 0.5,
+            "max_tokens": 256
+        }
+
         response = requests.post(
             "https://api.deepseek.com/v1/chat/completions",
             headers=headers,
             json=data,
-            verify=False,
             timeout=30
         )
         response.raise_for_status()
         
-        response_text = response.json()['choices'][0]['message']['content']
-        print("Rewritten question:", response_text)
-        return {"messages": [AIMessage(content=response_text)]}
-    except Exception as e:
-        error_msg = f"Rewrite Error: {str(e)}"
-        return {"messages": [AIMessage(content=error_msg)]}
-
-tools_pattern = re.compile(r"Action: .*")
-
-def custom_tools_condition(state: AgentState):
-    messages = state["messages"]
-    last_message = messages[-1]
-    content = last_message.content
+        rewritten = response.json()['choices'][0]['message']['content']
+        return {"messages": [AIMessage(content=f"Revised Query: {rewritten}")]}
 
-    print("Checking tools condition:", content)
-    if tools_pattern.match(content):
-        print("Moving to retrieve...")
-        return "tools"
-    print("Moving to END...")
-    return END
+    except Exception as e:
+        return {"messages": [AIMessage(content=f"Rewriting Error: {str(e)}")]}
 
 # ------------------------------
-# Workflow Configuration using LangGraph
+# Workflow Configuration
 # ------------------------------
 workflow = StateGraph(AgentState)
 
-# Add nodes
+# Node Registration
 workflow.add_node("agent", agent)
-retrieve_node = ToolNode(tools)
-workflow.add_node("retrieve", retrieve_node)
-workflow.add_node("rewrite", rewrite)
+workflow.add_node("retrieve", ToolNode(tools))
 workflow.add_node("generate", generate)
+workflow.add_node("rewrite", rewrite)
 
-# Set entry point
+# Workflow Structure
 workflow.set_entry_point("agent")
 
-# Define transitions
 workflow.add_conditional_edges(
     "agent",
-    custom_tools_condition,
-    {
-        "tools": "retrieve",
-        END: END
-    }
+    lambda state: "tools" if any(
+        tool.name in state["messages"][-1].content
+        for tool in tools
+    ) else END,
+    {"tools": "retrieve", END: END}
 )
 
 workflow.add_conditional_edges(
     "retrieve",
     simple_grade_documents,
-    {
-        "generate": "generate",
-        "rewrite": "rewrite"
-    }
+    {"generate": "generate", "rewrite": "rewrite"}
 )
 
 workflow.add_edge("generate", END)
 workflow.add_edge("rewrite", "agent")
 
-# Compile the workflow
 app = workflow.compile()
 
 # ------------------------------
-# Processing Function
-# ------------------------------
-def process_question(user_question, app, config):
-    """Process user question through the workflow"""
-    events = []
-    for event in app.stream({"messages": [("user", user_question)]}, config):
-        events.append(event)
-    return events
-
-# ------------------------------
-# Streamlit App UI (Dark Theme)
+# Streamlit UI Implementation
 # ------------------------------
 def main():
+    """Main application interface"""
     st.set_page_config(
-        page_title="AI Research & Development Assistant",
-        layout="wide",
+        page_title="AI Research Assistant",
+        layout="centered",
         initial_sidebar_state="expanded"
     )
 
+    # Dark Theme Configuration
     st.markdown("""
     <style>
     .stApp {
-        background-color: #1a1a1a;
-        color: #ffffff;
+        background-color: #0E1117;
+        color: #FAFAFA;
     }
     
     .stTextArea textarea {
-        background-color: #2d2d2d !important;
-        color: #ffffff !important;
+        background-color: #262730 !important;
+        color: #FAFAFA !important;
+        border: 1px solid #3D4051;
     }
     
-    .stButton > button {
-        background-color: #4CAF50;
+    .stButton>button {
+        background-color: #2E8B57;
         color: white;
+        border-radius: 4px;
+        padding: 0.5rem 1rem;
         transition: all 0.3s;
     }
     
-    .stButton > button:hover {
-        background-color: #45a049;
+    .stButton>button:hover {
+        background-color: #3CB371;
         transform: scale(1.02);
     }
     
-    .data-box {
-        background-color: #2d2d2d;
-        border-left: 5px solid #2196F3;
+    .stAlert {
+        background-color: #1A1D23 !important;
+        border: 1px solid #3D4051;
     }
     
-    .dev-box {
-        border-left: 5px solid #4CAF50;
+    .stExpander {
+        background-color: #1A1D23;
+        border: 1px solid #3D4051;
     }
     
-    .st-expander {
-        background-color: #2d2d2d;
-        border: 1px solid #3d3d3d;
+    .data-source {
+        padding: 0.5rem;
+        margin: 0.5rem 0;
+        background-color: #1A1D23;
+        border-left: 3px solid #2E8B57;
+        border-radius: 4px;
     }
     </style>
     """, unsafe_allow_html=True)
 
+    # Sidebar Configuration
     with st.sidebar:
-        st.header("📚 Available Data")
-        st.subheader("Research Database")
-        for text in research_texts:
-            st.markdown(f'<div class="data-box research-box" style="padding: 15px; margin: 10px 0; border-radius: 5px;">{text}</div>', unsafe_allow_html=True)
+        st.header("Technical Databases")
+        with st.expander("Research Corpus", expanded=True):
+            st.markdown("""
+            - AI Model Architectures
+            - Machine Learning Advances
+            - Quantum Computing Applications
+            - Algorithmic Breakthroughs
+            """)
             
-        st.subheader("Development Database")
-        for text in development_texts:
-            st.markdown(f'<div class="data-box dev-box" style="padding: 15px; margin: 10px 0; border-radius: 5px;">{text}</div>', unsafe_allow_html=True)
-
-    st.title("🤖 AI Research & Development Assistant")
-    st.markdown("---")
-
-    query = st.text_area("Enter your question:", height=100, placeholder="e.g., What is the latest advancement in AI research?")
-
-    col1, col2 = st.columns([1, 2])
-    with col1:
-        if st.button("🔍 Get Answer", use_container_width=True):
-            if query:
-                try:
-                    with st.spinner('Processing your question...'):
-                        events = process_question(query, app, {"configurable": {"thread_id": "1"}})
+        with st.expander("Development Tracking", expanded=True):
+            st.markdown("""
+            - Project Milestones
+            - System Architecture
+            - Deployment Status
+            - Performance Metrics
+            """)
+
+    # Main Interface
+    st.title("🧠 AI Research Assistant")
+    st.caption("Technical Analysis and Development Tracking System")
+
+    query = st.text_area(
+        "Enter Technical Query:",
+        height=150,
+        placeholder="Example: Compare transformer architectures for medical imaging analysis..."
+    )
+
+    if st.button("Execute Analysis", use_container_width=True):
+        if not query:
+            st.warning("Please input a technical query")
+            return
+            
+        with st.status("Processing...", expanded=True) as status:
+            try:
+                events = []
+                for event in app.stream({"messages": [HumanMessage(content=query)]}):
+                    events.append(event)
+                    
+                    if 'agent' in event:
+                        status.update(label="Decision Making", state="running")
+                        st.session_state.agent_step = event['agent']
+                        
+                    if 'retrieve' in event:
+                        status.update(label="Document Retrieval", state="running")
+                        st.session_state.retrieved = event['retrieve']
                         
-                        for event in events:
-                            if 'agent' in event:
-                                with st.expander("🔄 Processing Step", expanded=True):
-                                    content = event['agent']['messages'][0].content
-                                    if "Error" in content:
-                                        st.error(content)
-                                    elif "Results:" in content:
-                                        st.markdown("### 📑 Retrieved Documents:")
-                                        docs_start = content.find("Results:")
-                                        docs = content[docs_start:]
-                                        st.info(docs)
-                            elif 'generate' in event:
-                                content = event['generate']['messages'][0].content
-                                if "Error" in content:
-                                    st.error(content)
-                                else:
-                                    st.markdown("### ✨ Final Answer:")
-                                    st.success(content)
-                except Exception as e:
-                    st.error(f"""
-                    **Processing Error**  
-                    {str(e)}  
-                    Please check:  
-                    - API key configuration  
-                    - Account balance  
-                    - Network connection
-                    """)
-            else:
-                st.warning("⚠️ Please enter a question first!")
-
-    with col2:
-        st.markdown("""
-        ### 🎯 How to Use
-        1. Enter your question in the text box
-        2. Click the search button
-        3. Review processing steps
-        4. See final answer
-
-        ### 💡 Example Questions
-        - What's new in AI image recognition?
-        - How is Project B progressing?
-        - Recent machine learning trends?
-        """)
+                    if 'generate' in event:
+                        status.update(label="Synthesizing Answer", state="running")
+                        st.session_state.final_answer = event['generate']
+                        
+                status.update(label="Analysis Complete", state="complete")
+                
+            except Exception as e:
+                status.update(label="Processing Failed", state="error")
+                st.error(f"""
+                **System Error**  
+                {str(e)}  
+                Please verify:
+                - API key validity
+                - Network connectivity
+                - Query complexity
+                """)
+
+        if 'final_answer' in st.session_state:
+            answer = st.session_state.final_answer['messages'][0].content
+            
+            with st.container():
+                st.subheader("Technical Analysis")
+                st.markdown("---")
+                st.markdown(answer)
+                
+                if "Sources:" in answer:
+                    st.markdown("""
+                    <div class="data-source">
+                    ℹ️ Document sources are derived from the internal research database
+                    </div>
+                    """, unsafe_allow_html=True)
 
 if __name__ == "__main__":
     main()