Update app.py

app.py

@@ -1,6 +1,6 @@
-from langchain_openai import OpenAIEmbeddings
-from langchain_community.vectorstores import Chroma
-from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
+from langchain_openai import OpenAIEmbeddings  # Updated import path
+from langchain.vectorstores import Chroma
+from langchain.schema import HumanMessage, AIMessage
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langgraph.graph import END, StateGraph, START
 from langgraph.prebuilt import ToolNode
@@ -31,21 +31,16 @@ development_texts = [
 # --------------------------
 # Process the Data
 # --------------------------
-# Text splitting settings
 splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=10)
 
-# Generate Document objects from text
 research_docs = splitter.create_documents(research_texts)
 development_docs = splitter.create_documents(development_texts)
 
-# Create vector embeddings
 embeddings = OpenAIEmbeddings(
-    model="text-embedding-3-large",
-    # For text-embedding-3 class models, you can specify dimensions if needed.
-    # dimensions=1024
+    model="text-embedding-3-large"
+    # dimensions=1024  # Uncomment if needed
 )
 
-# Create vector stores
 research_vectorstore = Chroma.from_documents(
     documents=research_docs,
     embedding=embeddings,
@@ -58,7 +53,6 @@ development_vectorstore = Chroma.from_documents(
     collection_name="development_collection"
 )
 
-# Create retrievers from the vector stores
 research_retriever = research_vectorstore.as_retriever()
 development_retriever = development_vectorstore.as_retriever()
 
@@ -66,9 +60,9 @@ development_retriever = development_vectorstore.as_retriever()
 # Create Retriever Tools
 # --------------------------
 research_tool = create_retriever_tool(
-    research_retriever,  # Retriever object
-    "research_db_tool",  # Name of the tool to create
-    "Search information from the research database."  # Description of the tool
+    research_retriever,
+    "research_db_tool",
+    "Search information from the research database."
 )
 
 development_tool = create_retriever_tool(
@@ -77,25 +71,19 @@ development_tool = create_retriever_tool(
     "Search information from the development database."
 )
 
-# Combine the tools into a list
 tools = [research_tool, development_tool]
 
 # --------------------------
 # Define the Agent Function
 # --------------------------
 class AgentState(TypedDict):
-    messages: Annotated[Sequence[AIMessage | HumanMessage | ToolMessage], add_messages]
+    messages: Annotated[Sequence[AIMessage | HumanMessage], 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>
@@ -108,7 +96,7 @@ Otherwise, just answer directly.
 
     headers = {
         "Accept": "application/json",
-        "Authorization": f"Bearer sk-1cddf19f9dc4466fa3ecea6fe10abec0",
+        "Authorization": "Bearer sk-1cddf19f9dc4466fa3ecea6fe10abec0",
         "Content-Type": "application/json"
     }
     
@@ -130,16 +118,13 @@ Otherwise, just answer directly.
         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()
-            # Use direct call to research retriever
             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()
-            # Use direct call to development retriever
             results = development_retriever.invoke(query)
             return {"messages": [AIMessage(content=f'Action: development_db_tool\n{{"query": "{query}"}}\n\nResults: {str(results)}')]}
 
@@ -156,7 +141,6 @@ def simple_grade_documents(state: AgentState):
     last_message = messages[-1]
     print("Evaluating message:", last_message.content)
     
-    # Check if the content contains retrieved documents
     if "Results: [Document" in last_message.content:
         print("---DOCS FOUND, GO TO GENERATE---")
         return "generate"
@@ -170,10 +154,9 @@ def simple_grade_documents(state: AgentState):
 def generate(state: AgentState):
     print("---GENERATE FINAL ANSWER---")
     messages = state["messages"]
-    question = messages[0].content if isinstance(messages[0], tuple) else messages[0].content
+    question = messages[0].content if not isinstance(messages[0], tuple) else messages[0][1]
     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: [")
@@ -182,7 +165,7 @@ def generate(state: AgentState):
 
     headers = {
         "Accept": "application/json",
-        "Authorization": f"Bearer sk-1cddf19f9dc4466fa3ecea6fe10abec0",
+        "Authorization": "Bearer sk-1cddf19f9dc4466fa3ecea6fe10abec0",
         "Content-Type": "application/json"
     }
     
@@ -194,10 +177,7 @@ 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
     }
@@ -227,21 +207,17 @@ def rewrite(state: AgentState):
 
     headers = {
         "Accept": "application/json",
-        "Authorization": f"Bearer sk-1cddf19f9dc4466fa3ecea6fe10abec0",
+        "Authorization": "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,
@@ -249,9 +225,6 @@ def rewrite(state: AgentState):
         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)
@@ -270,50 +243,37 @@ def custom_tools_condition(state: AgentState):
     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
+    return "tools" if tools_pattern.match(content) else END
 
 # --------------------------
 # LangGraph Workflow Setup
 # --------------------------
 workflow = StateGraph(AgentState)
 
-# Define the workflow 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 up the initial edge
 workflow.add_edge(START, "agent")
 
-# Conditional edge from agent to either retrieve (if tool is called) or END
 workflow.add_conditional_edges(
     "agent",
     custom_tools_condition,
-    {
-        "tools": "retrieve",
-        END: END
-    }
+    {"tools": "retrieve", END: END}
 )
 
-# After retrieval, decide to generate or rewrite based on document grading
 workflow.add_conditional_edges("retrieve", simple_grade_documents)
 workflow.add_edge("generate", END)
 workflow.add_edge("rewrite", "agent")
 
-# Compile the workflow to make it executable
 app = workflow.compile()
 
 # --------------------------
 # Process Question 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)
@@ -329,60 +289,37 @@ def main():
         initial_sidebar_state="expanded"
     )
 
-    # Custom CSS
     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 Data Display
     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)
 
-    # Main Content
     st.title("🤖 AI Research & Development Assistant")
     st.markdown("---")
 
-    # Query Input
-    query = st.text_area("Enter your question:", height=100, placeholder="e.g., What is the latest advancement in AI research?")
+    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...'):
-                    # Process query through workflow
                     events = process_question(query, app, {"configurable": {"thread_id": "1"}})
-                    
-                    # Display results
                     for event in events:
                         if 'agent' in event:
                             with st.expander("🔄 Processing Step", expanded=True):
@@ -412,4 +349,4 @@ def main():
         """)
 
 if __name__ == "__main__":
-    main()
+    main()