From 8b94df72799ae3e87cd24dc66c3265991eb40f2a Mon Sep 17 00:00:00 2001 From: Zhongxuan Wang Date: Thu, 6 Aug 2026 11:40:29 -0700 Subject: [PATCH] feat: add LangGraph example workflows Signed-off-by: Zhongxuan Wang --- examples/README.md | 7 + examples/langgraph/README.md | 107 ++++++++++++++ examples/langgraph/__init__.py | 4 + examples/langgraph/calculator_mcp.py | 139 ++++++++++++++++++ examples/langgraph/config.py | 128 ++++++++++++++++ .../langgraph/configs/calculator_mcp.yaml | 36 +++++ .../configs/email_phishing_analyzer.yaml | 24 +++ examples/langgraph/email_phishing_analyzer.py | 98 ++++++++++++ examples/langgraph/mcp_math_server.py | 58 ++++++++ tests/examples/test_langgraph_examples.py | 101 +++++++++++++ 10 files changed, 702 insertions(+) create mode 100644 examples/langgraph/README.md create mode 100644 examples/langgraph/__init__.py create mode 100644 examples/langgraph/calculator_mcp.py create mode 100644 examples/langgraph/config.py create mode 100644 examples/langgraph/configs/calculator_mcp.yaml create mode 100644 examples/langgraph/configs/email_phishing_analyzer.yaml create mode 100644 examples/langgraph/email_phishing_analyzer.py create mode 100644 examples/langgraph/mcp_math_server.py create mode 100644 tests/examples/test_langgraph_examples.py diff --git a/examples/README.md b/examples/README.md index 105b091d..acf0f98d 100644 --- a/examples/README.md +++ b/examples/README.md @@ -37,6 +37,13 @@ just build-all --input "Reply with exactly: NeMo Fabric works" ``` +## LangGraph examples + +[`langgraph`](langgraph/README.md) contains native LangGraph calculator MCP and +email phishing analyzer examples. They establish the factory, model, MCP, state, +and structured-output boundaries that a reusable NeMo Fabric LangGraph adapter +must implement. + ## Harbor [`harbor`](harbor/README.md) demonstrates how to evaluate NeMo Fabric agents with diff --git a/examples/langgraph/README.md b/examples/langgraph/README.md new file mode 100644 index 00000000..5b09333c --- /dev/null +++ b/examples/langgraph/README.md @@ -0,0 +1,107 @@ + + +# LangGraph Examples + +This directory contains two native LangGraph examples: + +- `calculator_mcp.py` implements a per-user ReAct agent with a local + `current_timezone` tool and a per-user streamable-HTTP MCP calculator client. +- `email_phishing_analyzer.py` implements a purpose-built state graph that returns + a structured phishing assessment. + +They use `meta/llama-3.1-70b-instruct` through the OpenAI-compatible NIM endpoint. +NVIDIA NeMo Fabric does not yet ship a LangGraph adapter. The YAML files are +validated by the examples and intentionally use `langgraph:` entry points; they +are not accepted by `FabricConfig` yet. + +## Set Up the Python Environment + +From the repository root, install the adapter and test dependency groups that +provide LangGraph, LangChain MCP adapters, and the local MCP server: + +```bash +uv sync --no-default-groups --group adapters --group adapter-tests --group test +``` + +## Run the Calculator Example + +Set an NVIDIA API key and start the included MCP server in one terminal: + +```bash +export NVIDIA_API_KEY= +.venv/bin/python -m examples.langgraph.mcp_math_server --port 9901 +``` + +Run the calculator graph for a user in another terminal: + +```bash +.venv/bin/python -m examples.langgraph.calculator_mcp \ + --user-id alice \ + --input "What is 9 multiplied by 7?" +``` + +Each user ID creates a separate graph, `InMemorySaver`, and +`MultiServerMCPClient`. Reusing a user ID resumes only that user's conversation. +The example passes `verbose` to LangGraph's `debug` setting. +`retry_parsing_errors` has no direct LangGraph equivalent, so it is retained in +the source configuration as an adapter requirement rather than silently applied. + +## Run the Phishing Analyzer Example + +Set an NVIDIA API key, then run the structured-output graph: + +```bash +export NVIDIA_API_KEY= +.venv/bin/python -m examples.langgraph.email_phishing_analyzer \ + --input "Provide your account and routing numbers to receive a refund." +``` + +The workflow has application-owned state (`body` and `assessment`) and returns a +JSON-safe result. The NIM binding uses OpenAI-compatible function calling to +produce the structured result. It does not use MCP tools or per-user +checkpointers, which makes it the contrasting workflow required to define a +reusable adapter boundary. + +## Validate the Examples + +Run the offline checks without an API key or a live MCP server: + +```bash +.venv/bin/python -m pytest tests/examples/test_langgraph_examples.py +``` + +## Work Needed for a Full Adapter + +The examples validate the two workflow shapes, but a generic adapter needs the +following additional work before it can support them through NVIDIA NeMo Fabric: + +1. Add a LangGraph adapter package, descriptor, installation extra, wheel data, + and lifecycle host. The package needs a `langgraph_factory` entry point, not an + arbitrary compiled-graph import. +2. Define static and dynamic workflow contracts. A contract must validate + `workflow.settings`, declare each accepted normalized capability, define model, + instruction, local-tool, and MCP injection points, and contribute a digest to + the run plan. +3. Map a selected Fabric model alias to an NIM `ChatOpenAI` binding. The adapter + must validate the NIM base URL and credential environment variable before + starting a runtime. +4. Build a policy-aware complete tool inventory. It must retain MCP origins such + as `mcp_math__calculator__multiply`, enforce `tools.enabled` and + `tools.blocked` for local and MCP tools, and reject unknown selectors. +5. Create MCP clients and graph/checkpoint resources with the correct lifetime. + The calculator shows per-user state, but Fabric currently scopes a runtime to a + Fabric runtime ID rather than an end-user identity. A full adapter needs an + explicit, authenticated user/session contract before claiming per-user support. +6. Define the phishing graph's input/output projection, structured-output failure + behavior, and retry semantics. `retry_parsing_errors` needs a documented + LangGraph-equivalent policy that avoids retrying side-effecting tool calls. +7. Add conformance coverage for contract resolution, capability rejection, + multi-invocation state isolation, MCP filtering and tool policy, normalized + failures, cleanup, and opt-in NIM/MCP end-to-end tests. + +These examples show why the factory-and-contract boundary is necessary: the +workflows share model and tool injection but do not share state, output projection, +retry behavior, or resource lifetime. diff --git a/examples/langgraph/__init__.py b/examples/langgraph/__init__.py new file mode 100644 index 00000000..f5c4717e --- /dev/null +++ b/examples/langgraph/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native LangGraph examples for NVIDIA NeMo Fabric.""" diff --git a/examples/langgraph/calculator_mcp.py b/examples/langgraph/calculator_mcp.py new file mode 100644 index 00000000..0dbf7cb1 --- /dev/null +++ b/examples/langgraph/calculator_mcp.py @@ -0,0 +1,139 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""A per-user LangGraph ReAct agent with local and MCP calculator tools.""" + +from __future__ import annotations + +import argparse +import asyncio +import os +from collections.abc import Callable +from typing import Any + +from langchain.agents import create_agent +from langchain_core.tools import tool +from langchain_mcp_adapters.client import MultiServerMCPClient +from langgraph.checkpoint.memory import InMemorySaver + +from examples.langgraph.config import LangGraphExampleConfig +from examples.langgraph.config import McpServerConfig +from examples.langgraph.config import build_nim_chat_model +from examples.langgraph.config import load_config + + +@tool +def current_timezone() -> str: + """Return the server's configured IANA time zone, or ``UTC`` when unset.""" + + return os.environ.get("TZ", "UTC") + + +def mcp_connection(server: McpServerConfig) -> dict[str, str]: + """Translate the example's hyphenated transport into LangChain MCP syntax.""" + + return {"transport": "streamable_http", "url": str(server.url)} + + +def _selected_mcp_tools( + tools: list[Any], server: McpServerConfig +) -> list[Any]: + if not server.include: + return tools + by_name = {item.name: item for item in tools} + missing = set(server.include) - set(by_name) + if missing: + raise RuntimeError( + "MCP server did not expose configured tool(s): " + ", ".join(sorted(missing)) + ) + return [by_name[name] for name in server.include] + + +class PerUserReactAgent: + """Create an isolated LangGraph, MCP client, and checkpoint store per user.""" + + def __init__( + self, + config: LangGraphExampleConfig, + *, + model_factory: Callable[[Any], Any] = build_nim_chat_model, + mcp_client_factory: Callable[[dict[str, Any]], Any] = MultiServerMCPClient, + graph_factory: Callable[..., Any] = create_agent, + ) -> None: + if config.workflow.entrypoint != "langgraph:per_user_react_agent": + raise ValueError("calculator example requires langgraph:per_user_react_agent") + if config.mcp is None or "mcp_math" not in config.mcp.servers: + raise ValueError("calculator example requires mcp.servers.mcp_math") + if "current_timezone" not in config.tools: + raise ValueError("calculator example requires tools.current_timezone") + + self._config = config + self._model_factory = model_factory + self._mcp_client_factory = mcp_client_factory + self._graph_factory = graph_factory + self._sessions: dict[str, Any] = {} + + async def graph_for(self, user_id: str) -> Any: + """Return the user-owned graph, creating it and its MCP client on first use.""" + + if not user_id.strip(): + raise ValueError("user_id must be a non-empty string") + graph = self._sessions.get(user_id) + if graph is not None: + return graph + + server = self._config.mcp.servers["mcp_math"] # validated in __init__ + client = self._mcp_client_factory( + {"mcp_math": mcp_connection(server)}, tool_name_prefix=False + ) + mcp_tools = _selected_mcp_tools(list(await client.get_tools()), server) + model = self._model_factory(self._config.selected_model()) + graph = self._graph_factory( + model, + [current_timezone, *mcp_tools], + checkpointer=InMemorySaver(), + debug=bool(self._config.workflow.settings.get("verbose", False)), + name="per_user_calculator", + ) + self._sessions[user_id] = graph + return graph + + async def ainvoke(self, user_id: str, message: str) -> dict[str, Any]: + """Run a message in the graph and persisted conversation for ``user_id``.""" + + graph = await self.graph_for(user_id) + return await graph.ainvoke( + {"messages": [{"role": "user", "content": message}]}, + {"configurable": {"thread_id": user_id}}, + ) + + +def build_per_user_react_agent(config: LangGraphExampleConfig) -> PerUserReactAgent: + """Build the calculator example's declared LangGraph factory entry point.""" + + return PerUserReactAgent(config) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", default="examples/langgraph/configs/calculator_mcp.yaml") + parser.add_argument("--user-id", required=True) + parser.add_argument("--input", required=True) + return parser.parse_args() + + +async def _run() -> None: + args = _parse_args() + graph = build_per_user_react_agent(load_config(args.config)) + result = await graph.ainvoke(args.user_id, args.input) + print(result["messages"][-1].content) + + +def main() -> None: + """Run the calculator example from the command line.""" + + asyncio.run(_run()) + + +if __name__ == "__main__": + main() diff --git a/examples/langgraph/config.py b/examples/langgraph/config.py new file mode 100644 index 00000000..9b3ef083 --- /dev/null +++ b/examples/langgraph/config.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Configuration shared by the native LangGraph examples. + +This is application configuration for the examples, not a NeMo Fabric adapter +descriptor or a supported Fabric configuration format. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any +from typing import Literal + +import yaml +from langchain_openai import ChatOpenAI +from pydantic import BaseModel +from pydantic import Field +from pydantic import HttpUrl +from pydantic import model_validator + +NIM_OPENAI_BASE_URL = "https://integrate.api.nvidia.com/v1" + + +class NimModelConfig(BaseModel): + """One NVIDIA NIM model binding used by a LangGraph example.""" + + provider: Literal["nim"] + model: str = Field(min_length=1) + api_key_env: str = "NVIDIA_API_KEY" + base_url: HttpUrl = NIM_OPENAI_BASE_URL + temperature: float = 0.0 + max_tokens: int = Field(default=1024, gt=0) + + +class McpServerConfig(BaseModel): + """The streamable HTTP MCP server exposed to the calculator workflow.""" + + transport: Literal["streamable-http"] + url: HttpUrl + include: list[str] = Field(default_factory=list) + + +class McpConfig(BaseModel): + """MCP servers available to a LangGraph example.""" + + servers: dict[str, McpServerConfig] = Field(default_factory=dict) + + +class LocalToolConfig(BaseModel): + """A local tool supplied by the workflow rather than an MCP server.""" + + kind: Literal["local"] + description: str = Field(min_length=1) + + +class WorkflowConfig(BaseModel): + """Select a LangGraph graph and its workflow-owned settings.""" + + entrypoint: Literal[ + "langgraph:per_user_react_agent", "langgraph:email_phishing_analyzer" + ] + settings: dict[str, Any] = Field(default_factory=dict) + + +class LangGraphExampleConfig(BaseModel): + """Validated source configuration for the two native LangGraph examples.""" + + models: dict[str, NimModelConfig] = Field(min_length=1) + mcp: McpConfig | None = None + tools: dict[str, LocalToolConfig] = Field(default_factory=dict) + workflow: WorkflowConfig + + @model_validator(mode="after") + def _validate_workflow_references(self) -> "LangGraphExampleConfig": + llm_name = self.workflow.settings.get("llm_name") + if not isinstance(llm_name, str) or llm_name not in self.models: + raise ValueError("workflow.settings.llm_name must name a configured model") + + tool_names = self.workflow.settings.get("tool_names", []) + if not isinstance(tool_names, list) or not all( + isinstance(name, str) for name in tool_names + ): + raise ValueError("workflow.settings.tool_names must be a list of strings") + + available = set(self.tools) + if self.mcp is not None: + available.update(self.mcp.servers) + unknown = set(tool_names) - available + if unknown: + raise ValueError( + "workflow.settings.tool_names names unknown tool source(s): " + + ", ".join(sorted(unknown)) + ) + return self + + def selected_model(self) -> NimModelConfig: + """Return the NIM model selected by the workflow.""" + + return self.models[self.workflow.settings["llm_name"]] + + +def load_config(path: str | Path) -> LangGraphExampleConfig: + """Load and validate one LangGraph example YAML configuration file.""" + + raw = yaml.safe_load(Path(path).read_text(encoding="utf-8")) + if not isinstance(raw, dict): + raise ValueError("LangGraph example configuration must be a mapping") + return LangGraphExampleConfig.model_validate(raw) + + +def build_nim_chat_model(model: NimModelConfig) -> ChatOpenAI: + """Build the OpenAI-compatible NIM chat-model binding for a graph.""" + + api_key = os.environ.get(model.api_key_env) + if not api_key: + raise RuntimeError( + f"Set {model.api_key_env} before running a workflow that uses {model.model}." + ) + return ChatOpenAI( + model=model.model, + api_key=api_key, + base_url=str(model.base_url), + temperature=model.temperature, + max_tokens=model.max_tokens, + ) diff --git a/examples/langgraph/configs/calculator_mcp.yaml b/examples/langgraph/configs/calculator_mcp.yaml new file mode 100644 index 00000000..f95e37f9 --- /dev/null +++ b/examples/langgraph/configs/calculator_mcp.yaml @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +models: + nim_llm: + provider: nim + model: meta/llama-3.1-70b-instruct + temperature: 0.0 + max_tokens: 1024 + +mcp: + servers: + mcp_math: + transport: streamable-http + url: http://localhost:9901/mcp + include: + - calculator__add + - calculator__subtract + - calculator__multiply + - calculator__divide + +tools: + current_timezone: + kind: local + description: Return the IANA time zone configured for this server. + +workflow: + entrypoint: langgraph:per_user_react_agent + settings: + llm_name: nim_llm + tool_names: + - current_timezone + - mcp_math + verbose: true + retry_parsing_errors: true + max_retries: 3 diff --git a/examples/langgraph/configs/email_phishing_analyzer.yaml b/examples/langgraph/configs/email_phishing_analyzer.yaml new file mode 100644 index 00000000..34f33ef8 --- /dev/null +++ b/examples/langgraph/configs/email_phishing_analyzer.yaml @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +models: + nim_llm: + provider: nim + model: meta/llama-3.1-70b-instruct + temperature: 0.0 + max_tokens: 512 + +workflow: + entrypoint: langgraph:email_phishing_analyzer + settings: + llm_name: nim_llm + verbose: true + retry_parsing_errors: true + max_retries: 3 + prompt: | + Examine the email below for signs of phishing. Look for requests for + personal information, suspicious tone, urgency, unusual payment requests, + and emotional manipulation. + + Email content: + {body} diff --git a/examples/langgraph/email_phishing_analyzer.py b/examples/langgraph/email_phishing_analyzer.py new file mode 100644 index 00000000..b6ae06cd --- /dev/null +++ b/examples/langgraph/email_phishing_analyzer.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""A LangGraph workflow that returns a structured phishing assessment.""" + +from __future__ import annotations + +import argparse +import asyncio +import json +from typing import Any +from typing import TypedDict + +from langgraph.graph import END +from langgraph.graph import START +from langgraph.graph import StateGraph +from pydantic import BaseModel +from pydantic import Field + +from examples.langgraph.config import LangGraphExampleConfig +from examples.langgraph.config import build_nim_chat_model +from examples.langgraph.config import load_config + +DEFAULT_PROMPT = """Examine the email below for signs of phishing. + +Look for suspicious requests, urgency, generic greetings, grammar mistakes, +unusual payment requests, and emotional manipulation. Return a structured, +evidence-based assessment. + +Email content: +{body} +""" + + +class PhishingAssessment(BaseModel): + """The JSON-safe result projected from the email-analysis graph.""" + + is_likely_phishing: bool + explanation: str = Field(min_length=1) + + +class EmailPhishingState(TypedDict): + """Application-owned graph state for one email analysis.""" + + body: str + assessment: dict[str, Any] + + +def build_email_phishing_analyzer( + config: LangGraphExampleConfig, *, model: Any | None = None +) -> Any: + """Build the phishing graph selected by the example workflow configuration.""" + + if config.workflow.entrypoint != "langgraph:email_phishing_analyzer": + raise ValueError("phishing example requires langgraph:email_phishing_analyzer") + prompt = str(config.workflow.settings.get("prompt", DEFAULT_PROMPT)) + chat_model = model or build_nim_chat_model(config.selected_model()) + structured_model = chat_model.with_structured_output( + PhishingAssessment, method="function_calling" + ) + + async def analyze_email(state: EmailPhishingState) -> dict[str, Any]: + assessment = await structured_model.ainvoke(prompt.format(body=state["body"])) + if not isinstance(assessment, PhishingAssessment): + assessment = PhishingAssessment.model_validate(assessment) + return {"assessment": assessment.model_dump()} + + graph = StateGraph(EmailPhishingState) + graph.add_node("analyze_email", analyze_email) + graph.add_edge(START, "analyze_email") + graph.add_edge("analyze_email", END) + return graph.compile(name="email_phishing_analyzer") + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--config", default="examples/langgraph/configs/email_phishing_analyzer.yaml" + ) + parser.add_argument("--input", required=True, help="Email body to analyze.") + return parser.parse_args() + + +async def _run() -> None: + args = _parse_args() + graph = build_email_phishing_analyzer(load_config(args.config)) + result = await graph.ainvoke({"body": args.input}) + print(json.dumps(result["assessment"], indent=2)) + + +def main() -> None: + """Run the phishing analyzer example from the command line.""" + + asyncio.run(_run()) + + +if __name__ == "__main__": + main() diff --git a/examples/langgraph/mcp_math_server.py b/examples/langgraph/mcp_math_server.py new file mode 100644 index 00000000..d2d77030 --- /dev/null +++ b/examples/langgraph/mcp_math_server.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""A local streamable-HTTP MCP server for the calculator example.""" + +from __future__ import annotations + +import argparse + +from mcp.server.fastmcp import FastMCP + + +def create_server(*, host: str = "127.0.0.1", port: int = 9901) -> FastMCP: + """Create an MCP server with the four calculator tools used by the example.""" + + server = FastMCP("NeMo Fabric calculator example", host=host, port=port) + + @server.tool() + def calculator__add(left: float, right: float) -> float: + """Add two numbers.""" + + return left + right + + @server.tool() + def calculator__subtract(left: float, right: float) -> float: + """Subtract ``right`` from ``left``.""" + + return left - right + + @server.tool() + def calculator__multiply(left: float, right: float) -> float: + """Multiply two numbers.""" + + return left * right + + @server.tool() + def calculator__divide(left: float, right: float) -> float: + """Divide ``left`` by ``right``.""" + + if right == 0: + raise ValueError("right must not be zero") + return left / right + + return server + + +def main() -> None: + """Run the server with streamable HTTP transport.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", default=9901, type=int) + args = parser.parse_args() + create_server(host=args.host, port=args.port).run(transport="streamable-http") + + +if __name__ == "__main__": + main() diff --git a/tests/examples/test_langgraph_examples.py b/tests/examples/test_langgraph_examples.py new file mode 100644 index 00000000..8ce22546 --- /dev/null +++ b/tests/examples/test_langgraph_examples.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Offline coverage for the native LangGraph examples.""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import AsyncMock +from unittest.mock import MagicMock + +from examples.langgraph.calculator_mcp import PerUserReactAgent +from examples.langgraph.calculator_mcp import current_timezone +from examples.langgraph.config import load_config +from examples.langgraph.email_phishing_analyzer import PhishingAssessment +from examples.langgraph.email_phishing_analyzer import build_email_phishing_analyzer + +ROOT_DIR = Path(__file__).resolve().parents[2] +CALCULATOR_CONFIG = ROOT_DIR / "examples/langgraph/configs/calculator_mcp.yaml" +PHISHING_CONFIG = ROOT_DIR / "examples/langgraph/configs/email_phishing_analyzer.yaml" + + +def test_calculator_config_preserves_requested_mcp_and_workflow_shape(): + config = load_config(CALCULATOR_CONFIG) + + assert config.selected_model().model == "meta/llama-3.1-70b-instruct" + assert config.mcp is not None + assert config.mcp.servers["mcp_math"].transport == "streamable-http" + assert config.workflow.entrypoint == "langgraph:per_user_react_agent" + assert config.workflow.settings["tool_names"] == ["current_timezone", "mcp_math"] + + +async def test_calculator_creates_isolated_graph_and_mcp_client_per_user(): + config = load_config(CALCULATOR_CONFIG) + mock_tools = [] + for name in config.mcp.servers["mcp_math"].include: # validated by the example + mock_tool = MagicMock() + mock_tool.name = name + mock_tools.append(mock_tool) + mock_client_factory = MagicMock() + mock_client_factory.side_effect = [ + MagicMock(get_tools=AsyncMock(return_value=mock_tools)) + for _ in range(2) + ] + mock_graph_factory = MagicMock(side_effect=[MagicMock(), MagicMock()]) + mock_model_factory = MagicMock(side_effect=[MagicMock(), MagicMock()]) + + agent = PerUserReactAgent( + config, + model_factory=mock_model_factory, + mcp_client_factory=mock_client_factory, + graph_factory=mock_graph_factory, + ) + alice_first = await agent.graph_for("alice") + alice_second = await agent.graph_for("alice") + hatter = await agent.graph_for("hatter") + + assert alice_first is alice_second + assert alice_first is not hatter + assert mock_model_factory.call_count == 2 + assert mock_client_factory.call_count == 2 + connection = mock_client_factory.call_args.args[0]["mcp_math"] + assert connection == { + "transport": "streamable_http", + "url": "http://localhost:9901/mcp", + } + for call in mock_graph_factory.call_args_list: + assert call.kwargs["checkpointer"] is not None + assert call.kwargs["name"] == "per_user_calculator" + + +def test_current_timezone_uses_the_explicit_server_configuration(restore_environ): + restore_environ["TZ"] = "America/Los_Angeles" + + assert current_timezone.invoke({}) == "America/Los_Angeles" + + +async def test_phishing_graph_projects_a_json_safe_structured_assessment(): + config = load_config(PHISHING_CONFIG) + mock_structured_model = MagicMock() + mock_structured_model.ainvoke = AsyncMock( + return_value=PhishingAssessment( + is_likely_phishing=True, + explanation="It asks for banking information to complete a refund.", + ) + ) + mock_model = MagicMock() + mock_model.with_structured_output.return_value = mock_structured_model + + graph = build_email_phishing_analyzer(config, model=mock_model) + result = await graph.ainvoke( + {"body": "Provide your routing number so we can issue a refund."} + ) + + assert result["assessment"] == { + "is_likely_phishing": True, + "explanation": "It asks for banking information to complete a refund.", + } + mock_model.with_structured_output.assert_called_once_with( + PhishingAssessment, method="function_calling" + )