Update app.py

app.py

@@ -1,12 +1,13 @@
+# pip install -r requirements.txt
+
+# ------------------------------
+# Imports & Dependencies
+# ------------------------------
 from langchain_openai import OpenAIEmbeddings
 from langchain_community.vectorstores import Chroma
-from langchain_core.messages import HumanMessage, AIMessage
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langgraph.graph import StateGraph  # Only import StateGraph from here
-# Define START and END manually
-START = "__start__"
-END = "__end__"
-
+from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langgraph.graph import END, StateGraph, START
 from langgraph.prebuilt import ToolNode
 from langgraph.graph.message import add_messages
 from typing_extensions import TypedDict, Annotated
@@ -17,9 +18,9 @@ import streamlit as st
 import requests
 from langchain.tools.retriever import create_retriever_tool
 
-# --------------------------
-# Create Dummy Data
-# --------------------------
+# ------------------------------
+# 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",
@@ -32,41 +33,41 @@ development_texts = [
     "Product Y: In the Performance Optimization Stage Before Release"
 ]
 
-# --------------------------
-# Process the Data
-# --------------------------
+# ------------------------------
+# 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
+    # You can uncomment and set dimensions if needed:
+    # dimensions=1024
 )
-
 research_vectorstore = Chroma.from_documents(
     documents=research_docs,
     embedding=embeddings,
     collection_name="research_collection"
 )
-
 development_vectorstore = Chroma.from_documents(
     documents=development_docs,
     embedding=embeddings,
     collection_name="development_collection"
 )
-
 research_retriever = research_vectorstore.as_retriever()
 development_retriever = development_vectorstore.as_retriever()
 
