|
| 1 | +"""LangChain/LangGraph agent with automatic per-turn tool filtering. |
| 2 | +
|
| 3 | +Wraps ``create_react_agent`` so the LLM only sees relevant tools each turn, |
| 4 | +cutting token usage dramatically on large tool sets. |
| 5 | +
|
| 6 | +Usage:: |
| 7 | +
|
| 8 | + from graph_tool_call.langchain import create_agent |
| 9 | +
|
| 10 | + agent = create_agent(llm, tools=all_200_tools, top_k=5) |
| 11 | + result = agent.invoke({"messages": [("user", "cancel my order")]}) |
| 12 | + # LLM saw only ~5 tools instead of 200 |
| 13 | +
|
| 14 | +This works by passing a dynamic model factory to ``create_react_agent``: |
| 15 | +each turn, the latest user message is used to retrieve relevant tools |
| 16 | +via ``ToolGraph``, and only those are bound to the model. |
| 17 | +""" |
| 18 | + |
| 19 | +from __future__ import annotations |
| 20 | + |
| 21 | +import logging |
| 22 | +from typing import TYPE_CHECKING, Any |
| 23 | + |
| 24 | +if TYPE_CHECKING: |
| 25 | + pass |
| 26 | + |
| 27 | +logger = logging.getLogger("graph-tool-call.langchain") |
| 28 | + |
| 29 | + |
| 30 | +def _extract_query_from_langchain_messages(messages: list[Any]) -> str | None: |
| 31 | + """Extract the latest user message text from LangChain BaseMessage list.""" |
| 32 | + for msg in reversed(messages): |
| 33 | + # LangChain BaseMessage |
| 34 | + if hasattr(msg, "type") and hasattr(msg, "content"): |
| 35 | + if msg.type == "human": |
| 36 | + content = msg.content |
| 37 | + if isinstance(content, str): |
| 38 | + return content |
| 39 | + if isinstance(content, list): |
| 40 | + texts = [] |
| 41 | + for part in content: |
| 42 | + if isinstance(part, dict) and part.get("type") == "text": |
| 43 | + texts.append(part.get("text", "")) |
| 44 | + elif isinstance(part, str): |
| 45 | + texts.append(part) |
| 46 | + if texts: |
| 47 | + return " ".join(texts) |
| 48 | + |
| 49 | + # Tuple format: ("user", "message") |
| 50 | + if isinstance(msg, (list, tuple)) and len(msg) >= 2: |
| 51 | + if msg[0] in ("user", "human"): |
| 52 | + return str(msg[1]) |
| 53 | + |
| 54 | + return None |
| 55 | + |
| 56 | + |
| 57 | +def create_agent( |
| 58 | + model: Any, |
| 59 | + tools: list[Any], |
| 60 | + *, |
| 61 | + top_k: int = 5, |
| 62 | + graph: Any | None = None, |
| 63 | + **kwargs: Any, |
| 64 | +) -> Any: |
| 65 | + """Create a ReAct agent with automatic per-turn tool filtering. |
| 66 | +
|
| 67 | + Each LLM turn, the latest user message is used to retrieve the ``top_k`` |
| 68 | + most relevant tools via ``ToolGraph``. The model only sees (and pays tokens |
| 69 | + for) those tools — not the full list. |
| 70 | +
|
| 71 | + Parameters |
| 72 | + ---------- |
| 73 | + model: |
| 74 | + A LangChain ``BaseChatModel`` (e.g. ``ChatOpenAI``, ``ChatAnthropic``). |
| 75 | + tools: |
| 76 | + Full list of tools (LangChain ``BaseTool``, callables, or dicts). |
| 77 | + top_k: |
| 78 | + Number of tools to show the LLM each turn (default: 5). |
| 79 | + graph: |
| 80 | + Optional pre-built ``ToolGraph``. If *None*, one is built from *tools*. |
| 81 | + **kwargs: |
| 82 | + Passed through to ``create_react_agent`` (prompt, checkpointer, etc.). |
| 83 | +
|
| 84 | + Returns |
| 85 | + ------- |
| 86 | + CompiledStateGraph |
| 87 | + A LangGraph agent that can be invoked with |
| 88 | + ``agent.invoke({"messages": [...]})``. |
| 89 | + """ |
| 90 | + try: |
| 91 | + from langgraph.prebuilt import create_react_agent |
| 92 | + except ImportError: |
| 93 | + raise ImportError( |
| 94 | + "langgraph is required for create_agent(). " |
| 95 | + "Install with: pip install langgraph" |
| 96 | + ) |
| 97 | + |
| 98 | + from graph_tool_call import ToolGraph |
| 99 | + from graph_tool_call.toolkit import _extract_name, _ingest_tools |
| 100 | + |
| 101 | + # Build tool graph |
| 102 | + if graph is None: |
| 103 | + graph = ToolGraph() |
| 104 | + |
| 105 | + tool_map: dict[str, Any] = {} |
| 106 | + for t in tools: |
| 107 | + name = _extract_name(t) |
| 108 | + if name: |
| 109 | + tool_map[name] = t |
| 110 | + |
| 111 | + existing = set(graph.tools.keys()) |
| 112 | + if not existing.intersection(tool_map.keys()): |
| 113 | + _ingest_tools(graph, tools) |
| 114 | + |
| 115 | + # Dynamic model factory: called each turn with (state, runtime) |
| 116 | + def model_factory(state: dict[str, Any], runtime: Any) -> Any: |
| 117 | + messages = state.get("messages", []) |
| 118 | + query = _extract_query_from_langchain_messages(messages) |
| 119 | + |
| 120 | + if query: |
| 121 | + results = graph.retrieve(query, top_k=top_k) |
| 122 | + result_names = [r.name for r in results] |
| 123 | + filtered = [tool_map[n] for n in result_names if n in tool_map] |
| 124 | + |
| 125 | + if filtered: |
| 126 | + logger.debug( |
| 127 | + "Turn filter: %d → %d tools for: %s", |
| 128 | + len(tools), |
| 129 | + len(filtered), |
| 130 | + query[:50], |
| 131 | + ) |
| 132 | + return model.bind_tools(filtered) |
| 133 | + |
| 134 | + # Fallback: bind all tools |
| 135 | + return model.bind_tools(tools) |
| 136 | + |
| 137 | + return create_react_agent( |
| 138 | + model=model_factory, |
| 139 | + tools=tools, |
| 140 | + **kwargs, |
| 141 | + ) |
0 commit comments