Update app.py

app.py

@@ -1,298 +1,462 @@
-"""
-AI Research Assistant
-"""
-
 # ------------------------------
-# Core Imports & Configuration
+# Imports & Dependencies
 # ------------------------------
-import os
-import re
-import time
-import chromadb
-import requests
-import streamlit as st
-from typing import Sequence, Tuple
-from typing_extensions import TypedDict, Annotated
-from langchain_core.messages import HumanMessage, AIMessage
-from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_openai import OpenAIEmbeddings
 from langchain_community.vectorstores import Chroma
-from langchain.tools.retriever import create_retriever_tool
+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 chromadb.config import Settings
-from langchain_openai import OpenAIEmbeddings
+from typing_extensions import TypedDict, Annotated
+from typing import Sequence
+import chromadb
+import re
+import os
+import streamlit as st
+import requests
+from langchain.tools.retriever import create_retriever_tool
 
 # ------------------------------
-# Type Definitions
+# Configuration
 # ------------------------------
-class AgentState(TypedDict):
-    messages: Annotated[Sequence[AIMessage | HumanMessage], add_messages]
+# Get DeepSeek API key from Hugging Face Space secrets
+DEEPSEEK_API_KEY = os.environ.get("DEEPSEEK_API_KEY")
+
+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
+    """)
+    st.stop()
+
+# Create directory for Chroma persistence
+os.makedirs("chroma_db", exist_ok=True)
 
 # ------------------------------
-# Configuration & Constants
+# ChromaDB Client Configuration
 # ------------------------------
-class Config:
-    API_KEY = os.environ.get("DEEPSEEK_API_KEY")
-    CHROMA_PATH = "chroma_db"
-    TEXT_SPLITTER_CONFIG = {
-        "chunk_size": 512,
-        "chunk_overlap": 128,
-        "separators": ["\n\n", "\n", ". ", "! ", "? "]
-    }
+chroma_client = chromadb.PersistentClient(path="chroma_db")
 
 # ------------------------------
-# Core System Components
+# Dummy Data: Research & Development Texts
 # ------------------------------
-class ResearchAssistant:
-    def __init__(self):
-        self.embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
-        self.vector_stores = self._init_vector_stores()
-        self.tools = self._create_tools()
-        self.workflow = self._build_workflow()
-
-    def _init_vector_stores(self) -> Tuple[Chroma, Chroma]:
-        """Initialize vector stores with proper document processing"""
-        splitter = RecursiveCharacterTextSplitter(**Config.TEXT_SPLITTER_CONFIG)
-        
-        research_docs = splitter.create_documents([
-            "Research Report: New AI Model Achieves 98% Image Recognition Accuracy",
-            "Transformers: The New NLP Architecture Standard",
-            "Quantum Machine Learning: Emerging Trends and Applications"
-        ])
-        
-        development_docs = splitter.create_documents([
-            "Project A: UI Design Finalized, API Integration Phase",
-            "Project B: Feature Testing and Bug Fixes",
-            "Product Y: Performance Optimization Pre-Release"
-        ])
-
-        client = chromadb.PersistentClient(
-            path=Config.CHROMA_PATH,
-            settings=Settings(anonymized_telemetry=False)
-        
-        return (
-            Chroma.from_documents(research_docs, self.embeddings, 
-                                client=client, collection_name="research"),
-            Chroma.from_documents(development_docs, self.embeddings,
-                                client=client, collection_name="development")
-        )
+research_texts = [
+    "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"
+]
 
-    def _create_tools(self):
-        """Create retrieval tools with optimized search parameters"""
-        research_retriever = self.vector_stores[0].as_retriever(
-            search_kwargs={"k": 3, "score_threshold": 0.7}
-        )
-        development_retriever = self.vector_stores[1].as_retriever(
-            search_kwargs={"k": 3, "score_threshold": 0.7}
+development_texts = [
+    "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"
+]
+
+# ------------------------------
+# 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)
+
+# ------------------------------
+# Creating Vector Stores with Embeddings
+# ------------------------------
+embeddings = OpenAIEmbeddings(
+    model="text-embedding-3-large",
+    # dimensions=1024  # Uncomment if needed
+)
+
+research_vectorstore = Chroma.from_documents(
+    documents=research_docs,
+    embedding=embeddings,
+    client=chroma_client,
+    collection_name="research_collection"
+)
+
+development_vectorstore = Chroma.from_documents(
+    documents=development_docs,
+    embedding=embeddings,
+    client=chroma_client,
+    collection_name="development_collection"
+)
+
+research_retriever = research_vectorstore.as_retriever()
+development_retriever = development_vectorstore.as_retriever()
+
+# ------------------------------
+# Creating Retriever Tools
+# ------------------------------
+research_tool = create_retriever_tool(
+    research_retriever,
+    "research_db_tool",
+    "Search information from the research database."
+)
+
+development_tool = create_retriever_tool(
+    development_retriever,
+    "development_db_tool",
+    "Search information from the development database."
+)
+
+tools = [research_tool, development_tool]
+
+# ------------------------------
+# Agent Function & Workflow Functions
+# ------------------------------
+class AgentState(TypedDict):
+    messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]
+
+def agent(state: AgentState):
+    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:
+        response = requests.post(
+            "https://api.deepseek.com/v1/chat/completions",
+            headers=headers,
+            json=data,
+            verify=False,
+            timeout=30
         )
+        response.raise_for_status()
         
-        return [
-            create_retriever_tool(
-                research_retriever,
-                "research_db",
-                "Access technical research papers and reports"
-            ),
-            create_retriever_tool(
-                development_retriever,
-                "development_db",
-                "Retrieve project development status updates"
-            )
-        ]
-
-    def _build_workflow(self):
-        """Construct and return the processing workflow"""
-        workflow = StateGraph(AgentState)
-        
-        workflow.add_node("analyze", self.analyze_query)
-        workflow.add_node("retrieve", ToolNode(self.tools))
-        workflow.add_node("synthesize", self.synthesize_response)
+        response_text = response.json()['choices'][0]['message']['content']
+        print("Raw response:", response_text)
         
-        workflow.set_entry_point("analyze")
-        
-        workflow.add_conditional_edges(
-            "analyze",
-            self._needs_retrieval,
-            {"retrieve": "retrieve", "direct": "synthesize"}
+        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)}')]}
+
+        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)}')]}
+
+        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."
+        return {"messages": [AIMessage(content=error_msg)]}
+
+def simple_grade_documents(state: AgentState):
+    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---")
+        return "generate"
+    else:
+        print("---NO DOCS FOUND, TRY REWRITE---")
+        return "rewrite"
+
+def generate(state: AgentState):
+    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:
+        print("Sending generate request to API...")
+        response = requests.post(
+            "https://api.deepseek.com/v1/chat/completions",
+            headers=headers,
+            json=data,
+            verify=False,
+            timeout=30
         )
+        response.raise_for_status()
         
-        workflow.add_edge("retrieve", "synthesize")
-        workflow.add_edge("synthesize", END)
-        
-        return workflow.compile()
-
-    def _needs_retrieval(self, state: AgentState) -> str:
-        """Determine if document retrieval is needed"""
-        query = state["messages"][-1].content.lower()
-        return "retrieve" if any(kw in query for kw in {
-            "research", "study", "project", "develop", "trend"
-        }) else "direct"
-
-    def analyze_query(self, state: AgentState):
-        """Analyze user query and determine next steps"""
-        try:
-            user_input = state["messages"][-1].content
-            
-            headers = {
-                "Authorization": f"Bearer {Config.API_KEY}",
-                "Content-Type": "application/json"
-            }
-            
-            response = requests.post(
-                "https://api.deepseek.com/v1/chat/completions",
-                headers=headers,
-                json={
-                    "model": "deepseek-chat",
-                    "messages": [{
-                        "role": "user",
-                        "content": f"""Analyze this query and format as:
-                        CATEGORY: [RESEARCH|DEVELOPMENT|GENERAL]
-                        KEY_TERMS: comma-separated list
-                        {user_input}"""
-                    }],
-                    "temperature": 0.3
-                },
-                timeout=15
-            )
-            
-            response.raise_for_status()
-            analysis = response.json()["choices"][0]["message"]["content"]
-            
-            return {"messages": [AIMessage(content=analysis)]}
-        
-        except Exception as e:
-            return {"messages": [AIMessage(
-                content=f"Analysis Error: {str(e)}. Please rephrase your question."
-            )]}
-
-    def synthesize_response(self, state: AgentState):
-        """Generate final response with citations"""
-        try:
-            context = "\n".join([msg.content for msg in state["messages"]])
-            
-            headers = {
-                "Authorization": f"Bearer {Config.API_KEY}",
-                "Content-Type": "application/json"
-            }
-            
-            response = requests.post(
-                "https://api.deepseek.com/v1/chat/completions",
-                headers=headers,
-                json={
-                    "model": "deepseek-chat",
-                    "messages": [{
-                        "role": "user",
-                        "content": f"""Synthesize this information:
-                        {context}
-                        
-                        Include:
-                        1. Key findings
-                        2. Supporting evidence
-                        3. Technical details
-                        4. Potential applications"""
-                    }],
-                    "temperature": 0.5
-                },
-                timeout=20
-            )
-            
-            response.raise_for_status()
-            return {"messages": [AIMessage(
-                content=response.json()["choices"][0]["message"]["content"]
-            )]}
+        response_text = response.json()['choices'][0]['message']['content']
+        print("Final Answer:", response_text)
+        return {"messages": [AIMessage(content=response_text)]}
+    except Exception as e:
+        error_msg = f"Generation Error: {str(e)}"
+        return {"messages": [AIMessage(content=error_msg)]}
+
+def rewrite(state: AgentState):
+    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:
+        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()
         
-        except Exception as e:
-            return {"messages": [AIMessage(
-                content=f"Synthesis Error: {str(e)}. Please try again later."
-            )]}
+        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
+
+    print("Checking tools condition:", content)
+    if tools_pattern.match(content):
+        print("Moving to retrieve...")
+        return "tools"
+    print("Moving to END...")
+    return END
+
+# ------------------------------
+# Workflow Configuration using LangGraph
+# ------------------------------
+workflow = StateGraph(AgentState)
+
+# Add nodes
+workflow.add_node("agent", agent)
+retrieve_node = ToolNode(tools)
+workflow.add_node("retrieve", retrieve_node)
+workflow.add_node("rewrite", rewrite)
+workflow.add_node("generate", generate)
+
+# Set entry point
+workflow.set_entry_point("agent")
+
+# Define transitions
+workflow.add_conditional_edges(
+    "agent",
+    custom_tools_condition,
+    {
+        "tools": "retrieve",
+        END: END
+    }
+)
+
+workflow.add_conditional_edges(
+    "retrieve",
+    simple_grade_documents,
+    {
+        "generate": "generate",
+        "rewrite": "rewrite"
+    }
+)
+
+workflow.add_edge("generate", END)
+workflow.add_edge("rewrite", "agent")
+
+# Compile the workflow
+app = workflow.compile()
 
 # ------------------------------
-# Professional UI Interface
+# 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():
     st.set_page_config(
-        page_title="Research Assistant Pro",
+        page_title="AI Research & Development Assistant",
         layout="wide",
         initial_sidebar_state="expanded"
     )
-    
-    # Dark theme implementation
+
     st.markdown("""
     <style>
     .stApp {
-        background-color: #0f1114;
+        background-color: #1a1a1a;
         color: #ffffff;
     }