-# --------------------------
-# Create Retriever Tools
-# --------------------------
+# ------------------------------
+# Creating Retriever Tools
+# ------------------------------
 research_tool = create_retriever_tool(
-    research_retriever,
-    "research_db_tool",
-    "Search information from the research database."
+    research_retriever,  # Retriever object
+    "research_db_tool",  # Name of the tool
+    "Search information from the research database."  # Tool description
 )
 
 development_tool = create_retriever_tool(
@@ -77,17 +78,22 @@ development_tool = create_retriever_tool(
 
 tools = [research_tool, development_tool]
 
-# --------------------------
-# Define the Agent Function
-# --------------------------
+# ------------------------------
+# Agent Function & Workflow Functions
+# ------------------------------
 class AgentState(TypedDict):
-    messages: Annotated[Sequence[AIMessage | HumanMessage], add_messages]
+    messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]
 
 def agent(state: AgentState):
     print("---CALL AGENT---")
     messages = state["messages"]
-    user_message = messages[0][1] if isinstance(messages[0], tuple) else messages[0].content
 
+    if isinstance(messages[0], tuple):
+        user_message = messages[0][1]
+    else:
+        user_message = messages[0].content
+
+    # Structure prompt for consistent text output
     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>
@@ -100,7 +106,7 @@ Otherwise, just answer directly.
 
     headers = {
         "Accept": "application/json",
-        "Authorization": "Bearer sk-1cddf19f9dc4466fa3ecea6fe10abec0",
+        "Authorization": f"Bearer sk-1cddf19f9dc4466fa3ecea6fe10abec0",
         "Content-Type": "application/json"
     }
     
@@ -114,13 +120,15 @@ Otherwise, just answer directly.
     response = requests.post(
         "https://api.deepseek.com/v1/chat/completions",
         headers=headers,
-        json=data
+        json=data,
+        verify=False
     )
     
     if response.status_code == 200:
         response_text = response.json()['choices'][0]['message']['content']
         print("Raw response:", response_text)
         
+        # Format the response into expected tool format
         if "SEARCH_RESEARCH:" in response_text:
             query = response_text.split("SEARCH_RESEARCH:")[1].strip()
             results = research_retriever.invoke(query)
@@ -136,9 +144,6 @@ Otherwise, just answer directly.
     else:
         raise Exception(f"API call failed: {response.text}")
 
-# --------------------------
-# Grading Function
-# --------------------------
 def simple_grade_documents(state: AgentState):
     messages = state["messages"]
     last_message = messages[-1]
@@ -151,15 +156,13 @@ def simple_grade_documents(state: AgentState):
         print("---NO DOCS FOUND, TRY REWRITE---")
         return "rewrite"
 
-# --------------------------
-# Generation Function
-# --------------------------
 def generate(state: AgentState):
     print("---GENERATE FINAL ANSWER---")
     messages = state["messages"]
-    question = messages[0].content if not isinstance(messages[0], tuple) else messages[0][1]
+    question = messages[0].content if isinstance(messages[0], tuple) else messages[0].content
     last_message = messages[-1]
 
+    # Extract the document content from the results
     docs = ""
     if "Results: [" in last_message.content:
         results_start = last_message.content.find("Results: [")
@@ -168,7 +171,7 @@ def generate(state: AgentState):
 
     headers = {
         "Accept": "application/json",
-        "Authorization": "Bearer sk-1cddf19f9dc4466fa3ecea6fe10abec0",
+        "Authorization": f"Bearer sk-1cddf19f9dc4466fa3ecea6fe10abec0",
         "Content-Type": "application/json"
     }
     
@@ -180,7 +183,10 @@ Focus on extracting and synthesizing the key findings from the research papers.
     
     data = {
         "model": "deepseek-chat",
-        "messages": [{"role": "user", "content": prompt}],
+        "messages": [{
+            "role": "user",
+            "content": prompt
+        }],
         "temperature": 0.7,
         "max_tokens": 1024
     }
@@ -189,7 +195,8 @@ Focus on extracting and synthesizing the key findings from the research papers.
     response = requests.post(
         "https://api.deepseek.com/v1/chat/completions",
         headers=headers,
-        json=data
+        json=data,
+        verify=False
     )
     
     if response.status_code == 200:
@@ -199,9 +206,6 @@ Focus on extracting and synthesizing the key findings from the research papers.
     else:
         raise Exception(f"API call failed: {response.text}")
 
-# --------------------------
-# Rewrite Function
-# --------------------------
 def rewrite(state: AgentState):
     print("---REWRITE QUESTION---")
     messages = state["messages"]
@@ -209,23 +213,31 @@ def rewrite(state: AgentState):
 
     headers = {
         "Accept": "application/json",
-        "Authorization": "Bearer sk-1cddf19f9dc4466fa3ecea6fe10abec0",
+        "Authorization": f"Bearer sk-1cddf19f9dc4466fa3ecea6fe10abec0",
         "Content-Type": "application/json"
     }
     
     data = {
         "model": "deepseek-chat",
-        "messages": [{"role": "user", "content": f"Rewrite this question to be more specific and clearer: {original_question}"}],
+        "messages": [{
+            "role": "user",
+            "content": f"Rewrite this question to be more specific and clearer: {original_question}"
+        }],
         "temperature": 0.7,
         "max_tokens": 1024
     }
     
+    print("Sending rewrite request...")
     response = requests.post(
         "https://api.deepseek.com/v1/chat/completions",
         headers=headers,
-        json=data
+        json=data,
+        verify=False
     )
     
+    print("Status Code:", response.status_code)
+    print("Response:", response.text)
+    
     if response.status_code == 200:
         response_text = response.json()['choices'][0]['message']['content']
         print("Rewritten question:", response_text)
@@ -233,9 +245,6 @@ def rewrite(state: AgentState):
     else:
         raise Exception(f"API call failed: {response.text}")
 
-# --------------------------
-# Tools Decision Function
-# --------------------------
 tools_pattern = re.compile(r"Action: .*")
 
 def custom_tools_condition(state: AgentState):
@@ -244,11 +253,15 @@ def custom_tools_condition(state: AgentState):
     content = last_message.content
 
     print("Checking tools condition:", content)
-    return "tools" if tools_pattern.match(content) else END
-
-# --------------------------
-# LangGraph Workflow Setup
-# --------------------------
+    if tools_pattern.match(content):
+        print("Moving to retrieve...")
+        return "tools"
+    print("Moving to END...")
+    return END
+
+# ------------------------------
+# Workflow Configuration using LangGraph
+# ------------------------------
 workflow = StateGraph(AgentState)
 
 workflow.add_node("agent", agent)
@@ -262,7 +275,10 @@ workflow.add_edge(START, "agent")
 workflow.add_conditional_edges(
     "agent",
     custom_tools_condition,
-    {"tools": "retrieve", END: END}
+    {
+        "tools": "retrieve",
+        END: END
+    }
 )
 
 workflow.add_conditional_edges("retrieve", simple_grade_documents)
@@ -271,18 +287,16 @@ workflow.add_edge("rewrite", "agent")
 
 app = workflow.compile()
 
-# --------------------------
-# Process Question Function
-# --------------------------
-def process_question(user_question):
+def process_question(user_question, app, config):
+    """Process user question through the workflow"""
     events = []
-    for event in app.stream({"messages": [("user", user_question)]}):
+    for event in app.stream({"messages": [("user", user_question)]}, config):
         events.append(event)
     return events
 
-# --------------------------
-# Streamlit Application
-# --------------------------
+# ------------------------------
+# Streamlit App UI
+# ------------------------------
 def main():
     st.set_page_config(
         page_title="AI Research & Development Assistant",
@@ -290,21 +304,40 @@ def main():
         initial_sidebar_state="expanded"
     )
 
+    # Custom CSS for styling
     st.markdown("""
     <style>
-    .stApp { background-color: #f8f9fa; }
-    .stButton > button { width: 100%; margin-top: 20px; }
-    .data-box { padding: 20px; border-radius: 10px; margin: 10px 0; }
-    .research-box { background-color: #e3f2fd; border-left: 5px solid #1976d2; }
-    .dev-box { background-color: #e8f5e9; border-left: 5px solid #43a047; }
+    .stApp {
+        background-color: #f8f9fa;
+    }
+    .stButton > button {
+        width: 100%;
+        margin-top: 20px;
+    }
+    .data-box {
+        padding: 20px;
+        border-radius: 10px;
+        margin: 10px 0;
+    }
+    .research-box {
+        background-color: #e3f2fd;
+        border-left: 5px solid #1976d2;
+    }
+    .dev-box {
+        background-color: #e8f5e9;
+        border-left: 5px solid #43a047;
+    }
     </style>
     """, unsafe_allow_html=True)
 
+    # Sidebar with available data
     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">{text}</div>', unsafe_allow_html=True)
+            
         st.subheader("Development Database")
         for text in development_texts:
             st.markdown(f'<div class="data-box dev-box">{text}</div>', unsafe_allow_html=True)
@@ -312,15 +345,16 @@ def main():
     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?")
+    # Query input box
+    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:
                 with st.spinner('Processing your question...'):
-                    events = process_question(query)
+                    events = process_question(query, app, {"configurable": {"thread_id": "1"}})
+                    
                     for event in events:
                         if 'agent' in event:
                             with st.expander("🔄 Processing Step", expanded=True):