Update app.py

app.py

@@ -9,30 +9,27 @@ 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, List, Dict
+from typing import Sequence
 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 environment variables
+# Get DeepSeek API key from Hugging Face Space secrets
 DEEPSEEK_API_KEY = os.environ.get("DEEPSEEK_API_KEY")
 
-# Validate API key configuration
 if not DEEPSEEK_API_KEY:
     st.error("""
-    **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
+    **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
     """)
     st.stop()
 
@@ -42,28 +39,10 @@ os.makedirs("chroma_db", exist_ok=True)
 # ------------------------------
 # ChromaDB Client Configuration
 # ------------------------------
-# After (corrected)
-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
+chroma_client = chromadb.PersistentClient(path="chroma_db")
 
 # ------------------------------
-# Data Preparation
+# Dummy Data: Research & Development Texts
 # ------------------------------
 research_texts = [
     "Research Report: Results of a New AI Model Improving Image Recognition Accuracy to 98%",
@@ -77,461 +56,407 @@ development_texts = [
     "Product Y: In the Performance Optimization Stage Before Release"
 ]
 
-# 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))]
-)
+# ------------------------------
+# 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)
 
 # ------------------------------
-# Vector Store Initialization
+# Creating Vector Stores with Embeddings
 # ------------------------------
 embeddings = OpenAIEmbeddings(
     model="text-embedding-3-large",
-    model_kwargs={"dimensions": 1024}
+    # dimensions=1024  # Uncomment if needed
 )
 
 research_vectorstore = Chroma.from_documents(
     documents=research_docs,
     embedding=embeddings,
     client=chroma_client,
-    collection_name="research_collection",
-    collection_metadata={"hnsw:space": "cosine"}
+    collection_name="research_collection"
 )
 
 development_vectorstore = Chroma.from_documents(
     documents=development_docs,
     embedding=embeddings,
     client=chroma_client,
-    collection_name="development_collection",
-    collection_metadata={"hnsw:space": "cosine"}
+    collection_name="development_collection"
 )
 
+research_retriever = research_vectorstore.as_retriever()
+development_retriever = development_vectorstore.as_retriever()
+
 # ------------------------------
-# Retriever Tools Configuration
+# Creating Retriever Tools
 # ------------------------------
-research_retriever = research_vectorstore.as_retriever(
-    search_type="mmr",
-    search_kwargs={"k": 5, "fetch_k": 10}
+research_tool = create_retriever_tool(
+    research_retriever,
+    "research_db_tool",
+    "Search information from the research database."
 )
 
-development_retriever = development_vectorstore.as_retriever(
-    search_type="similarity",
-    search_kwargs={"k": 5}
+development_tool = create_retriever_tool(
+    development_retriever,
+    "development_db_tool",
+    "Search information from the development database."
 )
 
-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"
-    )
-]
+tools = [research_tool, development_tool]
 
 # ------------------------------
-# Agent State Definition
+# Agent Function & Workflow Functions
 # ------------------------------
 class AgentState(TypedDict):
     messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]
 
-# ------------------------------
-# Core Agent Function
-# ------------------------------
 def agent(state: AgentState):
-    """Main decision-making agent handling user queries"""
-    print("\n--- AGENT EXECUTION START ---")
+    print("---CALL AGENT---")
     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(f"Agent Decision: {response_text}")
-
-        # Handle different response types
+        print("Raw response:", response_text)
+        
         if "SEARCH_RESEARCH:" in response_text:
             query = response_text.split("SEARCH_RESEARCH:")[1].strip()
             results = research_retriever.invoke(query)
-            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}
-                    )
-                ]
-            }
+            return {"messages": [AIMessage(content=f'Action: research_db_tool\n{{"query": "{query}"}}\n\nResults: {str(results)}')]}
 
         elif "SEARCH_DEV:" in response_text:
             query = response_text.split("SEARCH_DEV:")[1].strip()
             results = development_retriever.invoke(query)
-            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}
-                    )
-                ]
-            }
+            return {"messages": [AIMessage(content=f'Action: development_db_tool\n{{"query": "{query}"}}\n\nResults: {str(results)}')]}
 
         else:
             return {"messages": [AIMessage(content=response_text)]}
 
-    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)}")]}
+        error_msg = f"API Error: {str(e)}"
+        if "Insufficient Balance" in str(e):
+            error_msg += "\n\nPlease check your DeepSeek API account balance."
+        return {"messages": [AIMessage(content=error_msg)]}
 
-# ------------------------------
-# 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 last_message.additional_kwargs.get("documents"):
-        print("--- Relevant Documents Found ---")
+    if "Results: [Document" in last_message.content:
+        print("---DOCS FOUND, GO TO GENERATE---")
         return "generate"
     else:
-        print("--- No Valid Documents Found ---")
+        print("---NO DOCS FOUND, TRY REWRITE---")
         return "rewrite"
 
 def generate(state: AgentState):
