wehub-resource-sync
147 lines
4.5 KiB
Python
147 lines
4.5 KiB
Python
"""
|
|
This is the main entry point for the agent.
|
|
It defines the workflow graph, state, tools, nodes and edges.
|
|
"""
|
|
|
|
from typing_extensions import Literal, TypedDict, Dict, List, Any, Union, Optional
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain_core.runnables import RunnableConfig
|
|
from langgraph.graph import StateGraph, END
|
|
from langgraph.checkpoint.memory import MemorySaver
|
|
from langgraph.types import Command
|
|
from copilotkit import CopilotKitState
|
|
from langchain_mcp_adapters.client import MultiServerMCPClient
|
|
from langgraph.prebuilt import create_react_agent
|
|
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
|
|
import os
|
|
|
|
|
|
# Define the connection type structures
|
|
class StdioConnection(TypedDict):
|
|
command: str
|
|
args: List[str]
|
|
transport: Literal["stdio"]
|
|
|
|
|
|
class SSEConnection(TypedDict):
|
|
url: str
|
|
transport: Literal["sse"]
|
|
|
|
|
|
# Type for MCP configuration
|
|
MCPConfig = Dict[str, Union[StdioConnection, SSEConnection]]
|
|
|
|
|
|
class AgentState(CopilotKitState):
|
|
"""
|
|
Here we define the state of the agent
|
|
|
|
In this instance, we're inheriting from CopilotKitState, which will bring in
|
|
the CopilotKitState fields. We're also adding a custom field, `mcp_config`,
|
|
which will be used to configure MCP services for the agent.
|
|
"""
|
|
|
|
# Define mcp_config as an optional field without skipping validation
|
|
mcp_config: Optional[MCPConfig]
|
|
|
|
|
|
# Default MCP configuration to use when no configuration is provided in the state
|
|
# Uses relative paths that will work within the project structure
|
|
DEFAULT_MCP_CONFIG: MCPConfig = {
|
|
"math": {
|
|
"command": "python",
|
|
# Use a relative path that will be resolved based on the current working directory
|
|
"args": [os.path.join(os.path.dirname(__file__), "..", "math_server.py")],
|
|
"transport": "stdio",
|
|
},
|
|
}
|
|
|
|
# Define a custom ReAct prompt that encourages the use of multiple tools
|
|
MULTI_TOOL_REACT_PROMPT = ChatPromptTemplate.from_messages(
|
|
[
|
|
(
|
|
"system",
|
|
"""You are an assistant that can use multiple tools to solve problems.
|
|
You should use a step-by-step approach, using as many tools as needed to find the complete answer.
|
|
Don't hesitate to call different tools sequentially if that helps reach a better solution.
|
|
|
|
You have access to the following tools:
|
|
|
|
{{tools}}
|
|
|
|
To use a tool, please use the following format:
|
|
```
|
|
Thought: I need to use a tool to help with this.
|
|
Action: tool_name
|
|
Action Input: the input to the tool
|
|
```
|
|
|
|
The observation will be returned in the following format:
|
|
```
|
|
Observation: tool result
|
|
```
|
|
|
|
When you have the final answer, respond in the following format:
|
|
```
|
|
Thought: I can now provide the final answer.
|
|
Final Answer: the final answer to the original input
|
|
```
|
|
|
|
Begin!
|
|
""",
|
|
),
|
|
MessagesPlaceholder(variable_name="messages"),
|
|
]
|
|
)
|
|
|
|
|
|
async def chat_node(
|
|
state: AgentState, config: RunnableConfig
|
|
) -> Command[Literal["__end__"]]:
|
|
"""
|
|
This is an enhanced agent that uses a modified ReAct pattern to allow multiple tool use.
|
|
It handles both chat responses and sequential tool execution in one node.
|
|
"""
|
|
# Get MCP configuration from state, or use the default config if not provided
|
|
mcp_config = state.get("mcp_config", DEFAULT_MCP_CONFIG)
|
|
|
|
# Set up the MCP client and tools using the configuration from state
|
|
async with MultiServerMCPClient(mcp_config) as mcp_client:
|
|
# Get the tools
|
|
mcp_tools = mcp_client.get_tools()
|
|
print(f"mcp_tools: {mcp_tools}")
|
|
|
|
# Create a model instance
|
|
model = ChatOpenAI(model="gpt-4o")
|
|
|
|
# Create the enhanced multi-tool react agent with our custom prompt
|
|
react_agent = create_react_agent(
|
|
model, mcp_tools, prompt=MULTI_TOOL_REACT_PROMPT
|
|
)
|
|
|
|
# Prepare messages for the react agent
|
|
agent_input = {"messages": state["messages"]}
|
|
|
|
# Run the react agent subgraph with our input
|
|
agent_response = await react_agent.ainvoke(agent_input)
|
|
|
|
print(f"agent_response: {agent_response}")
|
|
|
|
# Update the state with the new messages
|
|
updated_messages = state["messages"] + agent_response.get("messages", [])
|
|
|
|
# End the graph with the updated messages
|
|
return Command(
|
|
goto=END,
|
|
update={"messages": updated_messages},
|
|
)
|
|
|
|
|
|
# Define the workflow graph with only a chat node
|
|
workflow = StateGraph(AgentState)
|
|
workflow.add_node("chat_node", chat_node)
|
|
workflow.set_entry_point("chat_node")
|
|
|
|
# Compile the workflow graph
|
|
graph = workflow.compile(MemorySaver())
|