Update app.py

app.py

@@ -1,462 +1,291 @@
+"""
+AI Research Assistant - Professional Edition
+"""
+
 # ------------------------------
-# Imports & Dependencies
+# Core Imports & Configuration
 # ------------------------------
-from langchain_openai import OpenAIEmbeddings
-from langchain_community.vectorstores import Chroma
-from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
+import os
+import re
+import time
+import chromadb
+import requests
+import streamlit as st
+from typing import Sequence, Tuple
+from langchain_core.messages import HumanMessage, AIMessage
 from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import Chroma
+from langchain.tools.retriever import create_retriever_tool
 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
-import chromadb
-import re
-import os
-import streamlit as st
-import requests
-from langchain.tools.retriever import create_retriever_tool
+from chromadb.config import Settings
 
 # ------------------------------
-# Configuration
+# Configuration & Constants
 # ------------------------------
-# 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)
-
-# ------------------------------
-# ChromaDB Client Configuration
-# ------------------------------
-chroma_client = chromadb.PersistentClient(path="chroma_db")
-
-# ------------------------------
-# Dummy Data: Research & Development Texts
-# ------------------------------
-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"
-]
-
-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)
+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", ". ", "! ", "? "]
+    }
 
 # ------------------------------
-# Creating Vector Stores with Embeddings
+# Core System Components
 # ------------------------------
-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()
+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 and return research & development vector stores"""
+        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"
+        ])
 
-# ------------------------------
-# Creating Retriever Tools
-# ------------------------------
-research_tool = create_retriever_tool(
-    research_retriever,
-    "research_db_tool",
-    "Search information from the research database."
-)
+        client = chromadb.PersistentClient(path=Config.CHROMA_PATH)
+        
+        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")
+        )
 
-development_tool = create_retriever_tool(
-    development_retriever,
-    "development_db_tool",
-    "Search information from the development database."
-)
+    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}
+        )
+        
+        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)
+        
+        workflow.set_entry_point("analyze")
+        
+        workflow.add_conditional_edges(
+            "analyze",
+            self._needs_retrieval,
+            {"retrieve": "retrieve", "direct": "synthesize"}
+        )
+        
+        workflow.add_edge("retrieve", "synthesize")
+        workflow.add_edge("synthesize", END)
+        
+        return workflow.compile()
 
-tools = [research_tool, development_tool]
+    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"
 
 # ------------------------------
-# Agent Function & Workflow Functions
+# Processing 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>
+    messages: Annotated[Sequence[AIMessage], add_messages]
 
-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
-    }
-    
+def analyze_query(state: AgentState):
+    """Analyze user query and determine next steps"""
     try:
-        response = requests.post(
-            "https://api.deepseek.com/v1/chat/completions",
-            headers=headers,
-            json=data,
-            verify=False,
-            timeout=30
-        )
-        response.raise_for_status()
+        user_input = state["messages"][-1].content
         
-        response_text = response.json()['choices'][0]['message']['content']
-        print("Raw response:", response_text)
+        headers = {
+            "Authorization": f"Bearer {Config.API_KEY}",
+            "Content-Type": "application/json"
+        }
         
-        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
+            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"]
         
-        response_text = response.json()['choices'][0]['message']['content']
-        print("Final Answer:", response_text)
-        return {"messages": [AIMessage(content=response_text)]}
+        return {"messages": [AIMessage(content=analysis)]}
+    
     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"
+        return {"messages": [AIMessage(
+            content=f"Analysis Error: {str(e)}. Please rephrase your question."
+        )]}
 
-    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
-    }
-    
+def synthesize_response(state: AgentState):
+    """Generate final response with citations"""
     try:
-        print("Sending rewrite request...")
+        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=data,
-            verify=False,
-            timeout=30
+            json={
+                "model": "deepseek-chat",
+                "messages": [{
+                    "role": "user",
+                    "content": f"""Synthesize this information into a professional report:
+                    {context}
+                    
+                    Include:
+                    1. Key findings
+                    2. Supporting evidence
+                    3. Technical details
+                    4. Potential applications"""
+                }],
+                "temperature": 0.5
+            },
+            timeout=20
         )
-        response.raise_for_status()
         
-        response_text = response.json()['choices'][0]['message']['content']
-        print("Rewritten question:", response_text)
-        return {"messages": [AIMessage(content=response_text)]}
+        response.raise_for_status()
+        return {"messages": [AIMessage(
+            content=response.json()["choices"][0]["message"]["content"]
+        )]}
+    
     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()
+        return {"messages": [AIMessage(
+            content=f"Synthesis Error: {str(e)}. Please try again later."
+        )]}
 
 # ------------------------------
-# 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)
+# Professional UI Interface
 # ------------------------------
 def main():
     st.set_page_config(
-        page_title="AI Research & Development Assistant",
+        page_title="Research Assistant Pro",
         layout="wide",
         initial_sidebar_state="expanded"
     )
-
+    
+    # Dark theme implementation
     st.markdown("""
     <style>
     .stApp {
-        background-color: #1a1a1a;
+        background-color: #0f1114;
         color: #ffffff;
     }
-    
-    .stTextArea textarea {
-        background-color: #2d2d2d !important;
+    .stTextInput input, .stTextArea textarea {
+        background-color: #1e1e24 !important;
         color: #ffffff !important;
     }
-    
-    .stButton > button {
-        background-color: #4CAF50;
-        color: white;
-        transition: all 0.3s;
+    .stButton>button {
+        background: #2563eb;
+        transition: all 0.2s;
     }
-    
-    .stButton > button:hover {
-        background-color: #45a049;
+    .stButton>button:hover {
+        background: #1d4ed8;
         transform: scale(1.02);
     }
-    
-    .data-box {
-        background-color: #2d2d2d;
-        border-left: 5px solid #2196F3;
-    }
-    
-    .dev-box {
-        border-left: 5px solid #4CAF50;
-    }
-    
-    .st-expander {
-        background-color: #2d2d2d;
-        border: 1px solid #3d3d3d;
+    .result-card {
+        background: #1a1a1f;
+        border-radius: 8px;
+        padding: 1.5rem;
+        margin: 1rem 0;
     }
     </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?")
-
+    
+    st.title("🔍 Research Assistant Pro")
+    st.write("Advanced AI-Powered Research Analysis")
+    
     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"}})
-                        
-                        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 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..."):
+                    assistant = ResearchAssistant()
+                    result = assistant.workflow.invoke({"messages": [
+                        HumanMessage(content=query)
+                    ]})
+                    
+                    with st.expander("Analysis Details", expanded=True):
+                        st.markdown(f"""
+                        <div class="result-card">
+                            {result['messages'][-1].content}
+                        </div>
+                        """, unsafe_allow_html=True)
+    
     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?
-        """)
+        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
+            """)
 
 if __name__ == "__main__":
+    if not Config.API_KEY:
+        st.error("""
+        🔑 Configuration Required:
+        Set DEEPSEEK_API_KEY in environment variables
+        """)
+        st.stop()
     main()