-    """Generate final answer from documents"""
-    print("\n--- GENERATING FINAL ANSWER ---")
+    print("---GENERATE 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:
-        # 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
+        print("Sending generate request to API...")
         response = requests.post(
             "https://api.deepseek.com/v1/chat/completions",
             headers=headers,
             json=data,
-            timeout=45
+            verify=False,
+            timeout=30
         )
         response.raise_for_status()
         
-        # Format final answer
         response_text = response.json()['choices'][0]['message']['content']
-        formatted_answer = f"{response_text}\n\nSources: {', '.join(sources)}"
-        
-        return {"messages": [AIMessage(content=formatted_answer)]}
-
+        print("Final Answer:", response_text)
+        return {"messages": [AIMessage(content=response_text)]}
     except Exception as e:
-        return {"messages": [AIMessage(content=f"Generation Error: {str(e)}")]}
+        error_msg = f"Generation Error: {str(e)}"
+        return {"messages": [AIMessage(content=error_msg)]}
 
 def rewrite(state: AgentState):
-    """Rewrite unclear queries"""
-    print("\n--- REWRITING QUERY ---")
+    print("---REWRITE QUESTION---")
     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:
-        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
-        }
-
+        print("Sending rewrite request...")
         response = requests.post(
             "https://api.deepseek.com/v1/chat/completions",
             headers=headers,
             json=data,
+            verify=False,
             timeout=30
         )
         response.raise_for_status()
         
-        rewritten = response.json()['choices'][0]['message']['content']
-        return {"messages": [AIMessage(content=f"Revised Query: {rewritten}")]}
-
+        response_text = response.json()['choices'][0]['message']['content']
+        print("Rewritten question:", response_text)
+        return {"messages": [AIMessage(content=response_text)]}
     except Exception as e:
-        return {"messages": [AIMessage(content=f"Rewriting Error: {str(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
+
+    print("Checking tools condition:", content)
+    if tools_pattern.match(content):
+        print("Moving to retrieve...")
+        return "tools"
+    print("Moving to END...")
+    return END
 
 # ------------------------------
-# Workflow Configuration
+# Workflow Configuration using LangGraph
 # ------------------------------
 workflow = StateGraph(AgentState)
 
-# Node Registration
+# Add nodes
 workflow.add_node("agent", agent)
-workflow.add_node("retrieve", ToolNode(tools))
-workflow.add_node("generate", generate)
+retrieve_node = ToolNode(tools)
+workflow.add_node("retrieve", retrieve_node)
 workflow.add_node("rewrite", rewrite)
+workflow.add_node("generate", generate)
 
-# Workflow Structure
+# Set entry point
 workflow.set_entry_point("agent")
 
+# Define transitions
 workflow.add_conditional_edges(
     "agent",
-    lambda state: "tools" if any(
-        tool.name in state["messages"][-1].content
-        for tool in tools
-    ) else END,
-    {"tools": "retrieve", END: END}
+    custom_tools_condition,
+    {
+        "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()
 
 # ------------------------------
-# Streamlit UI Implementation
+# 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)
 # ------------------------------
 def main():
-    """Main application interface"""
     st.set_page_config(
-        page_title="AI Research Assistant",
-        layout="centered",
+        page_title="AI Research & Development Assistant",
+        layout="wide",
         initial_sidebar_state="expanded"
     )
 
-    # Dark Theme Configuration
     st.markdown("""
     <style>
     .stApp {
-        background-color: #0E1117;
-        color: #FAFAFA;
+        background-color: #1a1a1a;
+        color: #ffffff;
     }
     
     .stTextArea textarea {
-        background-color: #262730 !important;
-        color: #FAFAFA !important;
-        border: 1px solid #3D4051;
+        background-color: #2d2d2d !important;
+        color: #ffffff !important;
     }
     
-    .stButton>button {
-        background-color: #2E8B57;
+    .stButton > button {
+        background-color: #4CAF50;
         color: white;
-        border-radius: 4px;
-        padding: 0.5rem 1rem;
         transition: all 0.3s;
     }
     
-    .stButton>button:hover {
-        background-color: #3CB371;
+    .stButton > button:hover {
+        background-color: #45a049;
         transform: scale(1.02);
     }
     
-    .stAlert {
-        background-color: #1A1D23 !important;
-        border: 1px solid #3D4051;
+    .data-box {
+        background-color: #2d2d2d;
+        border-left: 5px solid #2196F3;
     }
     
-    .stExpander {
-        background-color: #1A1D23;
-        border: 1px solid #3D4051;
+    .dev-box {
+        border-left: 5px solid #4CAF50;
     }
     
-    .data-source {
-        padding: 0.5rem;
-        margin: 0.5rem 0;
-        background-color: #1A1D23;
-        border-left: 3px solid #2E8B57;
-        border-radius: 4px;
+    .st-expander {
+        background-color: #2d2d2d;
+        border: 1px solid #3d3d3d;
     }
     </style>
     """, unsafe_allow_html=True)
 
-    # Sidebar Configuration
     with st.sidebar:
-        st.header("Technical Databases")
-        with st.expander("Research Corpus", expanded=True):
-            st.markdown("""
-            - AI Model Architectures
-            - Machine Learning Advances
-            - Quantum Computing Applications
-            - Algorithmic Breakthroughs
-            """)
-            
-        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
+        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)
             
-        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']
+        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"}})
                         
-                    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)
+                        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 __name__ == "__main__":
     main()