update

pyproject.toml

@@ -33,6 +33,7 @@ ruff = "^0.11.12"
 [tool.ruff]
 line-length = 120
 target-version = "py312"
+exclude = ["notebooks/", "*.ipynb"]
 
 [tool.ruff.lint]
 select = ["I", "E4", "E7", "E9", "F", "B", "UP"]

src/agents/langgraph_agent.py

@@ -1,22 +1,22 @@
-from typing import Annotated, Sequence, TypedDict, Literal
-from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, ToolMessage
-from langchain_core.runnables import RunnableConfig
-from langgraph.graph.message import add_messages
-from langgraph.graph import StateGraph, END
-from pydantic import BaseModel, Field
+from collections.abc import Sequence
+from typing import Annotated, Literal, TypedDict
+
 from langchain.chat_models import init_chat_model
-import json
 
 # Import tools
-from langchain_community.tools import DuckDuckGoSearchRun
+from langchain_community.tools import DuckDuckGoSearchRun, WikipediaQueryRun
+from langchain_community.tools.arxiv import ArxivQueryRun
 from langchain_community.tools.pubmed.tool import PubmedQueryRun
 from langchain_community.tools.semanticscholar.tool import SemanticScholarQueryRun
-from langchain_community.tools.arxiv import ArxivQueryRun
 from langchain_community.tools.wikidata.tool import WikidataQueryRun
-from langchain_community.tools import WikipediaQueryRun
 from langchain_community.utilities import WikipediaAPIWrapper
-from langchain_experimental.utilities import PythonREPL
+from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, ToolMessage
+from langchain_core.runnables import RunnableConfig
 from langchain_core.tools import Tool
+from langchain_experimental.utilities import PythonREPL
+from langgraph.graph import END, StateGraph
+from langgraph.graph.message import add_messages
+from pydantic import BaseModel, Field
 
 # Set up tools
 python_repl = PythonREPL()
@@ -34,7 +34,7 @@ tools = [
     ArxivQueryRun(),
     WikidataQueryRun(),
     WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper()),
-    repl_tool
+    repl_tool,
 ]
 
 # Initialize Gemini model
@@ -44,39 +44,46 @@ model_with_tools = model.bind_tools(tools)
 # Create tools lookup
 tools_by_name = {tool.name: tool for tool in tools}
 
+
 # Pydantic models for structured output
 class ToolSufficiencyResponse(BaseModel):
     """Response for tool sufficiency check"""
+
     sufficient: bool = Field(description="Whether the available tools are sufficient to answer the question")
     reasoning: str = Field(description="Brief reasoning for the decision")
 
+
 class FinalAnswer(BaseModel):
     """Final answer structure"""
+
     answer: str = Field(description="The comprehensive answer to the user's question")
     confidence: Literal["high", "medium", "low"] = Field(description="Confidence level in the answer")
     sources_used: list[str] = Field(description="List of tools/sources that were used to generate the answer")
 
+
 # Define graph state
 class AgentState(TypedDict):
     """The state of the agent."""
+
     messages: Annotated[Sequence[BaseMessage], add_messages]
     llm_call_count: int
     max_llm_calls: int
 
+
 # Node functions
 def check_tool_sufficiency(state: AgentState, config: RunnableConfig):
     """Check if available tools are sufficient to answer the question"""
-    
+
     # Get the user's question
     user_message = None
     for msg in state["messages"]:
         if msg.type == "human":
             user_message = msg.content
             break
-    
+
     # Create system prompt for sufficiency check
     available_tools_desc = "\n".join([f"- {tool.name}: {tool.description}" for tool in tools])
-    
+
     system_prompt = f"""You are an AI assistant that needs to determine if the available tools are sufficient to answer a user's question.
 
 Available tools:
@@ -93,27 +100,22 @@ Be generous in your assessment - if there's a reasonable path to answer the ques
 
     # Use structured output for sufficiency check
     structured_model = model.with_structured_output(ToolSufficiencyResponse)
-    
-    messages = [
-        SystemMessage(content=system_prompt),
-        HumanMessage(content=f"Question to analyze: {user_message}")
-    ]
-    
+
+    messages = [SystemMessage(content=system_prompt), HumanMessage(content=f"Question to analyze: {user_message}")]
+
     response = structured_model.invoke(messages, config)
-    
+
     # Add response to messages for context
     response_message = SystemMessage(
         content=f"Tool sufficiency check: {'Sufficient' if response.sufficient else 'Insufficient'}. Reasoning: {response.reasoning}"
     )
-    
-    return {
-        "messages": [response_message],
-        "tool_sufficiency": response.sufficient
-    }
+
+    return {"messages": [response_message], "tool_sufficiency": response.sufficient}
+
 
 def call_model(state: AgentState, config: RunnableConfig):
     """Call the model (ReAct agent LLM node)"""
-    
+
     system_prompt = SystemMessage(
         content="""You are a helpful AI assistant with access to various tools. Use the tools available to you to answer the user's question comprehensively.
 
@@ -124,22 +126,20 @@ Think step by step:
 
 Be thorough but efficient with your tool usage."""
     )
-    
+
     response = model_with_tools.invoke([system_prompt] + state["messages"], config)
-    
+
     # Increment LLM call count
     new_count = state.get("llm_call_count", 0) + 1
-    
-    return {
-        "messages": [response],
-        "llm_call_count": new_count
-    }
+
+    return {"messages": [response], "llm_call_count": new_count}
+
 
 def tool_node(state: AgentState):
     """Execute tools based on the last message's tool calls"""
     outputs = []
     last_message = state["messages"][-1]
