agent.py

from typing import TypedDict, Annotated, Sequence
from langchain_core.messages import BaseMessage, HumanMessage
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.graph import END, StateGraph
from langgraph.prebuilt import ToolNode
from langchain.tools import DuckDuckGoSearchResults
from langchain_community.utilities import WikipediaAPIWrapper
from langchain.agents import create_tool_calling_agent
from langchain.agents import AgentExecutor
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
import operator
import json

load dotenv()

# Define the agent state
class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    sender: str

# Initialize tools
@tool
def wikipedia_search(query: str) -> str:
    """Search Wikipedia for information."""
    return WikipediaAPIWrapper().run(query)

@tool
def web_search(query: str, num_results: int = 3) -> list:
    """Search the web for current information."""
    return DuckDuckGoSearchResults(num_results=num_results).run(query)

@tool
def calculate(expression: str) -> str:
    """Evaluate mathematical expressions."""
    from langchain.chains import LLMMathChain
    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
    return LLMMathChain(llm=llm).run(expression)

class AdvancedAIAgent:
    def __init__(self, model_name="gpt-4-turbo"):
        # Initialize tools and LLM
        self.tools = [wikipedia_search, web_search, calculate]
        self.llm = ChatOpenAI(model=model_name, temperature=0.7)
        
        # Create the agent
        self.agent = self._create_agent()
        
        # Build the graph workflow
        self.workflow = self._build_graph()

    def _create_agent(self) -> AgentExecutor:
        """Create the agent with tools and prompt"""
        prompt = ChatPromptTemplate.from_messages([
            ("system", "You are a helpful AI assistant. Use tools when needed."),
            MessagesPlaceholder(variable_name="messages"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ])
        
        agent = create_tool_calling_agent(self.llm, self.tools, prompt)
        return AgentExecutor(agent=agent, tools=self.tools, verbose=True)

    def _build_graph(self):
        """Build the LangGraph workflow"""
        workflow = StateGraph(AgentState)
        
        # Define nodes
        workflow.add_node("agent", self._call_agent)
        workflow.add_node("tools", ToolNode(self.tools))  # Using ToolNode instead of ToolExecutor
        
        # Define edges
        workflow.set_entry_point("agent")
        workflow.add_conditional_edges(
            "agent",
            self._should_continue,
            {
                "continue": "tools",
                "end": END
            }
        )
        workflow.add_edge("tools", "agent")
        
        return workflow.compile()

    def _call_agent(self, state: AgentState):
        """Execute the agent"""
        response = self.agent.invoke({"messages": state["messages"]})
        return {"messages": [response["output"]]}

    def _should_continue(self, state: AgentState):
        """Determine if the workflow should continue"""
        last_message = state["messages"][-1]
        
        # If no tool calls, end
        if not last_message.additional_kwargs.get("tool_calls"):
            return "end"
        return "continue"

    def __call__(self, query: str) -> dict:
        """Process a user query"""
        # Initialize state
        state = AgentState(messages=[HumanMessage(content=query)], sender="user")
        
        # Execute the workflow
        for output in self.workflow.stream(state):
            for key, value in output.items():
                if key == "messages":
                    for message in value:
                        if isinstance(message, BaseMessage):
                            return {
                                "response": message.content,
                                "sources": self._extract_sources(state["messages"]),
                                "steps": self._extract_steps(state["messages"])
                            }

    def _extract_sources(self, messages: Sequence[BaseMessage]) -> list:
        """Extract sources from tool messages"""
        return [
            f"{msg.additional_kwargs.get('name', 'unknown')}: {msg.content}"
            for msg in messages
            if hasattr(msg, 'additional_kwargs') and 'name' in msg.additional_kwargs
        ]

    def _extract_steps(self, messages: Sequence[BaseMessage]) -> list:
        """Extract reasoning steps"""
        steps = []
        for msg in messages:
            if hasattr(msg, 'additional_kwargs') and 'tool_calls' in msg.additional_kwargs:
                for call in msg.additional_kwargs['tool_calls']:
                    steps.append(f"Used {call['function']['name']}: {call['function']['arguments']}")
        return steps