-    .stTextInput input, .stTextArea textarea {
-        background-color: #1e1e24 !important;
+    
+    .stTextArea textarea {
+        background-color: #2d2d2d !important;
         color: #ffffff !important;
     }
-    .stButton>button {
-        background: #2563eb;
-        transition: all 0.2s;
+    
+    .stButton > button {
+        background-color: #4CAF50;
+        color: white;
+        transition: all 0.3s;
     }
-    .stButton>button:hover {
-        background: #1d4ed8;
+    
+    .stButton > button:hover {
+        background-color: #45a049;
         transform: scale(1.02);
     }
-    .result-card {
-        background: #1a1a1f;
-        border-radius: 8px;
-        padding: 1.5rem;
-        margin: 1rem 0;
+    
+    .data-box {
+        background-color: #2d2d2d;
+        border-left: 5px solid #2196F3;
     }
-    </style>
-    """, unsafe_allow_html=True)
     
-    st.title("🔍 Research Assistant Pro")
-    st.write("Advanced AI-Powered Research Analysis")
+    .dev-box {
+        border-left: 5px solid #4CAF50;
+    }
     
+    .st-expander {
+        background-color: #2d2d2d;
+        border: 1px solid #3d3d3d;
+    }
+    </style>
+    """, unsafe_allow_html=True)
+
+    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.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:
-        with st.form("query_form"):
-            query = st.text_area("Research Query:", height=150,
-                                placeholder="Enter your research question...")
-            submitted = st.form_submit_button("Analyze")
-            
-            if submitted and query:
-                with st.spinner("Processing..."):
-                    try:
-                        assistant = ResearchAssistant()
-                        result = assistant.workflow.invoke({"messages": [
-                            HumanMessage(content=query)
-                        ]})
+        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("Analysis Details", expanded=True):
-                            st.markdown(f"""
-                            <div class="result-card">
-                                {result['messages'][-1].content}
-                            </div>
-                            """, unsafe_allow_html=True)
-                    except Exception as e:
-                        st.error(f"Processing Error: {str(e)}")
+                        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.subheader("Knowledge Base")
-        with st.expander("Research Documents"):
-            st.info("""
-            - Advanced Image Recognition Systems
-            - Transformer Architecture Analysis
-            - Quantum ML Research
-            """)
-        
-        with st.expander("Development Updates"):
-            st.info("""
-            - Project A: API Integration Phase
-            - Project B: Feature Testing
-            - Product Y: Optimization Stage
-            """)
+        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
 
-if __name__ == "__main__":
-    if not Config.API_KEY:
-        st.error("""
-        🔑 Configuration Required:
-        Set DEEPSEEK_API_KEY environment variable
+        ### 💡 Example Questions
+        - What's new in AI image recognition?
+        - How is Project B progressing?
+        - Recent machine learning trends?
         """)
-        st.stop()
+
+if __name__ == "__main__":
     main()