| import json |
| from pydantic import BaseModel, Field |
| from langchain_openai import ChatOpenAI |
| from langchain_core.messages import SystemMessage |
| from .state import AgentState, tools_config |
|
|
| class ToolSelection(BaseModel): |
| tool_name: str | None = Field( |
| description="Name of the selected tool, or null if no tool matches." |
| ) |
|
|
| llm = ChatOpenAI(model="gpt-4o") |
|
|
| selector_llm = llm.with_structured_output( |
| ToolSelection |
| ) |
|
|
|
|
| def select_tool(state: AgentState): |
| system_prompt = f""" |
| You are a tool selection agent. |
| Available tools: |
| {json.dumps(tools_config)} |
| |
| Current extracted parameters: |
| {json.dumps(state["global_params"])} |
| |
| Rules: |
| - Select the tool that matches the user's intent. |
| - Use conversation history for context. |
| - Do not select a tool only because its parameters |
| are available. |
| - Do not invent tool names. |
| - If no tool is appropriate, return null. |
| - Do not execute any tool. |
| """ |
|
|
| result = selector_llm.invoke([ |
| SystemMessage(content=system_prompt), |
| *state["messages"] |
| ]) |
|
|
| allowed_tools = { |
| tool["name"] for tool in tools_config |
| } |
|
|
| selected = ( |
| result.tool_name |
| if result.tool_name in allowed_tools |
| else None |
| ) |
|
|
| return { |
| "selected_tool": selected |
| } |
|
|
|
|
|
|
