Update app.py

app.py

@@ -1,79 +1,184 @@
-import streamlit as st
-from deep_research_system import run_deep_research_system
-import traceback
-import streamlit.components.v1 as components
-
-# Set page configuration
-st.set_page_config(page_title="Deep Research AI", layout="centered")
-
-# Title and instructions
-st.title("Deep Research AI Agentic System")
-st.write("Enter a question below to get the latest insights from web research. Click 'Reset' to start over.")
-
-# Initialize session state
-if "show_reset_button" not in st.session_state:
-    st.session_state.show_reset_button = False
-if "question" not in st.session_state:
-    st.session_state.question = ""
-if "reset_triggered" not in st.session_state:
-    st.session_state.reset_triggered = False
-
-# JavaScript to clear the input field
-clear_input_js = """
-<script>
-    const input = window.parent.document.querySelector('input[aria-label="Your Question"]');
-    if (input) {
-        input.value = '';
-    }
-</script>
-"""
-
-# Use a form to manage the question input and submission
-with st.form(key="question_form"):
-    st.session_state.question = st.text_input(
-        "Your Question",
-        placeholder="e.g., What are the latest advancements in quantum computing?",
-        value=st.session_state.question,
-        key="question_input"
-    )
-    submit_button = st.form_submit_button("Get Answer")
-
-# Process the form submission
-if submit_button:
-    if st.session_state.question:
-        st.write(f"Research Agent: Searching for '{st.session_state.question}'...")
-        try:
-            with st.spinner("Gathering research data..."):
-                answer = run_deep_research_system(st.session_state.question)
-            st.write("Research Agent: Found 5 relevant sources.")
-            st.write("Answer Drafter Agent: Drafted the final answer.")
-            st.write("**Final Answer:**")
-            st.write(answer)
-            st.session_state.show_reset_button = True
-        except Exception as e:
-            st.error(f"An error occurred: {str(e)}\n{traceback.format_exc()}")
-            st.session_state.show_reset_button = True
-    else:
-        st.warning("Please enter a question!")
-    # Reset the trigger flag after submission
-    st.session_state.reset_triggered = False
-
-# Function to clear the input state
-def clear_input():
-    st.session_state.show_reset_button = False
-    st.session_state.question = ""
-    st.session_state.pop("question_input", None)
-    st.session_state.pop("question_form", None)
-    st.session_state.reset_triggered = True
-
-# Show Reset button only if show_reset_button is True
-if st.session_state.show_reset_button:
-    if st.button("Reset"):
-        # Clear the input and reset state
-        clear_input()
-        # Refresh the webpage
-        st.rerun()
-
-# Execute JavaScript to clear the input field if reset was triggered
-if st.session_state.reset_triggered:
+
+import os
+import traceback
+import streamlit as st
+from dotenv import load_dotenv
+from tavily import TavilyClient
+from langchain_openai import ChatOpenAI
+from typing import Dict, TypedDict, List
+from langgraph.graph import StateGraph, END
+import streamlit.components.v1 as components
+from langchain_core.prompts import ChatPromptTemplate
+
+# Load environment variables
+load_dotenv()
+
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
+
+# Initialize the LLM (using OpenAI as an example, replace with your preferred model)
+llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+
+# Initialize Tavily client for web search
+tavily_client = TavilyClient(api_key=os.environ["TAVILY_API_KEY"])
+
+# Define the shared state for the agents
+class ResearchState(TypedDict):
+    query: str
+    research_data: List[Dict]
+    final_answer: str
+
+# Research Agent: Crawls the web and gathers information
+def research_agent(state: ResearchState) -> ResearchState:
+    query = state["query"]
+    print(f"Research Agent: Searching for '{query}'...")
+
+    # Use Tavily to search the web
+    search_results = tavily_client.search(query, max_results=5)
+
+    # Extract relevant information from search results
+    research_data = []
+    for result in search_results["results"]:
+        research_data.append({
+            "title": result["title"],
+            "url": result["url"],
+            "content": result["content"][:500]  # Limit content length for brevity
+        })
+
+    print(f"Research Agent: Found {len(research_data)} relevant sources.")
+    return {"research_data": research_data}
+
+# Answer Drafter Agent: Processes research data and drafts a response
+def answer_drafter_agent(state: ResearchState) -> ResearchState:
+    research_data = state["research_data"]
+    query = state["query"]
+
+    # Create a prompt for the answer drafter
+    prompt = ChatPromptTemplate.from_template(
+        """
+        You are an expert at drafting concise and accurate answers. Based on the following research data, provide a clear and informative response to the query: "{query}".
+        Research Data:
+        {research_data}
+        Provide a well-structured answer in 3-5 sentences, citing the sources where relevant.
+        """
+    )
+
+    # Format the research data for the prompt
+    research_text = "\n".join([f"- {item['title']}: {item['content']} (Source: {item['url']})" for item in research_data])
+    chain = prompt | llm
+
+    # Generate the final answer
+    response = chain.invoke({"query": query, "research_data": research_text})
+    final_answer = response.content
+
+    print("Answer Drafter Agent: Drafted the final answer.")
+    return {"final_answer": final_answer}
+
+# Define the LangGraph workflow
+def create_workflow():
+    workflow = StateGraph(ResearchState)
+
+    # Add nodes for each agent
+    workflow.add_node("research_agent", research_agent)
+    workflow.add_node("answer_drafter_agent", answer_drafter_agent)
+
+    # Define the flow: Research Agent -> Answer Drafter Agent -> End
+    workflow.add_edge("research_agent", "answer_drafter_agent")
+    workflow.add_edge("answer_drafter_agent", END)
+
+    # Set the entry point
+    workflow.set_entry_point("research_agent")
+
+    return workflow.compile()
+
+# Main function to run the system
+def run_deep_research_system(query: str) -> str:
+    # Initialize the workflow
+    app = create_workflow()
+
+    # Initial state
+    initial_state = {
+        "query": query,
+        "research_data": [],
+        "final_answer": ""
+    }
+
+    # Run the workflow
+    final_state = app.invoke(initial_state)
+
+    return final_state["final_answer"]
+
+# Set page configuration
+st.set_page_config(page_title="Deep Research AI", layout="centered")
+
+# Title and instructions
+st.title("Deep Research AI Agentic System")
+st.write("Enter a question below to get the latest insights from web research. Click 'Reset' to start over.")
+
+# Initialize session state
+if "show_reset_button" not in st.session_state:
+    st.session_state.show_reset_button = False
+if "question" not in st.session_state:
+    st.session_state.question = ""
+if "reset_triggered" not in st.session_state:
+    st.session_state.reset_triggered = False
+
+# JavaScript to clear the input field
+clear_input_js = """
+<script>
+    const input = window.parent.document.querySelector('input[aria-label="Your Question"]');
+    if (input) {
+        input.value = '';
+    }
+</script>
+"""
+
+# Use a form to manage the question input and submission
+with st.form(key="question_form"):
+    st.session_state.question = st.text_input(
+        "Your Question",
+        placeholder="e.g., What are the latest advancements in quantum computing?",
+        value=st.session_state.question,
+        key="question_input"
+    )
+    submit_button = st.form_submit_button("Get Answer")
+
+# Process the form submission
+if submit_button:
+    if st.session_state.question:
+        st.write(f"Research Agent: Searching for '{st.session_state.question}'...")
+        try:
+            with st.spinner("Gathering research data..."):
+                answer = run_deep_research_system(st.session_state.question)
+            st.write("Research Agent: Found 5 relevant sources.")
+            st.write("Answer Drafter Agent: Drafted the final answer.")
+            st.write("**Final Answer:**")
+            st.write(answer)
+            st.session_state.show_reset_button = True
+        except Exception as e:
+            st.error(f"An error occurred: {str(e)}\n{traceback.format_exc()}")
+            st.session_state.show_reset_button = True
+    else:
+        st.warning("Please enter a question!")
+    # Reset the trigger flag after submission
+    st.session_state.reset_triggered = False
+
+# Function to clear the input state
+def clear_input():
+    st.session_state.show_reset_button = False
+    st.session_state.question = ""
+    st.session_state.pop("question_input", None)
+    st.session_state.pop("question_form", None)
+    st.session_state.reset_triggered = True
+
+# Show Reset button only if show_reset_button is True
+if st.session_state.show_reset_button:
+    if st.button("Reset"):
+        # Clear the input and reset state
+        clear_input()
+        # Refresh the webpage
+        st.rerun()
+
+# Execute JavaScript to clear the input field if reset was triggered
+if st.session_state.reset_triggered:
     components.html(clear_input_js, height=0)