-    
+
     for tool_call in last_message.tool_calls:
         try:
             tool_result = tools_by_name[tool_call["name"]].invoke(tool_call["args"])
@@ -158,12 +158,13 @@ def tool_node(state: AgentState):
                     tool_call_id=tool_call["id"],
                 )
             )
-    
+
     return {"messages": outputs}
 
+
 def final_answer_node(state: AgentState, config: RunnableConfig):
     """Generate final structured answer based on conversation history"""
-    
+
     system_prompt = SystemMessage(
         content="""You are tasked with providing a final, comprehensive answer based on the conversation history and tool usage.
 
@@ -174,29 +175,27 @@ Analyze all the information gathered from the tools and provide:
 
 Be honest about limitations and indicate your confidence level appropriately."""
     )
-    
+
     # Get the original user question
     user_question = None
     for msg in state["messages"]:
         if msg.type == "human":
             user_question = msg.content
             break
-    
+
     # Create structured output model
     structured_model = model.with_structured_output(FinalAnswer)
-    
+
     messages = [
         system_prompt,
         HumanMessage(content=f"Original question: {user_question}"),
-        SystemMessage(content="Based on the following conversation history, provide your final answer:")
+        SystemMessage(content="Based on the following conversation history, provide your final answer:"),
     ] + state["messages"]
-    
+
     response = structured_model.invoke(messages, config)
-    
-    return {
-        "messages": [SystemMessage(content=f"Final Answer: {response.answer}")],
-        "final_answer": response
-    }
+
+    return {"messages": [SystemMessage(content=f"Final Answer: {response.answer}")], "final_answer": response}
+
 
 # Edge functions
 def should_continue_sufficiency(state: AgentState):
@@ -209,95 +208,84 @@ def should_continue_sufficiency(state: AgentState):
             return "insufficient"
     return "insufficient"  # Default to insufficient if unclear
 
+
 def should_continue_react(state: AgentState):
     """Decide whether to continue with ReAct loop or move to final answer"""
     messages = state["messages"]
     last_message = messages[-1]
     llm_call_count = state.get("llm_call_count", 0)
     max_calls = state.get("max_llm_calls", 4)
-    
+
     # If we've reached the maximum number of LLM calls, force stop
     if llm_call_count >= max_calls:
         return "final_answer"
-    
+
     # If there are no tool calls, we're done with ReAct loop
-    if not hasattr(last_message, 'tool_calls') or not last_message.tool_calls:
+    if not hasattr(last_message, "tool_calls") or not last_message.tool_calls:
         return "final_answer"
-    
+
     # Otherwise continue with tools
     return "continue"
 
+
 # Build the graph
 def create_react_agent_graph():
     """Create and return the compiled ReAct agent graph"""
-    
+
     workflow = StateGraph(AgentState)
-    
+
     # Add nodes
     workflow.add_node("check_sufficiency", check_tool_sufficiency)
     workflow.add_node("agent", call_model)
     workflow.add_node("tools", tool_node)
     workflow.add_node("final_answer", final_answer_node)
-    
+
     # Set entry point
     workflow.set_entry_point("check_sufficiency")
-    
+
     # Add conditional edge from sufficiency check
     workflow.add_conditional_edges(
-        "check_sufficiency",
-        should_continue_sufficiency,
-        {
-            "sufficient": "agent",
-            "insufficient": END
-        }
+        "check_sufficiency", should_continue_sufficiency, {"sufficient": "agent", "insufficient": END}
     )
-    
+
     # Add conditional edge from agent
     workflow.add_conditional_edges(
-        "agent",
-        should_continue_react,
-        {
-            "continue": "tools",
-            "final_answer": "final_answer"
-        }
+        "agent", should_continue_react, {"continue": "tools", "final_answer": "final_answer"}
     )
-    
+
     # Add edge from tools back to agent
     workflow.add_edge("tools", "agent")
-    
+
     # Add edge from final_answer to END
     workflow.add_edge("final_answer", END)
-    
+
     return workflow.compile()
 
+
 # Helper function for running the agent
 def run_agent(question: str, max_llm_calls: int = 4):
     """Run the ReAct agent with a question"""
-    
+
     graph = create_react_agent_graph()
-    
-    initial_state = {
-        "messages": [HumanMessage(content=question)],
-        "llm_call_count": 0,
-        "max_llm_calls": max_llm_calls
-    }
-    
+
+    initial_state = {"messages": [HumanMessage(content=question)], "llm_call_count": 0, "max_llm_calls": max_llm_calls}
+
     # Stream the execution
     print(f"Question: {question}")
     print("=" * 50)
-    
+
     for step in graph.stream(initial_state):
         for node, output in step.items():
             print(f"\n--- {node.upper()} ---")
             if "messages" in output and output["messages"]:
                 for msg in output["messages"]:
-                    if hasattr(msg, 'content'):
+                    if hasattr(msg, "content"):
                         print(f"{msg.__class__.__name__}: {msg.content}")
-                    elif hasattr(msg, 'tool_calls') and msg.tool_calls:
+                    elif hasattr(msg, "tool_calls") and msg.tool_calls:
                         print(f"Tool calls: {[tc['name'] for tc in msg.tool_calls]}")
-            
+
             if "final_answer" in output:
-                print(f"\nFINAL STRUCTURED ANSWER:")
+                print("\nFINAL STRUCTURED ANSWER:")
                 print(f"Answer: {output['final_answer'].answer}")
                 print(f"Confidence: {output['final_answer'].confidence}")
                 print(f"Sources: {output['final_answer'].sources_used}")