From 0d0384d94c24613e3b4f953355c4f30ea50a9766 Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Tue, 22 Sep 2026 11:56:59 -0800 Subject: [PATCH 1/5] feat(AIC-3408): add Vercel AI SDK messages and agents adapters Route LaunchDarkly AI Configs through the official Python ai runtime and AI Gateway, including native graph execution and a direct-model opt-out. Co-authored-by: Cursor --- .env.example | 1 + .release-please-manifest.json | 4 +- .sdk_metadata.json | 25 + AGENTS.md | 9 +- README.md | 8 + examples/native_graph_vercel.py | 27 + examples/vercel_agents_example.py | 31 ++ examples/vercel_direct.py | 31 ++ examples/vercel_messages_example.py | 32 ++ main.py | 15 +- packages/vercel-agents/README.md | 42 ++ packages/vercel-agents/agents.md | 18 + packages/vercel-agents/pyproject.toml | 21 + .../launchdarkly_ai_vercel_agents/__init__.py | 18 + .../launchdarkly_ai_vercel_agents/graph.py | 20 + .../launchdarkly_ai_vercel_agents/handler.py | 453 +++++++++++++++ .../launchdarkly_ai_vercel_agents/model_id.py | 145 +++++ .../native_graph.py | 247 +++++++++ .../launchdarkly_ai_vercel_agents/py.typed | 1 + .../launchdarkly_ai_vercel_agents/spans.py | 78 +++ packages/vercel-agents/tests/test_handler.py | 368 +++++++++++++ packages/vercel-agents/tests/test_model_id.py | 47 ++ .../vercel-agents/tests/test_native_graph.py | 448 +++++++++++++++ packages/vercel-messages/README.md | 63 +++ packages/vercel-messages/agents.md | 19 + packages/vercel-messages/pyproject.toml | 21 + .../__init__.py | 11 + .../handler.py | 518 ++++++++++++++++++ .../model_id.py | 145 +++++ .../launchdarkly_ai_vercel_messages/py.typed | 1 + .../launchdarkly_ai_vercel_messages/spans.py | 78 +++ .../vercel-messages/tests/test_handler.py | 503 +++++++++++++++++ .../vercel-messages/tests/test_model_id.py | 55 ++ pyproject.toml | 3 + release-please-config.json | 22 + tests/test_cross_handler_parity.py | 2 + uv.lock | 114 ++++ 37 files changed, 3639 insertions(+), 5 deletions(-) create mode 100644 examples/native_graph_vercel.py create mode 100644 examples/vercel_agents_example.py create mode 100644 examples/vercel_direct.py create mode 100644 examples/vercel_messages_example.py create mode 100644 packages/vercel-agents/README.md create mode 100644 packages/vercel-agents/agents.md create mode 100644 packages/vercel-agents/pyproject.toml create mode 100644 packages/vercel-agents/src/launchdarkly_ai_vercel_agents/__init__.py create mode 100644 packages/vercel-agents/src/launchdarkly_ai_vercel_agents/graph.py create mode 100644 packages/vercel-agents/src/launchdarkly_ai_vercel_agents/handler.py create mode 100644 packages/vercel-agents/src/launchdarkly_ai_vercel_agents/model_id.py create mode 100644 packages/vercel-agents/src/launchdarkly_ai_vercel_agents/native_graph.py create mode 100644 packages/vercel-agents/src/launchdarkly_ai_vercel_agents/py.typed create mode 100644 packages/vercel-agents/src/launchdarkly_ai_vercel_agents/spans.py create mode 100644 packages/vercel-agents/tests/test_handler.py create mode 100644 packages/vercel-agents/tests/test_model_id.py create mode 100644 packages/vercel-agents/tests/test_native_graph.py create mode 100644 packages/vercel-messages/README.md create mode 100644 packages/vercel-messages/agents.md create mode 100644 packages/vercel-messages/pyproject.toml create mode 100644 packages/vercel-messages/src/launchdarkly_ai_vercel_messages/__init__.py create mode 100644 packages/vercel-messages/src/launchdarkly_ai_vercel_messages/handler.py create mode 100644 packages/vercel-messages/src/launchdarkly_ai_vercel_messages/model_id.py create mode 100644 packages/vercel-messages/src/launchdarkly_ai_vercel_messages/py.typed create mode 100644 packages/vercel-messages/src/launchdarkly_ai_vercel_messages/spans.py create mode 100644 packages/vercel-messages/tests/test_handler.py create mode 100644 packages/vercel-messages/tests/test_model_id.py diff --git a/.env.example b/.env.example index 5bb04800..bb3a50ba 100644 --- a/.env.example +++ b/.env.example @@ -4,6 +4,7 @@ LD_SDK_KEY=your-sdk-key-here # AI providers (set the ones you use) OPENAI_API_KEY=your-openai-key-here ANTHROPIC_API_KEY=your-anthropic-key-here +AI_GATEWAY_API_KEY=your-vercel-ai-gateway-key-here # LaunchDarkly endpoints (leave blank to use production defaults) LD_BASE_URI= diff --git a/.release-please-manifest.json b/.release-please-manifest.json index 26e16e54..5608d945 100644 --- a/.release-please-manifest.json +++ b/.release-please-manifest.json @@ -6,5 +6,7 @@ "packages/openai-agents": "0.2.2", "packages/openai-messages": "0.2.2", "packages/langchain-agents": "0.2.2", - "packages/langchain-messages": "0.2.2" + "packages/langchain-messages": "0.2.2", + "packages/vercel-agents": "0.1.0", + "packages/vercel-messages": "0.1.0" } diff --git a/.sdk_metadata.json b/.sdk_metadata.json index 5d7708f3..6d63f548 100644 --- a/.sdk_metadata.json +++ b/.sdk_metadata.json @@ -102,6 +102,31 @@ "aiAgentGraph": { "introduced": "0.1" }, "aiTrackMetrics": { "introduced": "0.1" } } + }, + "python-ai-vercel-messages": { + "name": "Vercel Messages AI SDK for Python", + "type": "ai-provider", + "path": "packages/vercel-messages", + "languages": ["Python"], + "aiSdkNames": ["launchdarkly-ai-vercel-messages"], + "releases": { "tag-prefix": "launchdarkly-ai-vercel-messages-" }, + "features": { + "aiCompletionConfig": { "introduced": "0.1" }, + "aiTrackMetrics": { "introduced": "0.1" } + } + }, + "python-ai-vercel-agents": { + "name": "Vercel Agents AI SDK for Python", + "type": "ai-provider", + "path": "packages/vercel-agents", + "languages": ["Python"], + "aiSdkNames": ["launchdarkly-ai-vercel-agents"], + "releases": { "tag-prefix": "launchdarkly-ai-vercel-agents-" }, + "features": { + "aiAgentConfig": { "introduced": "0.1" }, + "aiAgentGraph": { "introduced": "0.1" }, + "aiTrackMetrics": { "introduced": "0.1" } + } } } } diff --git a/AGENTS.md b/AGENTS.md index 305613a9..5f47254c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -87,6 +87,7 @@ graph TD claude["launchdarkly-ai-claude-agents"] openai["launchdarkly-ai-openai-agents"] langchain["launchdarkly-ai-langchain-agents"] + vercel["launchdarkly-ai-vercel-agents"] newHandler["launchdarkly-ai-new-provider\n(future)"] end subgraph tier0 ["Tier 0 — Core"] @@ -102,6 +103,7 @@ graph TD claude --> client openai --> client langchain --> client + vercel --> client newHandler --> client ai --> client ``` @@ -110,7 +112,7 @@ graph TD - **Tier 0 — Core** (`launchdarkly-ai-server`): The foundation. Owns all LaunchDarkly integration, telemetry orchestration, shared data types, and the primary entry points (`config()`, `graph()`, `resolve_graph()`). Has no dependency on any other `launchdarkly-ai-*` package. - **Tier 0 — Convenience barrel** (`launchdarkly-ai-python`): A pure re-export package that makes all of `launchdarkly-ai-server` available under a shorter install name. No new logic — intended as the default install for most Python applications. -- **Tier 1 — Handler packages** (`launchdarkly-ai-claude-agents`, `launchdarkly-ai-claude-messages`, `launchdarkly-ai-openai-agents`, `launchdarkly-ai-openai-messages`, `launchdarkly-ai-langchain-agents`, `launchdarkly-ai-langchain-messages`, …): Each wraps a specific AI provider SDK. Depends on `launchdarkly-ai-server` for shared types and utilities. Must not depend on other Tier 1 packages. +- **Tier 1 — Handler packages** (`launchdarkly-ai-claude-agents`, `launchdarkly-ai-claude-messages`, `launchdarkly-ai-openai-agents`, `launchdarkly-ai-openai-messages`, `launchdarkly-ai-langchain-agents`, `launchdarkly-ai-langchain-messages`, `launchdarkly-ai-vercel-agents`, `launchdarkly-ai-vercel-messages`, …): Each wraps a specific AI provider SDK. Depends on `launchdarkly-ai-server` for shared types and utilities. Must not depend on other Tier 1 packages. - **Tier 2 — Consumer applications** (e.g. `main.py`, downstream projects): Imports from one or more handler packages and either `launchdarkly-ai-python` or `launchdarkly-ai-server`. Owns tool implementations and orchestration logic. No `launchdarkly-ai-*` package should ever depend on Tier 2 code. ### Rules @@ -657,7 +659,7 @@ def my_provider( For example, `claude_agents(config_key, user_input, context)` is equivalent to `config(key=config_key, handler=create_claude_agents_handler()).invoke(user_input, context)`. -The naming convention matches the package suffix: `claude_agents`, `claude_messages`, `openai_agents`, `openai_messages`, `langchain_agents`, `langchain_messages`. +The naming convention matches the package suffix: `claude_agents`, `claude_messages`, `openai_agents`, `openai_messages`, `langchain_agents`, `langchain_messages`, `vercel_agents`, `vercel_messages`. ### Graph Export (optional) @@ -668,7 +670,7 @@ def claude_graph(key: str, **options) -> GraphInstance: return graph(key, handlers=[create_claude_agents_handler()], **options) ``` -Naming convention: `claude_graph`, `openai_graph`, `langchain_graph`. +Naming convention: `claude_graph`, `openai_graph`, `langchain_graph`, `vercel_graph`. ### Native Graph Adapter (optional) @@ -678,6 +680,7 @@ Current adapters: - `to_claude_agents(def_coro, opts)` — exported from `launchdarkly_ai_claude_agents` - `to_openai_agents(def_coro, opts)` — exported from `launchdarkly_ai_openai_agents` - `to_lang_graph(def_coro, opts)` — exported from `launchdarkly_ai_langchain_agents` +- `to_vercel_agents(def_coro, opts)` — exported from `launchdarkly_ai_vercel_agents` --- diff --git a/README.md b/README.md index 0e693665..262e3682 100644 --- a/README.md +++ b/README.md @@ -85,6 +85,8 @@ Tier 0 — Core Client (launchdarkly-ai-server) | [`launchdarkly-ai-claude-agents`](packages/claude-agents/README.md) | Anthropic | `agent` | Claude Agent SDK — agentic loop with MCP tool support | | [`launchdarkly-ai-langchain-messages`](packages/langchain-messages/README.md) | `*` (any) | `messages` | Any `BaseChatModel` via LangChain `bind_tools` loop | | [`launchdarkly-ai-langchain-agents`](packages/langchain-agents/README.md) | `*` (any) | `agent` | LangGraph `StateGraph` — managed ReAct loop | +| [`launchdarkly-ai-vercel-messages`](packages/vercel-messages/README.md) | `*` (any) | `messages` | Vercel AI SDK with AI Gateway routing | +| [`launchdarkly-ai-vercel-agents`](packages/vercel-agents/README.md) | `*` (any) | `agent` | Native Vercel `ai.Agent` loop and graph adapter | ## Quick Start @@ -149,6 +151,8 @@ asyncio.run(main()) | `claude_agents` | `launchdarkly-ai-claude-agents` | `claude-agent-sdk` | Claude Agent SDK (MCP) | | `langchain_messages` | `launchdarkly-ai-langchain-messages` | `langchain-core` | LangChain `bind_tools` loop | | `langchain_agents` | `launchdarkly-ai-langchain-agents` | `langgraph` | LangGraph `StateGraph` | +| `vercel_messages` | `launchdarkly-ai-vercel-messages` | `ai` | Vercel AI SDK stream runtime | +| `vercel_agents` | `launchdarkly-ai-vercel-agents` | `ai` | Native Vercel `ai.Agent` | --- @@ -597,6 +601,10 @@ uv run python main.py [example] [flag-key] [user-input] | `graph-history` | `uv run python main.py graph-history` | `graph().invoke()` with multimodal `history` forwarded to the root node | | `openai-only` | `uv run python main.py openai-only` | `config()` with a custom `Registry` restricted to OpenAI handlers | | `streaming` | `uv run python main.py streaming` | `config().stream()` — token-by-token output | +| `vercel-agents` | `uv run python main.py vercel-agents` | Native Vercel `ai.Agent` handler via AI Gateway | +| `vercel-messages` | `uv run python main.py vercel-messages` | Vercel AI SDK messages handler via AI Gateway | +| `vercel-direct` | `uv run python main.py vercel-direct` | Vercel messages handler with a constructed OpenAI model (no Gateway) | +| `native-graph-vercel` | `uv run python main.py native-graph-vercel` | Framework-native Vercel `ai.Agent` graph runner | **Examples:** diff --git a/examples/native_graph_vercel.py b/examples/native_graph_vercel.py new file mode 100644 index 00000000..f77670a0 --- /dev/null +++ b/examples/native_graph_vercel.py @@ -0,0 +1,27 @@ +"""Framework-native graph runner using Vercel ``ai.Agent`` instances.""" + +from __future__ import annotations + +import json + +from examples.tools import get_preferences, web_search +from examples.utils import new_context, write_output +from launchdarkly_ai_server import resolve_graph +from launchdarkly_ai_vercel_agents import to_vercel_agents + + +async def run(key: str, user_input: str) -> None: + context = new_context() + response = await to_vercel_agents( + resolve_graph(key, context=context), + { + "context": context, + "tool_handlers": { + "user-preferences-lookup": get_preferences, + "web-search-tool": web_search, + }, + }, + ).invoke(user_input, {"user_id": "user-123"}) + + print(json.dumps(response, indent=2, default=str)) + write_output(response) diff --git a/examples/vercel_agents_example.py b/examples/vercel_agents_example.py new file mode 100644 index 00000000..0d2a06b9 --- /dev/null +++ b/examples/vercel_agents_example.py @@ -0,0 +1,31 @@ +"""Run the wildcard native Vercel AI SDK agents adapter.""" + +from __future__ import annotations + +import json + +from examples.tools import ( + fetch_launchdarkly_documentation, + get_preferences, + search_ld_documentation, + web_search, +) +from examples.utils import new_context, write_output +from launchdarkly_ai_vercel_agents import vercel_agents + + +async def run(key: str, user_input: str) -> None: + response = await vercel_agents( + key, + user_input, + new_context(), + tool_handlers={ + "get-user-preferences": get_preferences, + "search-ld-documentation": search_ld_documentation, + "fetch-ld-documentation": fetch_launchdarkly_documentation, + "fetch-launchdarkly-documentation": fetch_launchdarkly_documentation, + "web-search": web_search, + }, + ) + print(json.dumps(response, indent=2, default=str)) + write_output(response) diff --git a/examples/vercel_direct.py b/examples/vercel_direct.py new file mode 100644 index 00000000..f27e4bdd --- /dev/null +++ b/examples/vercel_direct.py @@ -0,0 +1,31 @@ +"""Inject a constructed OpenAI model so the Vercel adapter skips AI Gateway. + +Usage (via main.py): + python main.py vercel-direct launch-darkly-documentation-summarizer-messages-openai "What is the LaunchDarkly AI SDK?" +""" + +from __future__ import annotations + +import json + +import ai + +from examples.utils import new_context, write_output +from launchdarkly_ai_server import AiConfigRep +from launchdarkly_ai_vercel_messages import vercel_messages + + +def openai_model(config: AiConfigRep) -> ai.Model: + return ai.Model(id=str(config["model"]["name"]), provider=ai.get_provider("openai")) + + +async def run(key: str, user_input: str) -> None: + response = await vercel_messages( + key, + user_input, + new_context(), + model=openai_model, + variables={"user_input": user_input}, + ) + print(json.dumps(response, indent=2, default=str)) + write_output(response) diff --git a/examples/vercel_messages_example.py b/examples/vercel_messages_example.py new file mode 100644 index 00000000..ce627f6a --- /dev/null +++ b/examples/vercel_messages_example.py @@ -0,0 +1,32 @@ +"""Run the wildcard Vercel AI SDK messages adapter.""" + +from __future__ import annotations + +import json + +from examples.tools import ( + fetch_launchdarkly_documentation, + get_preferences, + search_ld_documentation, + web_search, +) +from examples.utils import new_context, write_output +from launchdarkly_ai_vercel_messages import vercel_messages + + +async def run(key: str, user_input: str) -> None: + response = await vercel_messages( + key, + user_input, + new_context(), + tool_handlers={ + "get-user-preferences": get_preferences, + "search-ld-documentation": search_ld_documentation, + "fetch-ld-documentation": fetch_launchdarkly_documentation, + "fetch-launchdarkly-documentation": fetch_launchdarkly_documentation, + "web-search": web_search, + }, + variables={"user_input": user_input}, + ) + print(json.dumps(response, indent=2, default=str)) + write_output(response) diff --git a/main.py b/main.py index fc5c1e5b..1678c67e 100644 --- a/main.py +++ b/main.py @@ -16,6 +16,8 @@ python main.py openai-agents launch-darkly-documentation-summarizer-open-ai-only "What is the LaunchDarkly AI SDK?" python main.py langchain-agents launch-darkly-documentation-summarizer "What is the LaunchDarkly AI SDK?" python main.py langchain-thinking launch-darkly-documentation-summarizer-messages-claude "Reason it out yourself without any tools: what is 17 times 23?" + python main.py vercel-messages launch-darkly-documentation-summarizer-messages "What is the LaunchDarkly AI SDK?" + python main.py vercel-direct launch-darkly-documentation-summarizer-messages-openai "What is the LaunchDarkly AI SDK?" python main.py native-graph travel-agent-flow "Book me a flight to Paris" python main.py native-graph-langchain travel-agent-flow "Book me a flight to Paris" """ @@ -58,8 +60,12 @@ "langchain-agents": "examples.langchain_agents_example", "langchain-messages": "examples.langchain_messages_example", "langchain-thinking": "examples.langchain_thinking", + "vercel-agents": "examples.vercel_agents_example", + "vercel-direct": "examples.vercel_direct", + "vercel-messages": "examples.vercel_messages_example", "native-graph": "examples.native_graph", "native-graph-langchain": "examples.native_graph_langchain", + "native-graph-vercel": "examples.native_graph_vercel", } DEFAULT_EXAMPLE = "agent" @@ -97,10 +103,17 @@ async def main() -> None: await shutdown() +def _error_message(exc: BaseException) -> str: + """Expose a sole nested provider error instead of an opaque TaskGroup wrapper.""" + if isinstance(exc, BaseExceptionGroup) and len(exc.exceptions) == 1: + return _error_message(exc.exceptions[0]) + return str(exc) + + if __name__ == "__main__": try: asyncio.run(main()) except Exception as exc: sys.stdout.flush() - print(f"Error: {exc}", file=sys.stderr) + print(f"Error: {_error_message(exc)}", file=sys.stderr) sys.exit(1) diff --git a/packages/vercel-agents/README.md b/packages/vercel-agents/README.md new file mode 100644 index 00000000..a689dec1 --- /dev/null +++ b/packages/vercel-agents/README.md @@ -0,0 +1,42 @@ +# `launchdarkly-ai-vercel-agents` + +Wildcard LaunchDarkly AI agent handler built on the official Vercel `ai.Agent` +runtime for Python. + +## Install + +```bash +pip install launchdarkly-ai-server launchdarkly-server-sdk launchdarkly-ai-vercel-agents +``` + +Default routing is AI Gateway via a mapped `creator/model` id (for example, +`xAI` + `grok-4.5` becomes `spacexai/grok-4.5`). Already-qualified ids remain +unchanged. Set `AI_GATEWAY_API_KEY` or Vercel OIDC. This package does not load a provider SDK from +`config.provider.name`. Inject a model or model factory to call a provider +directly. + +All LaunchDarkly providers are handled explicitly. Anthropic, OpenAI, Azure, +Gemini, Cohere, DeepSeek, Meta, Mistral, Perplexity, and Vertex map directly. +Bedrock, Cortex, Cursor, Databricks, and Fireworks AI infer the creator from the +model and fail locally if ambiguous. AI21 Labs and IBM Watson fail locally +because Vercel's current catalog has no corresponding creator. + +## Use + +```python +from launchdarkly_ai_server import config +from launchdarkly_ai_vercel_agents import create_vercel_agents_handler + +result = await config( + key="my-agent-config", + handler=create_vercel_agents_handler(), +).invoke("Help me", {"kind": "user", "key": "user-123"}) +``` + +The package also exports `vercel_graph()` and `to_vercel_agents()` for +LaunchDarkly and framework-native graph execution. The native adapter builds one +`ai.Agent` per graph node and follows only a `transfer_to_` tool that the +current node actually selects. + +This package advertises `("*", "agent")`; register only one wildcard agent +adapter in a handler pool. diff --git a/packages/vercel-agents/agents.md b/packages/vercel-agents/agents.md new file mode 100644 index 00000000..8254ede2 --- /dev/null +++ b/packages/vercel-agents/agents.md @@ -0,0 +1,18 @@ +# Agent Guide — Vercel Agents + +This Tier 1 package provides the wildcard `("*", "agent")` handler. + +- Default models use a mapped Gateway `creator/model` id. Preserve + slash-qualified ids, map `xAI` to `spacexai`, and split recognized LD dotted + creator prefixes once. Injected models are the only Gateway opt-out. +- `handler.py` constructs the native `ai.Agent` and maps models, messages, + `AgentTool`/`Tool`/`ToolSpec`, request parameters, streams, usage, and spans. +- `graph.py` pre-wires exactly one Vercel handler into the core graph API. +- `native_graph.py` creates one native agent per node. Handoff tools only select + a target; execution follows the selected target after the current run ends. +- Root caller history is never replayed into child nodes. +- Native graph spans end in `finally`, including `CancelledError`. +- Runtime calls are always async context-managed and never enable duplicate + runtime telemetry. + +`TESTING.md` §2.x and Appendix A.14 are authoritative. diff --git a/packages/vercel-agents/pyproject.toml b/packages/vercel-agents/pyproject.toml new file mode 100644 index 00000000..8ec1db6c --- /dev/null +++ b/packages/vercel-agents/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "launchdarkly-ai-vercel-agents" +version = "0.1.0" +requires-python = ">=3.12" +dependencies = [ + "launchdarkly-ai-server", + "opentelemetry-api>=1.25", + "ai>=0.7", +] +description = "Vercel AI SDK agent handler for LaunchDarkly AI SDK" +readme = "README.md" +license = "Apache-2.0" +authors = [{name = "LaunchDarkly", email = "team@launchdarkly.com"}] +keywords = ["launchdarkly", "ai", "vercel", "agents"] + +[build-system] +requires = ["hatchling>=1.32,<2"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/launchdarkly_ai_vercel_agents"] diff --git a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/__init__.py b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/__init__.py new file mode 100644 index 00000000..651cbfcd --- /dev/null +++ b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/__init__.py @@ -0,0 +1,18 @@ +"""LaunchDarkly AI SDK integration for native Vercel AI SDK agents.""" + +__version__ = "0.1.0" # x-release-please-version + +from launchdarkly_ai_server import register_ai_sdk_package + +from .graph import vercel_graph +from .handler import create_vercel_agents_handler, vercel_agents +from .native_graph import to_vercel_agents + +__all__ = [ + "create_vercel_agents_handler", + "to_vercel_agents", + "vercel_agents", + "vercel_graph", +] + +register_ai_sdk_package("launchdarkly-ai-vercel-agents", __version__) diff --git a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/graph.py b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/graph.py new file mode 100644 index 00000000..d5babf20 --- /dev/null +++ b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/graph.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +from typing import Any + +from launchdarkly_ai_server import graph + +from .handler import ModelSource, create_vercel_agents_handler + + +def vercel_graph( + key: str, + model: ModelSource | None = None, + *, + capture_content: bool = False, + **options: Any, +) -> Any: + """Create a LaunchDarkly graph with one wildcard Vercel handler.""" + options.pop("handlers", None) + handler = create_vercel_agents_handler(model=model, capture_content=capture_content) + return graph(key, handlers=[handler], **options) diff --git a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/handler.py b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/handler.py new file mode 100644 index 00000000..fe43b862 --- /dev/null +++ b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/handler.py @@ -0,0 +1,453 @@ +from __future__ import annotations + +import base64 +import inspect +import json +from collections.abc import AsyncGenerator, Callable +from typing import Any + +import ai +from ai.types.tools import ToolSpec +from opentelemetry import trace +from opentelemetry.trace import StatusCode +from pydantic import BaseModel, ConfigDict, create_model + +from launchdarkly_ai_server import ( + AiConfigRep, + LDContext, + NativeTool, + ProviderHandler, + compose_history, + config, + create_handler, + parse_template, + set_ld_span_attributes, +) + +from .model_id import gateway_model_id + +OWNED_PARAMETERS = { + "model", + "messages", + "prompt", + "system", + "tools", + "stream", + "output", + "outputformat", + "stopwhen", + "maxsteps", + "apikey", + "baseurl", +} +ModelSource = Any | Callable[[AiConfigRep], Any] + + +def _owned(name: str) -> bool: + return name.replace("_", "").lower() in OWNED_PARAMETERS + + +def build_request_params(config: AiConfigRep, runtime: Any = None) -> Any: + runtime = runtime or ai + raw = { + key: value + for key, value in (config.get("model", {}).get("parameters") or {}).items() + if not _owned(key) + } + sampling: dict[type[Any], Any] = {} + mappings = { + "temperature": ("TemperatureSamplerParams", "temperature"), + "top_p": ("TopPSamplerParams", "top_p"), + "topP": ("TopPSamplerParams", "top_p"), + "top_k": ("TopKSamplerParams", "top_k"), + "topK": ("TopKSamplerParams", "top_k"), + "min_p": ("MinPSamplerParams", "min_p"), + "minP": ("MinPSamplerParams", "min_p"), + "repetition_penalty": ( + "RepetitionPenaltyParams", + "repetition_penalty", + ), + "repetitionPenalty": ( + "RepetitionPenaltyParams", + "repetition_penalty", + ), + "seed": ("SeedSamplerParams", "seed"), + } + for key, (class_name, argument) in mappings.items(): + if key in raw: + cls = getattr(runtime, class_name) + sampling[cls] = cls(**{argument: raw.pop(key)}) + kwargs: dict[str, Any] = {} + if sampling: + kwargs["sampling"] = sampling + max_tokens = raw.pop("max_tokens", raw.pop("maxTokens", None)) + if max_tokens is not None: + kwargs["output"] = runtime.OutputParams(max_tokens=max_tokens) + reasoning_effort = raw.pop("reasoning_effort", raw.pop("reasoningEffort", None)) + if reasoning_effort is not None: + kwargs["reasoning"] = runtime.ReasoningParams(effort=reasoning_effort) + for direct in ("metadata", "safety_identifier", "extra_headers", "extra_query"): + if direct in raw: + kwargs[direct] = raw.pop(direct) + if raw: + kwargs["extra_body"] = raw + return runtime.InferenceRequestParams(**kwargs) + + +async def resolve_model( + source: ModelSource | None, cfg: AiConfigRep, runtime: Any = None +) -> Any: + runtime = runtime or ai + if source is None: + return runtime.get_model(gateway_model_id(cfg)) + value = source(cfg) if callable(source) else source + return await value if inspect.isawaitable(value) else value + + +def _image_part(block: dict[str, Any], runtime: Any) -> Any: + source = block.get("source") or {} + value: str | bytes + if source.get("type") == "base64": + value = base64.b64decode(source.get("data", "")) + else: + value = source.get("url") or block.get("url") or "" + return runtime.file_part(value, media_type=source.get("media_type")) + + +def _content(content: Any, runtime: Any) -> Any: + if not isinstance(content, list): + return content if isinstance(content, str) else "" + parts: list[Any] = [] + for block in content: + if block.get("type") == "text": + parts.append(block.get("text", "")) + elif block.get("type") in {"image", "file"}: + parts.append(_image_part(block, runtime)) + return parts + + +def build_messages( + cfg: AiConfigRep, + user_input: str | None, + variables: dict[str, Any], + history: list[dict[str, Any]] | None, + *, + include_instructions: bool = True, + runtime: Any = None, +) -> list[Any]: + runtime = runtime or ai + result: list[Any] = [] + config_messages: list[dict[str, Any]] = [] + if cfg.get("instructions"): + if include_instructions: + result.append( + runtime.system_message(parse_template(cfg["instructions"], variables)) + ) + else: + for message in cfg.get("messages") or []: + content = message.get("content", "") + mapped = ( + parse_template(content, variables) + if isinstance(content, str) + else content + ) + if message.get("role") == "system": + if include_instructions: + result.append(runtime.system_message(mapped)) + else: + config_messages.append({**message, "content": mapped}) + for turn in compose_history( + history=history or [], + user_input=user_input, + config_messages=config_messages, + ): + mapped_content = _content(turn.get("content", ""), runtime) + if turn.get("role") == "assistant": + args = ( + mapped_content if isinstance(mapped_content, list) else [mapped_content] + ) + result.append(runtime.assistant_message(*args)) + elif turn.get("role") == "user": + args = ( + mapped_content if isinstance(mapped_content, list) else [mapped_content] + ) + result.append(runtime.user_message(*args)) + return result + + +async def _call_tool(handler: Callable[..., Any], kwargs: dict[str, Any]) -> Any: + value = handler(kwargs) + return await value if inspect.isawaitable(value) else value + + +def build_agent_tools( + definitions: dict[str, Any] | None, + handlers: dict[str, Any] | None, + runtime: Any = None, +) -> list[Any]: + runtime = runtime or ai + result: list[Any] = [] + for name, definition in (definitions or {}).items(): + handler = (handlers or {}).get(name) + if not callable(handler) or isinstance(handler, NativeTool): + continue + + async def execute(_handler: Callable[..., Any] = handler, **kwargs: Any) -> Any: + return await _call_tool(_handler, kwargs) + + spec = ToolSpec( + description=definition.get("description"), + params=definition.get("parameters") or {}, + ) + model_tool = runtime.Tool(kind="function", name=name, spec=spec) + agent_tool = runtime.AgentTool(model_tool, execute) + try: + object.__setattr__(agent_tool, "name", name) + object.__setattr__(agent_tool, "execute", execute) + object.__setattr__(agent_tool, "input_schema", spec.params) + object.__setattr__(agent_tool, "description", spec.description) + except (AttributeError, TypeError): + pass + result.append(agent_tool) + return result + + +def _usage_values(usage: Any) -> tuple[int, int]: + if usage is None: + return 0, 0 + input_tokens = getattr(usage, "input_tokens", None) + output_tokens = getattr(usage, "output_tokens", None) + if input_tokens is None: + input_tokens = getattr(usage, "input", 0) + if output_tokens is None: + output_tokens = getattr(usage, "output", 0) + return int(input_tokens or 0), int(output_tokens or 0) + + +def usage_of(stream: Any) -> dict[str, int]: + """Total token usage for an agent run. + + ``AgentStream`` carries no aggregate usage — each message in the run holds its + own — so a tool loop only reports true totals when the per-message values are + summed rather than read off the stream. + """ + messages = getattr(stream, "messages", None) + if messages: + totals = [ + _usage_values(getattr(message, "usage", None)) for message in messages + ] + return { + "input_tokens": sum(total[0] for total in totals), + "output_tokens": sum(total[1] for total in totals), + } + input_tokens, output_tokens = _usage_values(getattr(stream, "usage", None)) + return {"input_tokens": input_tokens, "output_tokens": output_tokens} + + +def output_of(stream: Any) -> str: + """Final text of a completed run. + + ``AgentStream`` exposes ``output`` and no ``text``, so reading ``text`` + defensively would turn a real response into an empty one. + """ + output = stream.output + if isinstance(output, BaseModel): + return output.model_dump_json() + if isinstance(output, str): + return output + return json.dumps(output) + + +def _python_type(schema: dict[str, Any], name: str) -> Any: + kind = schema.get("type") + if kind == "integer": + return int + if kind == "number": + return float + if kind == "boolean": + return bool + if kind == "array": + items = schema.get("items") + item_type = ( + _python_type(items, f"{name}Item") if isinstance(items, dict) else str + ) + return list[item_type] # type: ignore[valid-type] + if kind == "object" or schema.get("properties"): + return build_output_type(schema, name=f"{name}Object") + # Strict structured output rejects an untyped member, so fall back to a string + # rather than emitting a schema the provider refuses. + return str + + +def build_output_type( + schema: dict[str, Any], *, name: str = "VercelAgentOutput" +) -> type[BaseModel]: + """Build the Pydantic model the run validates its final answer against. + + Every field is required and extras are forbidden because providers running + strict structured output reject a schema that allows either. + """ + fields: dict[str, Any] = { + field: (_python_type(field_schema, field.title().replace("_", "")), ...) + for field, field_schema in (schema.get("properties") or {}).items() + } + return create_model(name, __config__=ConfigDict(extra="forbid"), **fields) + + +def text_delta(event: Any) -> str | None: + if getattr(event, "kind", None) == "text_delta": + return str(getattr(event, "chunk", "")) + events = getattr(ai, "events", None) + delta_type = getattr(events, "TextDelta", None) + if delta_type is not None and isinstance(event, delta_type): + return str(event.chunk) + return None + + +def _start_span(cfg: AiConfigRep, variables: dict[str, Any]) -> Any: + span = trace.get_tracer("@launchdarkly/ai-vercel-agents").start_span("invoke_agent") + provider = str(cfg.get("provider", {}).get("name") or "vercel").lower() + model = str(cfg.get("model", {}).get("name") or "") + span.set_attribute("gen_ai.operation.name", "invoke_agent") + span.set_attribute("gen_ai.system", provider) + span.set_attribute("gen_ai.provider.name", provider) + span.set_attribute("gen_ai.request.model", model) + set_ld_span_attributes(span, variables) + return span + + +def _start_model_span(cfg: AiConfigRep, root: Any) -> Any: + model = str(cfg.get("model", {}).get("name") or "") + provider = str(cfg.get("provider", {}).get("name") or "vercel").lower() + span = trace.get_tracer("@launchdarkly/ai-vercel-agents").start_span( + f"chat {model}", context=trace.set_span_in_context(root) + ) + span.set_attribute("gen_ai.operation.name", "chat") + span.set_attribute("gen_ai.system", provider) + span.set_attribute("gen_ai.provider.name", provider) + span.set_attribute("gen_ai.request.model", model) + return span + + +def _set_usage(span: Any, usage: dict[str, int]) -> None: + input_tokens = usage["input_tokens"] + output_tokens = usage["output_tokens"] + span.set_attribute("gen_ai.usage.input_tokens", input_tokens) + span.set_attribute("gen_ai.usage.output_tokens", output_tokens) + span.set_attribute("gen_ai.usage.total_tokens", input_tokens + output_tokens) + span.set_attribute("gen_ai.usage.prompt_tokens", input_tokens) + span.set_attribute("gen_ai.usage.completion_tokens", output_tokens) + span.set_attribute("gen_ai.usage.cache_read.input_tokens", 0) + span.set_attribute("gen_ai.usage.cache_creation.input_tokens", 0) + + +def create_vercel_agents_handler( + model: ModelSource | None = None, + *, + capture_content: bool = False, +) -> ProviderHandler: + async def run( + cfg: AiConfigRep, + user_input: str | None = None, + tool_handlers: dict[str, Any] | None = None, + variables: dict[str, Any] | None = None, + history: list[dict[str, Any]] | None = None, + ) -> dict[str, Any]: + vs = variables or {} + span = _start_span(cfg, vs) + model_span = _start_model_span(cfg, span) + try: + agent = ai.Agent(tools=build_agent_tools(cfg.get("tools"), tool_handlers)) + run_kwargs: dict[str, Any] = { + "model": await resolve_model(model, cfg), + "messages": build_messages(cfg, user_input, vs, history), + "params": build_request_params(cfg), + } + if cfg.get("outputFormat"): + run_kwargs["output_type"] = build_output_type(cfg["outputFormat"]) + async with agent.run(**run_kwargs) as provider_stream: + async for _ in provider_stream: + pass + usage = usage_of(provider_stream) + output = output_of(provider_stream) + _set_usage(model_span, usage) + _set_usage(span, usage) + model_span.set_status(StatusCode.OK) + span.set_status(StatusCode.OK) + return {"output": output, "usage": usage} + except BaseException as exc: + if isinstance(exc, Exception): + model_span.record_exception(exc) + model_span.set_status(StatusCode.ERROR, str(exc)) + span.record_exception(exc) + span.set_status(StatusCode.ERROR, str(exc)) + raise + finally: + model_span.end() + span.end() + + async def stream( + cfg: AiConfigRep, + user_input: str | None = None, + tool_handlers: dict[str, Any] | None = None, + variables: dict[str, Any] | None = None, + history: list[dict[str, Any]] | None = None, + ) -> AsyncGenerator[dict[str, Any], None]: + vs = variables or {} + span = _start_span(cfg, vs) + model_span = _start_model_span(cfg, span) + completed = False + provider_stream: Any = None + try: + agent = ai.Agent(tools=build_agent_tools(cfg.get("tools"), tool_handlers)) + async with agent.run( + model=await resolve_model(model, cfg), + messages=build_messages(cfg, user_input, vs, history), + params=build_request_params(cfg), + ) as provider_stream: + async for event in provider_stream: + text = text_delta(event) + if text is not None: + yield {"type": "chunk", "text": text} + completed = True + usage = usage_of(provider_stream) + _set_usage(model_span, usage) + _set_usage(span, usage) + model_span.set_status(StatusCode.OK) + span.set_status(StatusCode.OK) + yield { + "type": "done", + "output": output_of(provider_stream), + "usage": usage, + } + except BaseException as exc: + if isinstance(exc, Exception): + model_span.record_exception(exc) + model_span.set_status(StatusCode.ERROR, str(exc)) + span.record_exception(exc) + span.set_status(StatusCode.ERROR, str(exc)) + raise + finally: + if not completed: + model_span.set_attribute("launchdarkly.stream.abandoned", True) + span.set_attribute("launchdarkly.stream.abandoned", True) + model_span.end() + span.end() + + return create_handler(("*", "agent"), run, stream, capture_content=capture_content) + + +def vercel_agents( + config_key: str, + user_input: str | None, + context: LDContext, + *, + model: ModelSource | None = None, + capture_content: bool = False, + variables: dict[str, Any] | None = None, + **options: Any, +) -> Any: + handler = create_vercel_agents_handler(model=model, capture_content=capture_content) + return config(key=config_key, handler=handler, **options).invoke( + user_input, context, variables=variables + ) diff --git a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/model_id.py b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/model_id.py new file mode 100644 index 00000000..8a23b660 --- /dev/null +++ b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/model_id.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +import re + +from launchdarkly_ai_server import AiConfigRep + +GATEWAY_CREATORS = { + "alibaba": "alibaba", + "amazon": "amazon", + "anthropic": "anthropic", + "arceeai": "arcee-ai", + "bfl": "bfl", + "bytedance": "bytedance", + "cohere": "cohere", + "deepseek": "deepseek", + "fishaudio": "fish-audio", + "gemini": "google", + "google": "google", + "googleai": "google", + "inception": "inception", + "inclusionai": "inclusionai", + "inferencenet": "inference-net", + "interfaze": "interfaze", + "klingai": "klingai", + "meta": "meta", + "minimax": "minimax", + "mistral": "mistral", + "mistralai": "mistral", + "mixedbread": "mixedbread", + "moonshotai": "moonshotai", + "morph": "morph", + "nvidia": "nvidia", + "openai": "openai", + "perplexity": "perplexity", + "poolside": "poolside", + "prodia": "prodia", + "quiverai": "quiverai", + "recraft": "recraft", + "sakana": "sakana", + "spacexai": "spacexai", + "stepfun": "stepfun", + "tencent": "tencent", + "thinkingmachines": "thinkingmachines", + "typesafe": "typesafe-ai", + "typesafeai": "typesafe-ai", + "voyage": "voyage", + "xiaomi": "xiaomi", + "xai": "spacexai", + "zai": "zai", +} + +PROVIDER_CREATORS: dict[str, str | None] = { + "anthropic": "anthropic", + "openai": "openai", + "bedrock": None, + "azure": "openai", + "gemini": "google", + "ai21labs": None, + "cohere": "cohere", + "cortex": None, + "cursor": None, + "databricks": None, + "deepseek": "deepseek", + "fireworksai": None, + "ibmwatson": None, + "meta": "meta", + "mistral": "mistral", + "perplexity": "perplexity", + "vertex": "google", + # Compatibility aliases found in existing configs. + "google": "google", + "googleai": "google", + "mistralai": "mistral", + "spacexai": "spacexai", + "typesafe": "typesafe-ai", + "typesafeai": "typesafe-ai", + "xai": "spacexai", +} + +MODEL_FAMILY_CREATORS = ( + (re.compile(r"^(?:gpt|o[1-9])(?:[-.]|$)", re.I), "openai"), + (re.compile(r"^claude(?:[-.]|$)", re.I), "anthropic"), + (re.compile(r"^gemini(?:[-.]|$)", re.I), "google"), + (re.compile(r"^grok(?:[-.]|$)", re.I), "spacexai"), + (re.compile(r"^(?:command|aya)(?:[-.]|$)", re.I), "cohere"), + (re.compile(r"^deepseek(?:[-.]|$)", re.I), "deepseek"), + (re.compile(r"^llama(?:[-.]|$)", re.I), "meta"), + ( + re.compile(r"^(?:mistral|mixtral|codestral|pixtral)(?:[-.]|$)", re.I), + "mistral", + ), + (re.compile(r"^sonar(?:[-.]|$)", re.I), "perplexity"), + (re.compile(r"^(?:nova|titan)(?:[-.]|$)", re.I), "amazon"), + (re.compile(r"^qwen(?:[-.]|$)", re.I), "alibaba"), +) + + +def _provider_key(value: str) -> str: + return re.sub(r"[^a-z0-9]", "", value.lower()) + + +def _creator_from_model(model: str) -> tuple[str, str] | None: + parts = model.split(".") + for index, part in enumerate(parts[:-1]): + creator = GATEWAY_CREATORS.get(_provider_key(part)) + if creator: + return creator, ".".join(parts[index + 1 :]) + for pattern, creator in MODEL_FAMILY_CREATORS: + if pattern.search(model): + return creator, model + return None + + +def gateway_model_id(config: AiConfigRep) -> str: + """Build the AI Gateway ``creator/model`` id for an evaluated config.""" + model = str(config["model"]["name"]) + if "/" in model: + return model + + provider = str((config.get("provider") or {}).get("name") or "") + provider_name = _provider_key(provider) + inferred = _creator_from_model(model) + if inferred: + return f"{inferred[0]}/{inferred[1]}" + + creator = PROVIDER_CREATORS.get(provider_name) + if creator: + return f"{creator}/{model}" + if provider_name in {"ai21labs", "ibmwatson"}: + raise ValueError( + f'Vercel AI Gateway currently exposes no models created by "{provider}". ' + "Inject a direct provider model with model/model_factory instead." + ) + if provider_name in PROVIDER_CREATORS: + raise ValueError( + f'LaunchDarkly provider "{provider}" hosts models from multiple ' + f'creators, so "{model}" cannot be converted to a Vercel creator/model ' + "id. Store an explicit creator/model id or inject a model/model_factory." + ) + else: + raise ValueError( + f'Cannot map LaunchDarkly provider "{provider or "unknown"}" to a ' + "Vercel AI Gateway creator. Pass a creator/model id in " + "config.model.name or inject a model/model factory." + ) diff --git a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/native_graph.py b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/native_graph.py new file mode 100644 index 00000000..b7f47313 --- /dev/null +++ b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/native_graph.py @@ -0,0 +1,247 @@ +from __future__ import annotations + +import time +import uuid +from types import SimpleNamespace +from typing import Any + +import ai +from ai.types.tools import ToolSpec +from opentelemetry import trace +from opentelemetry.trace import StatusCode + +from launchdarkly_ai_server import ( + GraphDefinition, + GraphNode, + get_client, + make_track_data, + to_ld_context, +) + +from .handler import ( + ModelSource, + build_agent_tools, + build_messages, + build_request_params, + output_of, + resolve_model, + usage_of, +) + + +def _handoff_tool( + source: str, + target: str, + description: str | None, + selected: dict[str, str | None], +) -> Any: + async def transfer() -> str: + selected[source] = target + return target + + model_tool = ai.Tool( + kind="function", + name=f"transfer_to_{target}", + spec=ToolSpec( + description=description, params={"type": "object", "properties": {}} + ), + ) + agent_tool = ai.AgentTool(model_tool, transfer) + try: + object.__setattr__(agent_tool, "name", f"transfer_to_{target}") + object.__setattr__(agent_tool, "execute", transfer) + except (AttributeError, TypeError): + pass + return agent_tool + + +def _track( + name: str, + context: Any, + node: GraphNode, + definition: GraphDefinition, + run_id: str, + value: int | float, +) -> None: + client = get_client() + client.track( + name, + to_ld_context(client, context), + make_track_data(node, definition.key, run_id), + value, + ) + + +def to_vercel_agents( + definition_promise: Any, + opts: dict[str, Any] | None = None, +) -> Any: + options = opts or {} + + async def invoke( + input_text: str = "", + variables: dict[str, Any] | None = None, + history: list[dict[str, Any]] | None = None, + ) -> dict[str, Any]: + definition: GraphDefinition = await definition_promise + if not definition.enabled: + raise ValueError(f'Agent graph "{definition.key}" is disabled') + if definition.root is None: + raise ValueError(f'Graph "{definition.key}" has no root node') + + vs = variables or {} + tool_handlers = options.get("tool_handlers") or {} + model_source: ModelSource | None = options.get( + "model", options.get("model_factory") + ) + selected: dict[str, str | None] = {} + agents: dict[str, Any] = {} + models: dict[str, Any] = {} + reachable: list[GraphNode] = [] + visited: set[str] = set() + + async def build(node: GraphNode) -> None: + if node.key in visited: + return + visited.add(node.key) + reachable.append(node) + selected[node.key] = None + model = await resolve_model(model_source, node.config, ai) + models[node.key] = model + tools = build_agent_tools(node.config.get("tools"), tool_handlers, ai) + for edge in definition.edges_from(node.key): + tools.append( + _handoff_tool( + node.key, + edge.target_key, + (edge.handoff or {}).get("description"), + selected, + ) + ) + agents[node.key] = ai.Agent(tools=tools) + for edge in definition.edges_from(node.key): + child = definition.get_node(edge.target_key) + if child is not None: + await build(child) + + await build(definition.root) + + span = trace.get_tracer("@launchdarkly/ai-vercel-agents").start_span( + "ld.ai.graph" + ) + span.set_attribute("ld.ai.graph.key", definition.key) + context = options.get("context") + run_id = str(uuid.uuid4()) + start = time.monotonic() + path: list[str] = [] + input_tokens = 0 + output_tokens = 0 + current = definition.root + current_input = input_text + final_text = "" + try: + while current is not None and current.key not in path: + path.append(current.key) + selected[current.key] = None + node_history = history if current.key == definition.root.key else None + messages = build_messages( + current.config, + current_input, + vs, + node_history, + runtime=ai, + ) + async with agents[current.key].run( + model=models[current.key], + messages=messages, + params=build_request_params(current.config, ai), + ) as stream: + async for _ in stream: + pass + usage = usage_of(stream) + input_tokens += usage["input_tokens"] + output_tokens += usage["output_tokens"] + final_text = output_of(stream) + if context is not None: + _track( + "$ld:ai:generation:success", + context, + current, + definition, + run_id, + 1, + ) + target = selected[current.key] + if target is None: + break + if context is not None: + _track( + "$ld:ai:graph:handoff_success", + context, + current, + definition, + run_id, + 1, + ) + next_node = definition.get_node(target) + if next_node is None: + break + current = next_node + current_input = final_text + + total = input_tokens + output_tokens + span.set_attribute("ld.ai.graph.path", "->".join(path)) + span.set_status(StatusCode.OK) + if context is not None: + _track( + "$ld:ai:graph:invocation_success", + context, + definition.root, + definition, + run_id, + 1, + ) + _track( + "$ld:ai:graph:duration:total", + context, + definition.root, + definition, + run_id, + (time.monotonic() - start) * 1000, + ) + _track( + "$ld:ai:graph:total_tokens", + context, + definition.root, + definition, + run_id, + total, + ) + return { + "response": final_text, + "usage": { + "input": input_tokens, + "output": output_tokens, + "total": total, + }, + } + except BaseException as exc: + if isinstance(exc, Exception): + span.record_exception(exc) + span.set_status(StatusCode.ERROR, str(exc)) + if context is not None: + _track( + "$ld:ai:graph:invocation_failure", + context, + definition.root, + definition, + run_id, + 1, + ) + else: + span.set_attribute("launchdarkly.run.cancelled", True) + raise + finally: + span.end() + + return SimpleNamespace(invoke=invoke) diff --git a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/py.typed b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/py.typed new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/py.typed @@ -0,0 +1 @@ + diff --git a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/spans.py b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/spans.py new file mode 100644 index 00000000..d782f2b2 --- /dev/null +++ b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/spans.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from typing import Any + +from opentelemetry import trace +from opentelemetry.trace import StatusCode + +from launchdarkly_ai_server import ( + AiConfigRep, + SpanUsage, + set_ld_span_attributes, + set_model_identity_attributes, + set_usage_span_attributes, +) + +TRACER_NAME = "@launchdarkly/ai-vercel-agents" + + +def model_name(config: AiConfigRep) -> str: + return str(config.get("model", {}).get("name") or "") + + +def serving_provider(config: AiConfigRep) -> str: + return str(config.get("provider", {}).get("name") or "vercel").lower() + + +def start_root_span(config: AiConfigRep, variables: dict[str, Any]) -> Any: + span = trace.get_tracer(TRACER_NAME).start_span("invoke_agent") + span.set_attribute("gen_ai.operation.name", "invoke_agent") + set_model_identity_attributes(span, serving_provider(config), model_name(config)) + set_ld_span_attributes(span, variables) + return span + + +def parent_context_of(span: Any) -> Any: + return trace.set_span_in_context(span) if span is not None else None + + +def start_model_span(config: AiConfigRep, parent: Any) -> Any: + span = trace.get_tracer(TRACER_NAME).start_span( + f"chat {model_name(config)}", context=parent + ) + span.set_attribute("gen_ai.operation.name", "chat") + set_model_identity_attributes(span, serving_provider(config), model_name(config)) + return span + + +def start_tool_span(name: str, call_id: str, parent: Any) -> Any: + span = trace.get_tracer(TRACER_NAME).start_span( + f"execute_tool {name}", context=parent + ) + span.set_attribute("gen_ai.operation.name", "execute_tool") + span.set_attribute("gen_ai.tool.name", name) + span.set_attribute("gen_ai.tool.call.id", call_id) + return span + + +def finish_root_span(span: Any, response_model: str, usage: SpanUsage) -> None: + span.set_attribute("gen_ai.response.model", response_model) + set_usage_span_attributes(span, usage) + + +def succeed_span(span: Any) -> None: + span.set_status(StatusCode.OK) + span.end() + + +def mark_ok(span: Any) -> None: + span.set_status(StatusCode.OK) + + +def fail_span(span: Any, error: BaseException, tracker: set[int] | None = None) -> None: + span.record_exception(error) + span.set_status(StatusCode.ERROR, str(error)) + if tracker is None or id(span) not in tracker: + if tracker is not None: + tracker.add(id(span)) + span.end() diff --git a/packages/vercel-agents/tests/test_handler.py b/packages/vercel-agents/tests/test_handler.py new file mode 100644 index 00000000..12aa97db --- /dev/null +++ b/packages/vercel-agents/tests/test_handler.py @@ -0,0 +1,368 @@ +"""Test-first contract for the wildcard native Vercel ``ai.Agent`` adapter.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from pydantic import BaseModel + +import launchdarkly_ai_vercel_agents.graph as graph_mod +import launchdarkly_ai_vercel_agents.handler as handler_mod +from launchdarkly_ai_vercel_agents import ( + create_vercel_agents_handler, + vercel_graph, +) + +CONFIG: dict[str, Any] = { + "model": { + "name": "anthropic/claude-sonnet-4", + "parameters": {"temperature": 0.2}, + }, + "provider": {"name": "Anthropic"}, + "instructions": "Be concise.", +} + + +class TextDelta: + def __init__(self, chunk: str) -> None: + self.chunk = chunk + + +class AgentStream: + """Mirrors ``ai.agents.AgentStream``. + + The real type exposes ``output`` and per-message usage, and has no ``text`` + or aggregate ``usage``; a double carrying those would hide a handler that + reads them and silently returns an empty response. + """ + + def __init__( + self, + events: list[Any] | None = None, + *, + text: str = "answer", + input_tokens: int = 8, + output_tokens: int = 3, + turns: int = 1, + ) -> None: + self.events = list(events or [TextDelta(text)]) + self.output = text + self.messages = [ + SimpleNamespace( + usage=SimpleNamespace( + input_tokens=input_tokens, output_tokens=output_tokens + ) + ) + for _ in range(turns) + ] + self.entered = 0 + self.exited = 0 + self.consumed = 0 + + async def __aenter__(self) -> AgentStream: + self.entered += 1 + return self + + async def __aexit__(self, *args: Any) -> None: + self.exited += 1 + + def __aiter__(self) -> AgentStream: + return self + + async def __anext__(self) -> Any: + if not self.events: + raise StopAsyncIteration + self.consumed += 1 + return self.events.pop(0) + + +@pytest.fixture +def ai_runtime() -> MagicMock: + runtime = MagicMock() + runtime.events.TextDelta = TextDelta + runtime.get_model = MagicMock(return_value=object()) + runtime.system_message = MagicMock( + side_effect=lambda content: {"role": "system", "content": content} + ) + runtime.user_message = MagicMock( + side_effect=lambda *content: { + "role": "user", + "content": content[0] if len(content) == 1 else list(content), + } + ) + runtime.assistant_message = MagicMock( + side_effect=lambda *content: { + "role": "assistant", + "content": content[0] if len(content) == 1 else list(content), + } + ) + runtime.file_part = MagicMock( + side_effect=lambda value, media_type=None: { + "type": "file", + "value": value, + "media_type": media_type, + } + ) + runtime.tool = MagicMock( + side_effect=lambda fn=None, **kwargs: ( + SimpleNamespace(execute=fn, **kwargs) if fn is not None else kwargs + ) + ) + runtime.InferenceRequestParams = MagicMock(return_value=object()) + agent = MagicMock() + agent.run = MagicMock(return_value=AgentStream()) + runtime.Agent = MagicMock(return_value=agent) + with patch.object(handler_mod, "ai", runtime): + yield runtime + + +class TestNativeAgent: + def test_advertises_wildcard_agent(self) -> None: + assert create_vercel_agents_handler().provides_for == ("*", "agent") + + @pytest.mark.asyncio + async def test_constructs_native_ai_agent_and_runs_it( + self, ai_runtime: MagicMock + ) -> None: + await create_vercel_agents_handler()(CONFIG, "hello") + ai_runtime.Agent.assert_called_once() + agent = ai_runtime.Agent.return_value + agent.run.assert_called_once() + assert agent.run.call_args.kwargs["model"] is ai_runtime.get_model.return_value + + @pytest.mark.asyncio + async def test_does_not_fall_back_to_messages_stream( + self, ai_runtime: MagicMock + ) -> None: + ai_runtime.stream = MagicMock() + await create_vercel_agents_handler()(CONFIG, "hello") + ai_runtime.Agent.assert_called_once() + ai_runtime.stream.assert_not_called() + + @pytest.mark.asyncio + async def test_preserves_gateway_model_id(self, ai_runtime: MagicMock) -> None: + await create_vercel_agents_handler()(CONFIG, "hello") + ai_runtime.get_model.assert_called_once_with("anthropic/claude-sonnet-4") + + @pytest.mark.asyncio + async def test_builds_gateway_creator_model_id(self, ai_runtime: MagicMock) -> None: + config = { + **CONFIG, + "model": {"name": "grok-4.5"}, + "provider": {"name": "xAI"}, + } + await create_vercel_agents_handler()(config, "hello") + ai_runtime.get_model.assert_called_once_with("spacexai/grok-4.5") + + @pytest.mark.asyncio + async def test_converts_dotted_creator_prefix(self, ai_runtime: MagicMock) -> None: + config = { + **CONFIG, + "model": {"name": "openai.gpt-5.6-sol"}, + "provider": {"name": "OpenAI"}, + } + await create_vercel_agents_handler()(config, "hello") + ai_runtime.get_model.assert_called_once_with("openai/gpt-5.6-sol") + + @pytest.mark.asyncio + async def test_model_factory_receives_config_once( + self, ai_runtime: MagicMock + ) -> None: + model = object() + factory = AsyncMock(return_value=model) + await create_vercel_agents_handler(model=factory)(CONFIG, "hello") + factory.assert_awaited_once_with(CONFIG) + assert ai_runtime.Agent.return_value.run.call_args.kwargs["model"] is model + + @pytest.mark.asyncio + async def test_system_history_and_input_use_native_messages( + self, ai_runtime: MagicMock + ) -> None: + history = [ + {"role": "user", "content": "earlier"}, + {"role": "assistant", "content": "reply"}, + ] + await create_vercel_agents_handler()(CONFIG, "latest", {}, {}, history) + messages = ai_runtime.Agent.return_value.run.call_args.kwargs["messages"] + assert [message["role"] for message in messages] == [ + "system", + "user", + "assistant", + "user", + ] + assert messages[-1]["content"] == "latest" + + @pytest.mark.asyncio + async def test_multimodal_history_uses_file_part( + self, ai_runtime: MagicMock + ) -> None: + history = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + { + "type": "image", + "source": { + "type": "url", + "url": "https://example.invalid/image.png", + "media_type": "image/png", + }, + }, + ], + } + ] + await create_vercel_agents_handler()(CONFIG, None, {}, {}, history) + ai_runtime.file_part.assert_called_once() + assert "example.invalid/image.png" in str(ai_runtime.file_part.call_args) + + @pytest.mark.asyncio + async def test_only_callable_tools_reach_agent(self, ai_runtime: MagicMock) -> None: + config = { + **CONFIG, + "tools": { + "lookup": { + "description": "Look up a value", + "parameters": {"type": "object"}, + }, + "missing": {"description": "No implementation", "parameters": {}}, + }, + } + lookup = AsyncMock(return_value="found") + await create_vercel_agents_handler()(config, "hello", {"lookup": lookup}) + tools = ai_runtime.Agent.call_args.kwargs["tools"] + assert [tool.name for tool in tools] == ["lookup"] + assert await tools[0].execute(key="x") == "found" + lookup.assert_awaited_once_with({"key": "x"}) + + @pytest.mark.asyncio + async def test_agent_receives_model_settings_as_request_params( + self, ai_runtime: MagicMock + ) -> None: + await create_vercel_agents_handler()(CONFIG, "hello") + run_kwargs = ai_runtime.Agent.return_value.run.call_args.kwargs + assert run_kwargs["params"] is ai_runtime.InferenceRequestParams.return_value + ai_runtime.TemperatureSamplerParams.assert_called_once_with(temperature=0.2) + + @pytest.mark.asyncio + async def test_blocking_result_normalizes_text_and_usage( + self, ai_runtime: MagicMock + ) -> None: + ai_runtime.Agent.return_value.run.return_value = AgentStream( + text="done", input_tokens=21, output_tokens=5 + ) + result = await create_vercel_agents_handler()(CONFIG, "hello") + assert result == { + "output": "done", + "usage": {"input_tokens": 21, "output_tokens": 5}, + } + + @pytest.mark.asyncio + async def test_sums_usage_across_every_turn_of_a_tool_loop( + self, ai_runtime: MagicMock + ) -> None: + ai_runtime.Agent.return_value.run.return_value = AgentStream( + text="done", input_tokens=21, output_tokens=5, turns=3 + ) + result = await create_vercel_agents_handler()(CONFIG, "hello") + assert result["usage"] == {"input_tokens": 63, "output_tokens": 15} + + @pytest.mark.asyncio + async def test_structured_output_is_requested_and_serialized( + self, ai_runtime: MagicMock + ) -> None: + schema = { + "type": "object", + "properties": {"answer": {"type": "string"}}, + "required": ["answer"], + } + + class Answer(BaseModel): + answer: str + + stream = AgentStream(text="") + stream.output = Answer(answer="yes") + ai_runtime.Agent.return_value.run.return_value = stream + + result = await create_vercel_agents_handler()( + {**CONFIG, "outputFormat": schema}, "question" + ) + + run_kwargs = ai_runtime.Agent.return_value.run.call_args.kwargs + assert run_kwargs["output_type"] is not None + assert result["output"] == '{"answer":"yes"}' + + @pytest.mark.asyncio + async def test_stream_forwards_text_and_closes_context( + self, ai_runtime: MagicMock + ) -> None: + provider_stream = AgentStream([TextDelta("a"), TextDelta("b")], text="ab") + ai_runtime.Agent.return_value.run.return_value = provider_stream + events = [ + event + async for event in await create_vercel_agents_handler().stream( + CONFIG, "hello" + ) + ] + assert [event.get("text") for event in events[:-1]] == ["a", "b"] + assert events[-1]["type"] == "done" + assert provider_stream.exited == 1 + + @pytest.mark.asyncio + async def test_early_stream_exit_closes_context( + self, ai_runtime: MagicMock + ) -> None: + provider_stream = AgentStream([TextDelta("a"), TextDelta("b")]) + ai_runtime.Agent.return_value.run.return_value = provider_stream + stream = await create_vercel_agents_handler().stream(CONFIG, "hello") + async for _event in stream: + break + await stream.aclose() + assert provider_stream.exited == 1 + assert provider_stream.consumed == 1 + + +class TestGraphWrapper: + def test_prewires_exactly_one_wildcard_handler(self) -> None: + graph_instance = object() + with ( + patch.object(graph_mod, "graph", return_value=graph_instance) as graph_fn, + patch.object( + graph_mod, + "create_vercel_agents_handler", + wraps=create_vercel_agents_handler, + ) as factory, + ): + assert vercel_graph("graph-key") is graph_instance + factory.assert_called_once() + handlers = graph_fn.call_args.kwargs["handlers"] + assert len(handlers) == 1 + assert handlers[0].provides_for == ("*", "agent") + + def test_forwards_model_and_capture_content_to_factory(self) -> None: + model = object() + with ( + patch.object(graph_mod, "graph", return_value=object()), + patch.object( + graph_mod, + "create_vercel_agents_handler", + return_value=MagicMock(provides_for=("*", "agent")), + ) as factory, + ): + vercel_graph( + "graph-key", + model=model, + capture_content=True, + context={"kind": "user", "key": "u"}, + ) + factory.assert_called_once_with(model=model, capture_content=True) + + def test_caller_cannot_override_handlers(self) -> None: + with patch.object(graph_mod, "graph", return_value=object()) as graph_fn: + vercel_graph("graph-key", handlers=[MagicMock()]) + handlers = graph_fn.call_args.kwargs["handlers"] + assert len(handlers) == 1 + assert handlers[0].provides_for == ("*", "agent") diff --git a/packages/vercel-agents/tests/test_model_id.py b/packages/vercel-agents/tests/test_model_id.py new file mode 100644 index 00000000..90c3735c --- /dev/null +++ b/packages/vercel-agents/tests/test_model_id.py @@ -0,0 +1,47 @@ +import pytest + +from launchdarkly_ai_vercel_agents.model_id import gateway_model_id + + +def config(provider: str, model: str = "model") -> dict: + return {"provider": {"name": provider}, "model": {"name": model}} + + +@pytest.mark.parametrize( + ("provider", "model", "expected"), + [ + ("Anthropic", "model", "anthropic/model"), + ("OpenAI", "model", "openai/model"), + ("Bedrock", "anthropic.claude-sonnet-4", "anthropic/claude-sonnet-4"), + ("Azure", "model", "openai/model"), + ("Gemini", "model", "google/model"), + ("Cohere", "model", "cohere/model"), + ("Cortex", "llama-4-scout", "meta/llama-4-scout"), + ("Cursor", "claude-sonnet-4", "anthropic/claude-sonnet-4"), + ("Databricks", "llama-4-maverick", "meta/llama-4-maverick"), + ("DeepSeek", "model", "deepseek/model"), + ("Fireworks AI", "qwen-3-235b", "alibaba/qwen-3-235b"), + ("Meta", "model", "meta/model"), + ("Mistral", "model", "mistral/model"), + ("Perplexity", "model", "perplexity/model"), + ("Vertex", "model", "google/model"), + ], +) +def test_maps_every_supported_launchdarkly_provider( + provider: str, model: str, expected: str +) -> None: + assert gateway_model_id(config(provider, model)) == expected + + +@pytest.mark.parametrize("provider", ["AI21 Labs", "IBM Watson"]) +def test_fails_locally_for_unsupported_creator(provider: str) -> None: + with pytest.raises(ValueError, match="currently exposes no models created by"): + gateway_model_id(config(provider)) + + +def test_preserves_explicit_and_dotted_gateway_ids() -> None: + assert gateway_model_id(config("Bedrock", "amazon/nova-pro")) == "amazon/nova-pro" + assert ( + gateway_model_id(config("Bedrock", "us.anthropic.claude-sonnet-4")) + == "anthropic/claude-sonnet-4" + ) diff --git a/packages/vercel-agents/tests/test_native_graph.py b/packages/vercel-agents/tests/test_native_graph.py new file mode 100644 index 00000000..d0558afd --- /dev/null +++ b/packages/vercel-agents/tests/test_native_graph.py @@ -0,0 +1,448 @@ +"""Test-first contract for ``to_vercel_agents``. + +The ``ai`` runtime is fully mocked. The suite describes topology, handoff, +history, usage, telemetry, and cleanup without gateway credentials or I/O. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +import launchdarkly_ai_vercel_agents.native_graph as native_graph_mod +from launchdarkly_ai_server import GraphDefinition, GraphEdge, GraphNode +from launchdarkly_ai_vercel_agents.native_graph import to_vercel_agents + + +def _edge(source: str, target: str) -> GraphEdge: + return GraphEdge( + key=f"{source}-{target}", + source_key=source, + target_key=target, + handoff={"description": f"Transfer to {target}"}, + ) + + +def _node( + key: str, + *, + edges: list[GraphEdge] | None = None, + tools: dict[str, Any] | None = None, +) -> GraphNode: + node_edges = edges or [] + return GraphNode( + key=key, + config={ + "model": {"name": f"gateway/{key}", "parameters": {"temperature": 0.1}}, + "provider": {"name": "Gateway"}, + "instructions": f"instructions for {key}", + "tools": tools or {}, + }, + meta={"variationKey": f"variation-{key}", "version": 1}, + edges=node_edges, + is_terminal=not node_edges, + ) + + +def _graph( + *, + enabled: bool = True, + include_node_tool: bool = False, +) -> GraphDefinition: + root_edge = _edge("root", "leaf") + root = _node( + "root", + edges=[root_edge], + tools=( + {"lookup": {"description": "Look up", "parameters": {"type": "object"}}} + if include_node_tool + else None + ), + ) + leaf = _node("leaf") + nodes = {"root": root, "leaf": leaf} + + async def _traverse(fn: Any, ctx: Any = None) -> None: + for node in (root, leaf): + value = fn(node) + if hasattr(value, "__await__"): + await value + + async def _unused(*args: Any, **kwargs: Any) -> Any: + raise AssertionError("native graph must not call GraphDefinition.run_node") + + return GraphDefinition( + key="vercel-graph", + enabled=enabled, + root=root, + get_node=lambda key: nodes.get(key), + get_child_nodes=lambda key: [leaf] if key == "root" else [], + get_parent_nodes=lambda key: [root] if key == "leaf" else [], + terminal_nodes=lambda: [leaf], + is_terminal=lambda key: key == "leaf", + edges_from=lambda key: [root_edge] if key == "root" else [], + run_node=_unused, + route=_unused, + traverse=_traverse, + reverse_traverse=_traverse, + ) + + +async def _definition(value: GraphDefinition) -> GraphDefinition: + return value + + +class RunStream: + def __init__( + self, + text: str, + *, + input_tokens: int, + output_tokens: int, + on_enter: Any = None, + ) -> None: + # ``AgentStream`` exposes ``output``, not ``text``. + self.output = text + self.messages: list[Any] = [ + SimpleNamespace( + usage=SimpleNamespace( + input_tokens=input_tokens, output_tokens=output_tokens + ) + ) + ] + self.exited = 0 + self.on_enter = on_enter + + async def __aenter__(self) -> RunStream: + if self.on_enter is not None: + value = self.on_enter() + if hasattr(value, "__await__"): + await value + return self + + async def __aexit__(self, *args: Any) -> None: + self.exited += 1 + + def __aiter__(self) -> RunStream: + return self + + async def __anext__(self) -> Any: + raise StopAsyncIteration + + +@pytest.fixture +def ai_runtime() -> MagicMock: + runtime = MagicMock() + runtime.get_model = MagicMock(side_effect=lambda name: f"model:{name}") + runtime.system_message = MagicMock( + side_effect=lambda content: {"role": "system", "content": content} + ) + runtime.user_message = MagicMock( + side_effect=lambda *content: { + "role": "user", + "content": content[0] if len(content) == 1 else list(content), + } + ) + runtime.assistant_message = MagicMock( + side_effect=lambda *content: { + "role": "assistant", + "content": content[0] if len(content) == 1 else list(content), + } + ) + runtime.file_part = MagicMock( + side_effect=lambda value, media_type=None: { + "type": "file", + "value": value, + "media_type": media_type, + } + ) + runtime.tool = MagicMock( + side_effect=lambda fn=None, **kwargs: ( + SimpleNamespace(execute=fn, **kwargs) if fn is not None else kwargs + ) + ) + runtime.Tool = MagicMock(side_effect=lambda **kwargs: SimpleNamespace(**kwargs)) + runtime.AgentTool = MagicMock( + side_effect=lambda tool, fn: SimpleNamespace(tool=tool, fn=fn, name=tool.name) + ) + runtime.InferenceRequestParams = MagicMock(return_value=object()) + runtime._follow_handoff = False + created: list[Any] = [] + + def _agent(**kwargs: Any) -> MagicMock: + agent = MagicMock() + agent.kwargs = kwargs + index = len(created) + + async def _select_handoff() -> None: + if index != 0 or not runtime._follow_handoff: + return + handoff = next( + tool + for tool in kwargs.get("tools", []) + if tool.name == "transfer_to_leaf" + ) + await handoff.execute() + + agent.run = MagicMock( + return_value=RunStream( + "root answer" if index == 0 else "leaf answer", + input_tokens=10 if index == 0 else 4, + output_tokens=3 if index == 0 else 2, + on_enter=_select_handoff, + ) + ) + created.append(agent) + return agent + + runtime.Agent = MagicMock(side_effect=_agent) + runtime._created = created + with patch.object(native_graph_mod, "ai", runtime): + yield runtime + + +class TestTopology: + @pytest.mark.asyncio + async def test_constructs_one_native_agent_per_node( + self, ai_runtime: MagicMock + ) -> None: + runner = to_vercel_agents(_definition(_graph())) + await runner.invoke("hello") + assert ai_runtime.Agent.call_count == 2 + assert ai_runtime.get_model.call_args_list[0].args == ("gateway/root",) + assert ai_runtime.get_model.call_args_list[1].args == ("gateway/leaf",) + + @pytest.mark.asyncio + async def test_each_agent_gets_its_own_instructions_and_settings( + self, ai_runtime: MagicMock + ) -> None: + ai_runtime._follow_handoff = True + await to_vercel_agents(_definition(_graph())).invoke("hello") + root_messages = ai_runtime._created[0].run.call_args.kwargs["messages"] + leaf_messages = ai_runtime._created[1].run.call_args.kwargs["messages"] + assert "instructions for root" in str(root_messages) + assert "instructions for leaf" in str(leaf_messages) + assert all( + call.kwargs["params"] is ai_runtime.InferenceRequestParams.return_value + for agent in ai_runtime._created + for call in agent.run.call_args_list + ) + assert ai_runtime.TemperatureSamplerParams.call_count == 2 + assert all( + call.kwargs == {"temperature": 0.1} + for call in ai_runtime.TemperatureSamplerParams.call_args_list + ) + + @pytest.mark.asyncio + async def test_non_terminal_node_has_transfer_tool_per_edge( + self, ai_runtime: MagicMock + ) -> None: + await to_vercel_agents(_definition(_graph())).invoke("hello") + root_tools = ai_runtime.Agent.call_args_list[0].kwargs["tools"] + assert [ + tool.name for tool in root_tools if tool.name.startswith("transfer_to_") + ] == ["transfer_to_leaf"] + + @pytest.mark.asyncio + async def test_terminal_node_has_no_handoff_tools( + self, ai_runtime: MagicMock + ) -> None: + await to_vercel_agents(_definition(_graph())).invoke("hello") + leaf_tools = ai_runtime.Agent.call_args_list[1].kwargs.get("tools", []) + assert not [tool for tool in leaf_tools if tool.name.startswith("transfer_to_")] + + @pytest.mark.asyncio + async def test_node_tools_are_included_with_global_handlers( + self, ai_runtime: MagicMock + ) -> None: + lookup = MagicMock(return_value="found") + await to_vercel_agents( + _definition(_graph(include_node_tool=True)), + {"tool_handlers": {"lookup": lookup}}, + ).invoke("hello") + root_tools = ai_runtime.Agent.call_args_list[0].kwargs["tools"] + node_tool = next(tool for tool in root_tools if tool.name == "lookup") + assert await node_tool.execute(key="x") == "found" + lookup.assert_called_once_with({"key": "x"}) + + @pytest.mark.asyncio + async def test_disabled_graph_raises(self, ai_runtime: MagicMock) -> None: + with pytest.raises(ValueError, match="disabled"): + await to_vercel_agents(_definition(_graph(enabled=False))).invoke("hello") + + +class TestInvocation: + @pytest.mark.asyncio + async def test_stops_at_root_when_no_handoff_tool_is_called( + self, ai_runtime: MagicMock + ) -> None: + result = await to_vercel_agents(_definition(_graph())).invoke("hello") + assert result["response"] == "root answer" + ai_runtime._created[0].run.assert_called_once() + ai_runtime._created[1].run.assert_not_called() + + @pytest.mark.asyncio + async def test_starts_at_root_and_applies_history_only_there( + self, ai_runtime: MagicMock + ) -> None: + ai_runtime._follow_handoff = True + history = [ + {"role": "user", "content": "past"}, + {"role": "assistant", "content": "reply"}, + ] + await to_vercel_agents(_definition(_graph())).invoke("latest", {}, history) + root_messages = ai_runtime._created[0].run.call_args.kwargs["messages"] + leaf_messages = ai_runtime._created[1].run.call_args.kwargs.get("messages", []) + assert "past" in str(root_messages) + assert "latest" in str(root_messages) + assert "past" not in str(leaf_messages) + + @pytest.mark.asyncio + async def test_multimodal_history_is_mapped_for_root( + self, ai_runtime: MagicMock + ) -> None: + history = [ + { + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": "YWJj", + }, + } + ], + } + ] + await to_vercel_agents(_definition(_graph())).invoke("describe", {}, history) + ai_runtime.file_part.assert_called_once() + root_messages = ai_runtime._created[0].run.call_args.kwargs["messages"] + assert "b'abc'" in str(root_messages) + + @pytest.mark.asyncio + async def test_returns_final_leaf_text_and_accumulated_usage( + self, ai_runtime: MagicMock + ) -> None: + ai_runtime._follow_handoff = True + result = await to_vercel_agents(_definition(_graph())).invoke("hello") + assert result["response"] == "leaf answer" + assert result["usage"] == {"input": 14, "output": 5, "total": 19} + + @pytest.mark.asyncio + async def test_handoff_path_has_no_duplicate_entries( + self, ai_runtime: MagicMock + ) -> None: + ai_runtime._follow_handoff = True + span = MagicMock() + tracer = MagicMock() + tracer.start_span.return_value = span + with patch.object(native_graph_mod.trace, "get_tracer", return_value=tracer): + await to_vercel_agents(_definition(_graph())).invoke("hello") + path_values = [ + call.args[1] + for call in span.set_attribute.call_args_list + if call.args[0] == "ld.ai.graph.path" + ] + assert path_values + path = path_values[-1].split("->") + assert path == ["root", "leaf"] + assert len(path) == len(set(path)) + + +class TestTelemetryAndCleanup: + @pytest.mark.asyncio + async def test_context_emits_graph_success_handoff_duration_and_tokens( + self, ai_runtime: MagicMock + ) -> None: + ai_runtime._follow_handoff = True + client = MagicMock() + context = {"kind": "user", "key": "u"} + with patch.object(native_graph_mod, "get_client", return_value=client): + await to_vercel_agents(_definition(_graph()), {"context": context}).invoke( + "hello" + ) + event_names = [call.args[0] for call in client.track.call_args_list] + assert "$ld:ai:graph:invocation_success" in event_names + assert "$ld:ai:graph:handoff_success" in event_names + assert "$ld:ai:graph:duration:total" in event_names + assert "$ld:ai:graph:total_tokens" in event_names + + @pytest.mark.asyncio + async def test_no_context_emits_no_launchdarkly_tracking( + self, ai_runtime: MagicMock + ) -> None: + client = MagicMock() + with patch.object(native_graph_mod, "get_client", return_value=client): + await to_vercel_agents(_definition(_graph())).invoke("hello") + client.track.assert_not_called() + + @pytest.mark.asyncio + async def test_graph_span_ends_on_success(self, ai_runtime: MagicMock) -> None: + span = MagicMock() + tracer = MagicMock() + tracer.start_span.return_value = span + with patch.object(native_graph_mod.trace, "get_tracer", return_value=tracer): + await to_vercel_agents(_definition(_graph())).invoke("hello") + tracer.start_span.assert_called_once_with("ld.ai.graph") + span.end.assert_called_once() + + @pytest.mark.asyncio + async def test_graph_span_ends_and_failure_is_tracked_on_error( + self, ai_runtime: MagicMock + ) -> None: + class BrokenContext: + async def __aenter__(self) -> Any: + raise RuntimeError("provider failed") + + async def __aexit__(self, *args: Any) -> None: + return None + + client = MagicMock() + span = MagicMock() + tracer = MagicMock() + tracer.start_span.return_value = span + ai_runtime.Agent.side_effect = None + broken_agent = MagicMock() + broken_agent.run.return_value = BrokenContext() + ai_runtime.Agent.return_value = broken_agent + with ( + patch.object(native_graph_mod, "get_client", return_value=client), + patch.object(native_graph_mod.trace, "get_tracer", return_value=tracer), + pytest.raises(RuntimeError, match="provider failed"), + ): + await to_vercel_agents( + _definition(_graph()), {"context": {"kind": "user", "key": "u"}} + ).invoke("hello") + assert "$ld:ai:graph:invocation_failure" in [ + call.args[0] for call in client.track.call_args_list + ] + span.end.assert_called_once() + + @pytest.mark.asyncio + async def test_graph_span_ends_on_cancellation(self, ai_runtime: MagicMock) -> None: + class CancelledContext: + async def __aenter__(self) -> Any: + raise asyncio.CancelledError + + async def __aexit__(self, *args: Any) -> None: + return None + + span = MagicMock() + tracer = MagicMock() + tracer.start_span.return_value = span + ai_runtime.Agent.side_effect = None + cancelled_agent = MagicMock() + cancelled_agent.run.return_value = CancelledContext() + ai_runtime.Agent.return_value = cancelled_agent + with ( + patch.object(native_graph_mod.trace, "get_tracer", return_value=tracer), + pytest.raises(asyncio.CancelledError), + ): + await to_vercel_agents(_definition(_graph())).invoke("hello") + span.end.assert_called_once() diff --git a/packages/vercel-messages/README.md b/packages/vercel-messages/README.md new file mode 100644 index 00000000..47610293 --- /dev/null +++ b/packages/vercel-messages/README.md @@ -0,0 +1,63 @@ +# `launchdarkly-ai-vercel-messages` + +Wildcard LaunchDarkly AI messages handler built on the official Vercel AI SDK +for Python. It builds Vercel's `creator/model` id from the evaluated provider +and model, while preserving already-qualified Gateway ids. + +## Install + +```bash +pip install launchdarkly-ai-server launchdarkly-server-sdk launchdarkly-ai-vercel-messages +``` + +Default routing is AI Gateway: for example, `xAI` + `grok-4.5` becomes +`spacexai/grok-4.5` before calling `ai.get_model`. Set `AI_GATEWAY_API_KEY`, or use Vercel OIDC with the +`ai[vercel]` extra. Provider keys such as `OPENAI_API_KEY` are ignored on that +path. This package does not import `langchain-openai` or other provider SDKs +from `config.provider.name`. To call a model directly, inject `model=` with an +already configured model or a sync/async factory: + +```python +import ai +from launchdarkly_ai_vercel_messages import vercel_messages + +result = await vercel_messages( + key, + user_input, + context, + model=lambda cfg: ai.Model( + id=cfg["model"]["name"], + provider=ai.get_provider("openai"), + ), +) +``` + +All LaunchDarkly providers are handled explicitly. Anthropic, OpenAI, Azure, +Gemini, Cohere, DeepSeek, Meta, Mistral, Perplexity, and Vertex map directly. +Bedrock, Cortex, Cursor, Databricks, and Fireworks AI infer the creator from the +model and fail locally if ambiguous. AI21 Labs and IBM Watson fail locally +because Vercel's current catalog has no corresponding creator. + +`experimental_evaluate` / `vercelEvaluate` is TypeScript-only until the official +Python `ai` package ships the same API. + +## Use + +```python +from launchdarkly_ai_server import config +from launchdarkly_ai_vercel_messages import create_vercel_messages_handler + +result = await config( + key="my-ai-config", + handler=create_vercel_messages_handler(), +).invoke("Hello", {"kind": "user", "key": "user-123"}) +``` + +The factory accepts `model=` and +`capture_content=False`. It advertises `("*", "messages")`; register only one +wildcard messages adapter (Vercel, LangChain, etc.) in a handler pool. + +The adapter maps native text and multimodal history, callable tools, streaming, +structured JSON output, usage, and LaunchDarkly telemetry. Flag parameters +cannot replace handler-owned model, credentials, messages, tools, stream +lifecycle, output contract, or loop controls. diff --git a/packages/vercel-messages/agents.md b/packages/vercel-messages/agents.md new file mode 100644 index 00000000..8b03372f --- /dev/null +++ b/packages/vercel-messages/agents.md @@ -0,0 +1,19 @@ +# Agent Guide — Vercel Messages + +This Tier 1 package provides the wildcard `("*", "messages")` handler. + +- `handler.py` owns model resolution, `InferenceRequestParams`, native message + and file-part conversion, native tool declarations, structured output, + streaming cleanup, usage, and spans. +- `__init__.py` exports the factory/wrapper and registers the package. +- Default models use a mapped Gateway `creator/model` id. Preserve + slash-qualified ids, map `xAI` to `spacexai`, and split recognized LD dotted + creator prefixes once before `ai.get_model`. +- Injected model instances and sync/async factories are the Gateway opt-out + and are scoped to their handler. +- Only callable tool handlers are exposed; JSON Schema is preserved in + `ToolSpec.params`. +- `ai.stream` must remain inside `async with`, including early consumer exit. +- Do not add a provider-native client or enable duplicate runtime telemetry. + +`TESTING.md` §1.x and Appendix A.14 are authoritative. diff --git a/packages/vercel-messages/pyproject.toml b/packages/vercel-messages/pyproject.toml new file mode 100644 index 00000000..e73b4bff --- /dev/null +++ b/packages/vercel-messages/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "launchdarkly-ai-vercel-messages" +version = "0.1.0" +requires-python = ">=3.12" +dependencies = [ + "launchdarkly-ai-server", + "opentelemetry-api>=1.25", + "ai>=0.7", +] +description = "Vercel AI SDK messages handler for LaunchDarkly AI SDK" +readme = "README.md" +license = "Apache-2.0" +authors = [{name = "LaunchDarkly", email = "team@launchdarkly.com"}] +keywords = ["launchdarkly", "ai", "vercel", "messages"] + +[build-system] +requires = ["hatchling>=1.32,<2"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/launchdarkly_ai_vercel_messages"] diff --git a/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/__init__.py b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/__init__.py new file mode 100644 index 00000000..8c0df2ad --- /dev/null +++ b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/__init__.py @@ -0,0 +1,11 @@ +"""LaunchDarkly AI SDK integration for Vercel AI SDK messages.""" + +__version__ = "0.1.0" # x-release-please-version + +from launchdarkly_ai_server import register_ai_sdk_package + +from .handler import create_vercel_messages_handler, vercel_messages + +__all__ = ["create_vercel_messages_handler", "vercel_messages"] + +register_ai_sdk_package("launchdarkly-ai-vercel-messages", __version__) diff --git a/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/handler.py b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/handler.py new file mode 100644 index 00000000..22f3a7ad --- /dev/null +++ b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/handler.py @@ -0,0 +1,518 @@ +from __future__ import annotations + +import base64 +import inspect +import json +from collections.abc import AsyncGenerator, Callable +from contextlib import aclosing +from typing import Any + +import ai +from ai.types.tools import ToolSpec +from opentelemetry import trace +from opentelemetry.trace import StatusCode +from pydantic import BaseModel, ConfigDict, create_model + +from launchdarkly_ai_server import ( + AiConfigRep, + LDContext, + NativeTool, + ProviderHandler, + compose_history, + config, + create_handler, + parse_template, + set_ld_span_attributes, +) + +from .model_id import gateway_model_id + +OWNED_PARAMETERS = { + "model", + "messages", + "prompt", + "system", + "tools", + "stream", + "output", + "outputformat", + "stopwhen", + "maxsteps", + "apikey", + "baseurl", +} + +ModelSource = Any | Callable[[AiConfigRep], Any] + +MAX_STEPS = 10 + + +def _owned(name: str) -> bool: + return name.replace("_", "").lower() in OWNED_PARAMETERS + + +def _request_params(config: AiConfigRep) -> Any: + raw = { + key: value + for key, value in (config.get("model", {}).get("parameters") or {}).items() + if not _owned(key) + } + sampling: dict[type[Any], Any] = {} + mappings = { + "temperature": ("TemperatureSamplerParams", "temperature"), + "top_p": ("TopPSamplerParams", "top_p"), + "topP": ("TopPSamplerParams", "top_p"), + "top_k": ("TopKSamplerParams", "top_k"), + "topK": ("TopKSamplerParams", "top_k"), + "min_p": ("MinPSamplerParams", "min_p"), + "minP": ("MinPSamplerParams", "min_p"), + "repetition_penalty": ( + "RepetitionPenaltyParams", + "repetition_penalty", + ), + "repetitionPenalty": ( + "RepetitionPenaltyParams", + "repetition_penalty", + ), + "seed": ("SeedSamplerParams", "seed"), + } + for key, (class_name, argument) in mappings.items(): + if key in raw: + cls = getattr(ai, class_name) + sampling[cls] = cls(**{argument: raw.pop(key)}) + kwargs: dict[str, Any] = {} + if sampling: + kwargs["sampling"] = sampling + max_tokens = raw.pop("max_tokens", raw.pop("maxTokens", None)) + if max_tokens is not None: + kwargs["output"] = ai.OutputParams(max_tokens=max_tokens) + reasoning_effort = raw.pop("reasoning_effort", raw.pop("reasoningEffort", None)) + if reasoning_effort is not None: + kwargs["reasoning"] = ai.ReasoningParams(effort=reasoning_effort) + for direct in ("metadata", "safety_identifier", "extra_headers", "extra_query"): + if direct in raw: + kwargs[direct] = raw.pop(direct) + if raw: + kwargs["extra_body"] = raw + return ai.InferenceRequestParams(**kwargs) + + +async def _resolve_model(source: ModelSource | None, cfg: AiConfigRep) -> Any: + if source is None: + return ai.get_model(gateway_model_id(cfg)) + value = source(cfg) if callable(source) else source + return await value if inspect.isawaitable(value) else value + + +def _image_part(block: dict[str, Any]) -> Any: + source = block.get("source") or {} + media_type = source.get("media_type") + data: str | bytes + if source.get("type") == "base64": + data = base64.b64decode(source.get("data", "")) + else: + data = source.get("url") or block.get("url") or "" + return ai.file_part(data, media_type=media_type) + + +def _content(content: Any) -> Any: + if not isinstance(content, list): + return content if isinstance(content, str) else "" + parts: list[Any] = [] + for block in content: + if block.get("type") == "text": + parts.append(block.get("text", "")) + elif block.get("type") in {"image", "file"}: + parts.append(_image_part(block)) + return parts + + +def _messages( + cfg: AiConfigRep, + user_input: str | None, + variables: dict[str, Any], + history: list[dict[str, Any]] | None, +) -> list[Any]: + result: list[Any] = [] + config_messages: list[dict[str, Any]] = [] + if cfg.get("instructions"): + result.append(ai.system_message(parse_template(cfg["instructions"], variables))) + else: + for message in cfg.get("messages") or []: + content = message.get("content", "") + mapped = ( + parse_template(content, variables) + if isinstance(content, str) + else content + ) + if message.get("role") == "system": + result.append(ai.system_message(mapped)) + else: + config_messages.append({**message, "content": mapped}) + + turns = compose_history( + history=history or [], + user_input=user_input, + config_messages=config_messages, + ) + for turn in turns: + content = _content(turn.get("content", "")) + args = content if isinstance(content, list) else [content] + if turn.get("role") == "assistant": + result.append(ai.assistant_message(*args)) + elif turn.get("role") == "user": + result.append(ai.user_message(*args)) + return result + + +async def _call_tool(handler: Callable[..., Any], kwargs: dict[str, Any]) -> Any: + value = handler(kwargs) + return await value if inspect.isawaitable(value) else value + + +def _tool_executors( + cfg: AiConfigRep, handlers: dict[str, Any] | None +) -> dict[str, Callable[..., Any]]: + executors: dict[str, Callable[..., Any]] = {} + for name in cfg.get("tools") or {}: + handler = (handlers or {}).get(name) + if callable(handler) and not isinstance(handler, NativeTool): + executors[name] = handler + return executors + + +def _tool_args(call: Any) -> dict[str, Any]: + args = getattr(call, "tool_args", None) + if isinstance(args, str): + try: + args = json.loads(args) + except json.JSONDecodeError: + return {} + return args if isinstance(args, dict) else {} + + +async def _tool_result(call: Any, executors: dict[str, Callable[..., Any]]) -> Any: + handler = executors.get(call.tool_name) + if handler is None: + return ai.tool_result_part( + call.tool_call_id, + tool_name=call.tool_name, + result=f"No handler registered for tool {call.tool_name!r}", + is_error=True, + ) + try: + result = await _call_tool(handler, _tool_args(call)) + except Exception as exc: # surfaced to the model so it can recover or explain + return ai.tool_result_part( + call.tool_call_id, + tool_name=call.tool_name, + result=str(exc), + is_error=True, + ) + return ai.tool_result_part( + call.tool_call_id, tool_name=call.tool_name, result=result + ) + + +def _final_output(stream: Any) -> str: + output = stream.output + if isinstance(output, BaseModel): + return output.model_dump_json() + if isinstance(output, str): + return output + return json.dumps(output) + + +def _tools(cfg: AiConfigRep, handlers: dict[str, Any] | None) -> list[Any]: + result: list[Any] = [] + for name, definition in (cfg.get("tools") or {}).items(): + handler = (handlers or {}).get(name) + if not callable(handler) or isinstance(handler, NativeTool): + continue + + async def execute(_handler: Callable[..., Any] = handler, **kwargs: Any) -> Any: + return await _call_tool(_handler, kwargs) + + spec = ToolSpec( + description=definition.get("description"), + params=definition.get("parameters") or {}, + ) + tool = ai.Tool(kind="function", name=name, spec=spec) + # The model-facing object remains a native Tool. The attribute is useful to + # custom executors and harmless to the frozen SDK type only when supported. + try: + object.__setattr__(tool, "name", name) + object.__setattr__(tool, "execute", execute) + object.__setattr__(tool, "input_schema", spec.params) + object.__setattr__(tool, "description", spec.description) + except (AttributeError, TypeError): + pass + result.append(tool) + return result + + +def _python_type(schema: dict[str, Any]) -> Any: + kind = schema.get("type") + if kind == "string": + return str + if kind == "integer": + return int + if kind == "number": + return float + if kind == "boolean": + return bool + if kind == "array": + return list[Any] + return Any + + +def _output_type(schema: dict[str, Any]) -> type[BaseModel]: + required = set(schema.get("required") or []) + fields: dict[str, Any] = {} + for name, field_schema in (schema.get("properties") or {}).items(): + typ = _python_type(field_schema) + fields[name] = (typ, ... if name in required else None) + return create_model( + "VercelStructuredOutput", + __config__=ConfigDict(extra="allow"), + **fields, + ) + + +def _usage(stream: Any) -> dict[str, int]: + usage = getattr(stream, "usage", None) + input_tokens = int(getattr(usage, "input_tokens", getattr(usage, "input", 0)) or 0) + output_tokens = int( + getattr(usage, "output_tokens", getattr(usage, "output", 0)) or 0 + ) + return {"input_tokens": input_tokens, "output_tokens": output_tokens} + + +def _text_delta(event: Any) -> str | None: + if getattr(event, "kind", None) == "text_delta": + return str(getattr(event, "chunk", "")) + events = getattr(ai, "events", None) + delta_type = getattr(events, "TextDelta", None) + if delta_type is not None and isinstance(event, delta_type): + return str(event.chunk) + return None + + +def _start_span(cfg: AiConfigRep, variables: dict[str, Any]) -> Any: + span = trace.get_tracer("@launchdarkly/ai-vercel-messages").start_span( + "invoke_agent" + ) + provider = str(cfg.get("provider", {}).get("name") or "vercel").lower() + model = str(cfg.get("model", {}).get("name") or "") + span.set_attribute("gen_ai.operation.name", "invoke_agent") + span.set_attribute("gen_ai.system", provider) + span.set_attribute("gen_ai.provider.name", provider) + span.set_attribute("gen_ai.request.model", model) + set_ld_span_attributes(span, variables) + return span + + +def _start_model_span(cfg: AiConfigRep, root: Any) -> Any: + model = str(cfg.get("model", {}).get("name") or "") + provider = str(cfg.get("provider", {}).get("name") or "vercel").lower() + span = trace.get_tracer("@launchdarkly/ai-vercel-messages").start_span( + f"chat {model}", context=trace.set_span_in_context(root) + ) + span.set_attribute("gen_ai.operation.name", "chat") + span.set_attribute("gen_ai.system", provider) + span.set_attribute("gen_ai.provider.name", provider) + span.set_attribute("gen_ai.request.model", model) + return span + + +def _set_usage(span: Any, usage: dict[str, int]) -> None: + input_tokens = usage["input_tokens"] + output_tokens = usage["output_tokens"] + span.set_attribute("gen_ai.usage.input_tokens", input_tokens) + span.set_attribute("gen_ai.usage.output_tokens", output_tokens) + span.set_attribute("gen_ai.usage.total_tokens", input_tokens + output_tokens) + span.set_attribute("gen_ai.usage.prompt_tokens", input_tokens) + span.set_attribute("gen_ai.usage.completion_tokens", output_tokens) + span.set_attribute("gen_ai.usage.cache_read.input_tokens", 0) + span.set_attribute("gen_ai.usage.cache_creation.input_tokens", 0) + + +async def _run_conversation( + cfg: AiConfigRep, + model: ModelSource | None, + user_input: str | None, + variables: dict[str, Any], + history: list[dict[str, Any]] | None, + tool_handlers: dict[str, Any] | None, + *, + structured: bool, +) -> AsyncGenerator[dict[str, Any], None]: + """Drive a request to its final answer, resolving tool calls between turns. + + ``ai.stream`` reports the tool calls a model asks for but never runs them, so a + config with tools stalls on an assistant message that still has calls pending + unless each round is executed and fed back. + """ + resolved_model = await _resolve_model(model, cfg) + messages = _messages(cfg, user_input, variables, history) + tools = _tools(cfg, tool_handlers) + executors = _tool_executors(cfg, tool_handlers) + params = _request_params(cfg) + output_type = ( + _output_type(cfg["outputFormat"]) + if structured and cfg.get("outputFormat") + else None + ) + + input_tokens = 0 + output_tokens = 0 + for _ in range(MAX_STEPS): + kwargs: dict[str, Any] = { + "model": resolved_model, + "messages": messages, + "tools": tools, + "params": params, + } + if output_type is not None: + kwargs["output_type"] = output_type + async with ai.stream(**kwargs) as provider_stream: + async for event in provider_stream: + text = _text_delta(event) + if text is not None: + yield {"type": "chunk", "text": text} + usage = _usage(provider_stream) + input_tokens += usage["input_tokens"] + output_tokens += usage["output_tokens"] + message = provider_stream.message + calls = list(message.tool_calls or []) + if not calls: + yield { + "type": "done", + "output": _final_output(provider_stream), + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + }, + } + return + results = [await _tool_result(call, executors) for call in calls] + messages = [*messages, message, ai.tool_message(*results)] + raise RuntimeError( + f"Vercel messages run did not reach a final response within {MAX_STEPS} steps" + ) + + +def create_vercel_messages_handler( + model: ModelSource | None = None, + *, + capture_content: bool = False, +) -> ProviderHandler: + async def invoke( + cfg: AiConfigRep, + user_input: str | None = None, + tool_handlers: dict[str, Any] | None = None, + variables: dict[str, Any] | None = None, + history: list[dict[str, Any]] | None = None, + ) -> dict[str, Any]: + vs = variables or {} + span = _start_span(cfg, vs) + model_span = _start_model_span(cfg, span) + try: + text = "" + usage = {"input_tokens": 0, "output_tokens": 0} + async with aclosing( + _run_conversation( + cfg, model, user_input, vs, history, tool_handlers, structured=True + ) + ) as events: + async for event in events: + if event["type"] == "done": + text = event["output"] + usage = event["usage"] + _set_usage(model_span, usage) + _set_usage(span, usage) + model_span.set_status(StatusCode.OK) + span.set_status(StatusCode.OK) + return {"output": text, "usage": usage} + except BaseException as exc: + if isinstance(exc, Exception): + model_span.record_exception(exc) + model_span.set_status(StatusCode.ERROR, str(exc)) + span.record_exception(exc) + span.set_status(StatusCode.ERROR, str(exc)) + raise + finally: + model_span.end() + span.end() + + async def stream( + cfg: AiConfigRep, + user_input: str | None = None, + tool_handlers: dict[str, Any] | None = None, + variables: dict[str, Any] | None = None, + history: list[dict[str, Any]] | None = None, + ) -> AsyncGenerator[dict[str, Any], None]: + vs = variables or {} + span = _start_span(cfg, vs) + model_span = _start_model_span(cfg, span) + completed = False + try: + output = "" + usage = {"input_tokens": 0, "output_tokens": 0} + # aclosing so that abandoning this generator also unwinds the provider + # stream's context rather than leaving it open until finalization. + async with aclosing( + _run_conversation( + cfg, model, user_input, vs, history, tool_handlers, structured=False + ) + ) as events: + async for event in events: + if event["type"] == "chunk": + yield event + continue + completed = True + output = event["output"] + usage = event["usage"] + _set_usage(model_span, usage) + _set_usage(span, usage) + model_span.set_status(StatusCode.OK) + span.set_status(StatusCode.OK) + yield {"type": "done", "output": output, "usage": usage} + except BaseException as exc: + if isinstance(exc, Exception): + model_span.record_exception(exc) + model_span.set_status(StatusCode.ERROR, str(exc)) + span.record_exception(exc) + span.set_status(StatusCode.ERROR, str(exc)) + raise + finally: + if not completed: + model_span.set_attribute("launchdarkly.stream.abandoned", True) + span.set_attribute("launchdarkly.stream.abandoned", True) + model_span.end() + span.end() + + return create_handler( + ("*", "messages"), + invoke, + stream, + capture_content=capture_content, + ) + + +def vercel_messages( + config_key: str, + user_input: str | None, + context: LDContext, + *, + model: ModelSource | None = None, + capture_content: bool = False, + variables: dict[str, Any] | None = None, + **options: Any, +) -> Any: + handler = create_vercel_messages_handler( + model=model, capture_content=capture_content + ) + return config(key=config_key, handler=handler, **options).invoke( + user_input, context, variables=variables + ) diff --git a/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/model_id.py b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/model_id.py new file mode 100644 index 00000000..8a23b660 --- /dev/null +++ b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/model_id.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +import re + +from launchdarkly_ai_server import AiConfigRep + +GATEWAY_CREATORS = { + "alibaba": "alibaba", + "amazon": "amazon", + "anthropic": "anthropic", + "arceeai": "arcee-ai", + "bfl": "bfl", + "bytedance": "bytedance", + "cohere": "cohere", + "deepseek": "deepseek", + "fishaudio": "fish-audio", + "gemini": "google", + "google": "google", + "googleai": "google", + "inception": "inception", + "inclusionai": "inclusionai", + "inferencenet": "inference-net", + "interfaze": "interfaze", + "klingai": "klingai", + "meta": "meta", + "minimax": "minimax", + "mistral": "mistral", + "mistralai": "mistral", + "mixedbread": "mixedbread", + "moonshotai": "moonshotai", + "morph": "morph", + "nvidia": "nvidia", + "openai": "openai", + "perplexity": "perplexity", + "poolside": "poolside", + "prodia": "prodia", + "quiverai": "quiverai", + "recraft": "recraft", + "sakana": "sakana", + "spacexai": "spacexai", + "stepfun": "stepfun", + "tencent": "tencent", + "thinkingmachines": "thinkingmachines", + "typesafe": "typesafe-ai", + "typesafeai": "typesafe-ai", + "voyage": "voyage", + "xiaomi": "xiaomi", + "xai": "spacexai", + "zai": "zai", +} + +PROVIDER_CREATORS: dict[str, str | None] = { + "anthropic": "anthropic", + "openai": "openai", + "bedrock": None, + "azure": "openai", + "gemini": "google", + "ai21labs": None, + "cohere": "cohere", + "cortex": None, + "cursor": None, + "databricks": None, + "deepseek": "deepseek", + "fireworksai": None, + "ibmwatson": None, + "meta": "meta", + "mistral": "mistral", + "perplexity": "perplexity", + "vertex": "google", + # Compatibility aliases found in existing configs. + "google": "google", + "googleai": "google", + "mistralai": "mistral", + "spacexai": "spacexai", + "typesafe": "typesafe-ai", + "typesafeai": "typesafe-ai", + "xai": "spacexai", +} + +MODEL_FAMILY_CREATORS = ( + (re.compile(r"^(?:gpt|o[1-9])(?:[-.]|$)", re.I), "openai"), + (re.compile(r"^claude(?:[-.]|$)", re.I), "anthropic"), + (re.compile(r"^gemini(?:[-.]|$)", re.I), "google"), + (re.compile(r"^grok(?:[-.]|$)", re.I), "spacexai"), + (re.compile(r"^(?:command|aya)(?:[-.]|$)", re.I), "cohere"), + (re.compile(r"^deepseek(?:[-.]|$)", re.I), "deepseek"), + (re.compile(r"^llama(?:[-.]|$)", re.I), "meta"), + ( + re.compile(r"^(?:mistral|mixtral|codestral|pixtral)(?:[-.]|$)", re.I), + "mistral", + ), + (re.compile(r"^sonar(?:[-.]|$)", re.I), "perplexity"), + (re.compile(r"^(?:nova|titan)(?:[-.]|$)", re.I), "amazon"), + (re.compile(r"^qwen(?:[-.]|$)", re.I), "alibaba"), +) + + +def _provider_key(value: str) -> str: + return re.sub(r"[^a-z0-9]", "", value.lower()) + + +def _creator_from_model(model: str) -> tuple[str, str] | None: + parts = model.split(".") + for index, part in enumerate(parts[:-1]): + creator = GATEWAY_CREATORS.get(_provider_key(part)) + if creator: + return creator, ".".join(parts[index + 1 :]) + for pattern, creator in MODEL_FAMILY_CREATORS: + if pattern.search(model): + return creator, model + return None + + +def gateway_model_id(config: AiConfigRep) -> str: + """Build the AI Gateway ``creator/model`` id for an evaluated config.""" + model = str(config["model"]["name"]) + if "/" in model: + return model + + provider = str((config.get("provider") or {}).get("name") or "") + provider_name = _provider_key(provider) + inferred = _creator_from_model(model) + if inferred: + return f"{inferred[0]}/{inferred[1]}" + + creator = PROVIDER_CREATORS.get(provider_name) + if creator: + return f"{creator}/{model}" + if provider_name in {"ai21labs", "ibmwatson"}: + raise ValueError( + f'Vercel AI Gateway currently exposes no models created by "{provider}". ' + "Inject a direct provider model with model/model_factory instead." + ) + if provider_name in PROVIDER_CREATORS: + raise ValueError( + f'LaunchDarkly provider "{provider}" hosts models from multiple ' + f'creators, so "{model}" cannot be converted to a Vercel creator/model ' + "id. Store an explicit creator/model id or inject a model/model_factory." + ) + else: + raise ValueError( + f'Cannot map LaunchDarkly provider "{provider or "unknown"}" to a ' + "Vercel AI Gateway creator. Pass a creator/model id in " + "config.model.name or inject a model/model factory." + ) diff --git a/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/py.typed b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/py.typed new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/py.typed @@ -0,0 +1 @@ + diff --git a/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/spans.py b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/spans.py new file mode 100644 index 00000000..648dc88d --- /dev/null +++ b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/spans.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from typing import Any + +from opentelemetry import trace +from opentelemetry.trace import StatusCode + +from launchdarkly_ai_server import ( + AiConfigRep, + SpanUsage, + set_ld_span_attributes, + set_model_identity_attributes, + set_usage_span_attributes, +) + +TRACER_NAME = "@launchdarkly/ai-vercel-messages" + + +def model_name(config: AiConfigRep) -> str: + return str(config.get("model", {}).get("name") or "") + + +def serving_provider(config: AiConfigRep) -> str: + return str(config.get("provider", {}).get("name") or "vercel").lower() + + +def start_root_span(config: AiConfigRep, variables: dict[str, Any]) -> Any: + span = trace.get_tracer(TRACER_NAME).start_span("invoke_agent") + span.set_attribute("gen_ai.operation.name", "invoke_agent") + set_model_identity_attributes(span, serving_provider(config), model_name(config)) + set_ld_span_attributes(span, variables) + return span + + +def parent_context_of(span: Any) -> Any: + return trace.set_span_in_context(span) if span is not None else None + + +def start_model_span(config: AiConfigRep, parent: Any) -> Any: + span = trace.get_tracer(TRACER_NAME).start_span( + f"chat {model_name(config)}", context=parent + ) + span.set_attribute("gen_ai.operation.name", "chat") + set_model_identity_attributes(span, serving_provider(config), model_name(config)) + return span + + +def start_tool_span(name: str, call_id: str, parent: Any) -> Any: + span = trace.get_tracer(TRACER_NAME).start_span( + f"execute_tool {name}", context=parent + ) + span.set_attribute("gen_ai.operation.name", "execute_tool") + span.set_attribute("gen_ai.tool.name", name) + span.set_attribute("gen_ai.tool.call.id", call_id) + return span + + +def finish_root_span(span: Any, response_model: str, usage: SpanUsage) -> None: + span.set_attribute("gen_ai.response.model", response_model) + set_usage_span_attributes(span, usage) + + +def succeed_span(span: Any) -> None: + span.set_status(StatusCode.OK) + span.end() + + +def mark_ok(span: Any) -> None: + span.set_status(StatusCode.OK) + + +def fail_span(span: Any, error: BaseException, tracker: set[int] | None = None) -> None: + span.record_exception(error) + span.set_status(StatusCode.ERROR, str(error)) + if tracker is None or id(span) not in tracker: + if tracker is not None: + tracker.add(id(span)) + span.end() diff --git a/packages/vercel-messages/tests/test_handler.py b/packages/vercel-messages/tests/test_handler.py new file mode 100644 index 00000000..577906b8 --- /dev/null +++ b/packages/vercel-messages/tests/test_handler.py @@ -0,0 +1,503 @@ +"""Test-first contract for the wildcard Vercel AI SDK messages adapter. + +These tests intentionally precede ``launchdarkly_ai_vercel_messages`` production +modules. They use only in-memory ``ai`` doubles and must never open a network. +""" + +from __future__ import annotations + +import json +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import launchdarkly_ai_vercel_messages.handler as handler_mod +from launchdarkly_ai_vercel_messages import create_vercel_messages_handler + +CONFIG: dict[str, Any] = { + "model": {"name": "anthropic/claude-sonnet-4"}, + "provider": {"name": "Anthropic"}, + "instructions": "Help {{name}}.", +} + + +class TextDelta: + def __init__(self, chunk: str) -> None: + self.chunk = chunk + + +class FakeStream: + """Mirrors ``ai.Stream``, including the ``message`` whose pending tool calls + decide whether the adapter has to run another turn.""" + + def __init__( + self, + events: list[Any] | None = None, + *, + text: str = "answer", + input_tokens: int = 10, + output_tokens: int = 4, + output: Any = None, + tool_calls: list[Any] | None = None, + ) -> None: + self.events = list(events or [TextDelta(text)]) + self.text = text + self.output = output if output is not None else text + self.usage = SimpleNamespace( + input_tokens=input_tokens, + output_tokens=output_tokens, + ) + self.message = SimpleNamespace( + role="assistant", + text=text, + tool_calls=list(tool_calls or []), + usage=self.usage, + ) + self.entered = 0 + self.exited = 0 + self.consumed = 0 + + async def __aenter__(self) -> FakeStream: + self.entered += 1 + return self + + async def __aexit__(self, *args: Any) -> None: + self.exited += 1 + + def __aiter__(self) -> FakeStream: + return self + + async def __anext__(self) -> Any: + if not self.events: + raise StopAsyncIteration + self.consumed += 1 + return self.events.pop(0) + + +@pytest.fixture +def ai_runtime() -> MagicMock: + """Patch every Vercel runtime entry point with an in-memory implementation.""" + runtime = MagicMock() + runtime.events.TextDelta = TextDelta + runtime.get_model = MagicMock(return_value=object()) + runtime.system_message = MagicMock( + side_effect=lambda content: {"role": "system", "content": content} + ) + runtime.user_message = MagicMock( + side_effect=lambda *content: { + "role": "user", + "content": content[0] if len(content) == 1 else list(content), + } + ) + runtime.assistant_message = MagicMock( + side_effect=lambda *content: { + "role": "assistant", + "content": content[0] if len(content) == 1 else list(content), + } + ) + runtime.file_part = MagicMock( + side_effect=lambda value, media_type=None: { + "type": "file", + "value": value, + "media_type": media_type, + } + ) + runtime.tool = MagicMock( + side_effect=lambda fn=None, **kwargs: ( + SimpleNamespace(execute=fn, **kwargs) if fn is not None else kwargs + ) + ) + runtime.InferenceRequestParams = MagicMock(return_value=object()) + runtime.stream = MagicMock(return_value=FakeStream()) + with patch.object(handler_mod, "ai", runtime): + yield runtime + + +def _stream_kwargs(runtime: MagicMock) -> dict[str, Any]: + return runtime.stream.call_args.kwargs + + +class TestFactoryAndModelResolution: + def test_advertises_wildcard_messages(self) -> None: + assert create_vercel_messages_handler().provides_for == ("*", "messages") + + @pytest.mark.asyncio + async def test_gateway_model_id_is_passed_unchanged_to_get_model( + self, ai_runtime: MagicMock + ) -> None: + await create_vercel_messages_handler()(CONFIG, "hello", {}, {"name": "Ada"}) + ai_runtime.get_model.assert_called_once_with("anthropic/claude-sonnet-4") + + @pytest.mark.asyncio + async def test_builds_gateway_creator_model_id(self, ai_runtime: MagicMock) -> None: + config = { + **CONFIG, + "model": {"name": "grok-4.5"}, + "provider": {"name": "xAI"}, + } + await create_vercel_messages_handler()(config, "hello") + ai_runtime.get_model.assert_called_once_with("spacexai/grok-4.5") + + @pytest.mark.asyncio + async def test_converts_dotted_creator_prefix(self, ai_runtime: MagicMock) -> None: + config = { + **CONFIG, + "model": {"name": "openai.gpt-5.6-sol"}, + "provider": {"name": "OpenAI"}, + } + await create_vercel_messages_handler()(config, "hello") + ai_runtime.get_model.assert_called_once_with("openai/gpt-5.6-sol") + + @pytest.mark.asyncio + async def test_does_not_construct_a_provider_native_client( + self, ai_runtime: MagicMock + ) -> None: + forbidden = { + "openai": MagicMock(), + "anthropic": MagicMock(), + "boto3": MagicMock(), + "google.generativeai": MagicMock(), + } + with patch.dict("sys.modules", forbidden): + await create_vercel_messages_handler()(CONFIG, "hello") + for module in forbidden.values(): + assert not module.mock_calls + + @pytest.mark.asyncio + async def test_injected_model_instance_bypasses_get_model( + self, ai_runtime: MagicMock + ) -> None: + model = object() + await create_vercel_messages_handler(model=model)(CONFIG, "hello") + ai_runtime.get_model.assert_not_called() + assert _stream_kwargs(ai_runtime)["model"] is model + + @pytest.mark.asyncio + async def test_sync_factory_receives_evaluated_config_once( + self, ai_runtime: MagicMock + ) -> None: + model = object() + factory = MagicMock(return_value=model) + await create_vercel_messages_handler(model=factory)(CONFIG, "hello") + factory.assert_called_once_with(CONFIG) + assert _stream_kwargs(ai_runtime)["model"] is model + + @pytest.mark.asyncio + async def test_async_factory_is_supported(self, ai_runtime: MagicMock) -> None: + model = object() + factory = AsyncMock(return_value=model) + await create_vercel_messages_handler(model=factory)(CONFIG, "hello") + factory.assert_awaited_once_with(CONFIG) + assert _stream_kwargs(ai_runtime)["model"] is model + + @pytest.mark.asyncio + async def test_handler_instances_keep_model_sources_isolated( + self, ai_runtime: MagicMock + ) -> None: + first, second = object(), object() + await create_vercel_messages_handler(model=first)(CONFIG, "one") + await create_vercel_messages_handler(model=second)(CONFIG, "two") + assert [call.kwargs["model"] for call in ai_runtime.stream.call_args_list] == [ + first, + second, + ] + + +class TestOwnedParameters: + @pytest.mark.asyncio + async def test_forwards_generation_settings_and_strips_owned_fields( + self, ai_runtime: MagicMock + ) -> None: + owned = { + "model": "evil", + "messages": [], + "prompt": "evil", + "system": "evil", + "tools": ["evil"], + "stream": False, + "output": "evil", + "outputFormat": {}, + "output_format": {}, + "stopWhen": "evil", + "stop_when": "evil", + "maxSteps": 999, + "max_steps": 999, + "apiKey": "secret", + "api_key": "secret", + "baseURL": "https://evil.invalid", + "base_url": "https://evil.invalid", + } + config = { + **CONFIG, + "model": { + "name": CONFIG["model"]["name"], + "parameters": {"temperature": 0.2, "top_p": 0.8, **owned}, + }, + } + await create_vercel_messages_handler()(config, "hello") + kwargs = _stream_kwargs(ai_runtime) + assert kwargs["params"] is ai_runtime.InferenceRequestParams.return_value + ai_runtime.TemperatureSamplerParams.assert_called_once_with(temperature=0.2) + ai_runtime.TopPSamplerParams.assert_called_once_with(top_p=0.8) + constructor_args = repr(ai_runtime.InferenceRequestParams.call_args) + for key in owned: + assert key not in constructor_args + + +class TestMessagesAndHistory: + @pytest.mark.asyncio + async def test_uses_ai_stream_and_native_message_constructors( + self, ai_runtime: MagicMock + ) -> None: + await create_vercel_messages_handler()(CONFIG, "hello", {}, {"name": "Ada"}) + ai_runtime.stream.assert_called_once() + ai_runtime.system_message.assert_called_once_with("Help Ada.") + ai_runtime.user_message.assert_called_with("hello") + + @pytest.mark.asyncio + async def test_config_history_and_input_are_ordered( + self, ai_runtime: MagicMock + ) -> None: + config = { + **CONFIG, + "instructions": None, + "messages": [{"role": "user", "content": "config turn"}], + } + history = [{"role": "assistant", "content": "history turn"}] + await create_vercel_messages_handler()(config, "latest", {}, {}, history) + messages = _stream_kwargs(ai_runtime)["messages"] + assert [message["content"] for message in messages] == [ + "config turn", + "history turn", + "latest", + ] + + @pytest.mark.asyncio + async def test_multimodal_history_uses_file_part( + self, ai_runtime: MagicMock + ) -> None: + history = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": "YWJj", + }, + }, + ], + } + ] + await create_vercel_messages_handler()(CONFIG, None, {}, {}, history) + ai_runtime.file_part.assert_called() + args = ai_runtime.file_part.call_args + assert "image/png" in str(args) + assert "b'abc'" in str(args) + + +class TestToolsAndOutput: + @pytest.mark.asyncio + async def test_only_tools_with_callable_handlers_are_offered( + self, ai_runtime: MagicMock + ) -> None: + config = { + **CONFIG, + "tools": { + "weather": { + "description": "Get weather", + "parameters": {"type": "object", "properties": {"city": {}}}, + }, + "missing": {"description": "No implementation", "parameters": {}}, + }, + } + weather = AsyncMock(return_value="sunny") + await create_vercel_messages_handler()(config, "hello", {"weather": weather}) + tools = _stream_kwargs(ai_runtime)["tools"] + assert [tool.name for tool in tools] == ["weather"] + assert tools[0].description == "Get weather" + assert tools[0].input_schema == config["tools"]["weather"]["parameters"] + + @pytest.mark.asyncio + async def test_runtime_tool_executes_registered_callable( + self, ai_runtime: MagicMock + ) -> None: + config = { + **CONFIG, + "tools": {"weather": {"description": "Get weather", "parameters": {}}}, + } + weather = AsyncMock(return_value="sunny") + await create_vercel_messages_handler()(config, "hello", {"weather": weather}) + tool = _stream_kwargs(ai_runtime)["tools"][0] + assert await tool.execute(city="Oakland") == "sunny" + weather.assert_awaited_once_with({"city": "Oakland"}) + + @pytest.mark.asyncio + async def test_executes_tool_calls_and_continues_until_a_final_answer( + self, ai_runtime: MagicMock + ) -> None: + config = { + **CONFIG, + "tools": {"weather": {"description": "Get weather", "parameters": {}}}, + } + call = SimpleNamespace( + tool_call_id="tc-1", tool_name="weather", tool_args={"city": "Oakland"} + ) + ai_runtime.stream.side_effect = [ + FakeStream(text="", tool_calls=[call], input_tokens=10, output_tokens=4), + FakeStream(text="sunny in Oakland", input_tokens=7, output_tokens=2), + ] + weather = AsyncMock(return_value="sunny") + + result = await create_vercel_messages_handler()( + config, "hello", {"weather": weather} + ) + + weather.assert_awaited_once_with({"city": "Oakland"}) + assert result["output"] == "sunny in Oakland" + assert result["usage"] == {"input_tokens": 17, "output_tokens": 6} + assert ai_runtime.stream.call_count == 2 + ai_runtime.tool_result_part.assert_called_once_with( + "tc-1", tool_name="weather", result="sunny" + ) + + @pytest.mark.asyncio + async def test_reports_a_failing_tool_back_to_the_model( + self, ai_runtime: MagicMock + ) -> None: + config = { + **CONFIG, + "tools": {"weather": {"description": "Get weather", "parameters": {}}}, + } + call = SimpleNamespace(tool_call_id="tc-1", tool_name="weather", tool_args={}) + ai_runtime.stream.side_effect = [ + FakeStream(text="", tool_calls=[call]), + FakeStream(text="could not check"), + ] + weather = AsyncMock(side_effect=RuntimeError("upstream down")) + + result = await create_vercel_messages_handler()( + config, "hello", {"weather": weather} + ) + + assert result["output"] == "could not check" + ai_runtime.tool_result_part.assert_called_once_with( + "tc-1", tool_name="weather", result="upstream down", is_error=True + ) + + @pytest.mark.asyncio + async def test_stops_a_runaway_tool_loop(self, ai_runtime: MagicMock) -> None: + config = { + **CONFIG, + "tools": {"weather": {"description": "Get weather", "parameters": {}}}, + } + call = SimpleNamespace(tool_call_id="tc-1", tool_name="weather", tool_args={}) + ai_runtime.stream.side_effect = lambda **_: FakeStream( + text="", tool_calls=[call] + ) + + with pytest.raises(RuntimeError, match="did not reach a final response"): + await create_vercel_messages_handler()( + config, "hello", {"weather": AsyncMock(return_value="sunny")} + ) + + @pytest.mark.asyncio + async def test_structured_output_uses_output_type_and_serializes_result( + self, ai_runtime: MagicMock + ) -> None: + schema = { + "type": "object", + "properties": {"answer": {"type": "string"}}, + "required": ["answer"], + } + ai_runtime.stream.return_value = FakeStream(text="", output={"answer": "yes"}) + result = await create_vercel_messages_handler()( + {**CONFIG, "outputFormat": schema}, "question" + ) + assert _stream_kwargs(ai_runtime)["output_type"] is not None + assert json.loads(result["output"]) == {"answer": "yes"} + + @pytest.mark.asyncio + async def test_streaming_ignores_output_format(self, ai_runtime: MagicMock) -> None: + handler = create_vercel_messages_handler() + events = [ + event + async for event in await handler.stream( + {**CONFIG, "outputFormat": {"type": "object"}}, "question" + ) + ] + assert "output_type" not in _stream_kwargs(ai_runtime) + assert events[-1]["type"] == "done" + + +class TestUsageAndStreaming: + @pytest.mark.asyncio + async def test_normalizes_usage_after_blocking_stream_completion( + self, ai_runtime: MagicMock + ) -> None: + ai_runtime.stream.return_value = FakeStream(input_tokens=12, output_tokens=7) + result = await create_vercel_messages_handler()(CONFIG, "hello") + assert result["usage"] == {"input_tokens": 12, "output_tokens": 7} + + @pytest.mark.asyncio + async def test_stream_forwards_deltas_and_emits_one_done( + self, ai_runtime: MagicMock + ) -> None: + ai_runtime.stream.return_value = FakeStream( + [TextDelta("hel"), TextDelta("lo")], text="hello" + ) + events = [ + event + async for event in await create_vercel_messages_handler().stream( + CONFIG, "hello" + ) + ] + assert [event["text"] for event in events if event["type"] == "chunk"] == [ + "hel", + "lo", + ] + assert sum(event["type"] == "done" for event in events) == 1 + assert events[-1]["output"] == "hello" + + @pytest.mark.asyncio + async def test_early_exit_closes_async_stream_context( + self, ai_runtime: MagicMock + ) -> None: + provider_stream = FakeStream([TextDelta("one"), TextDelta("two")]) + ai_runtime.stream.return_value = provider_stream + stream = await create_vercel_messages_handler().stream(CONFIG, "hello") + async for _event in stream: + break + await stream.aclose() + assert provider_stream.exited == 1 + assert provider_stream.consumed == 1 + + @pytest.mark.asyncio + async def test_runtime_telemetry_is_not_enabled( + self, ai_runtime: MagicMock + ) -> None: + await create_vercel_messages_handler()(CONFIG, "hello") + kwargs = _stream_kwargs(ai_runtime) + assert not kwargs.get("experimental_telemetry") + assert not kwargs.get("telemetry") + + @pytest.mark.asyncio + async def test_span_identity_uses_provider_and_preserves_gateway_model( + self, ai_runtime: MagicMock + ) -> None: + span = MagicMock() + tracer = MagicMock() + tracer.start_span.return_value = span + with patch.object(handler_mod.trace, "get_tracer", return_value=tracer): + await create_vercel_messages_handler()(CONFIG, "hello") + attributes = { + call.args[0]: call.args[1] for call in span.set_attribute.call_args_list + } + assert attributes["gen_ai.system"] == "anthropic" + assert attributes["gen_ai.provider.name"] == "anthropic" + assert attributes["gen_ai.request.model"] == "anthropic/claude-sonnet-4" diff --git a/packages/vercel-messages/tests/test_model_id.py b/packages/vercel-messages/tests/test_model_id.py new file mode 100644 index 00000000..6f226600 --- /dev/null +++ b/packages/vercel-messages/tests/test_model_id.py @@ -0,0 +1,55 @@ +import pytest + +from launchdarkly_ai_vercel_messages.model_id import gateway_model_id + + +def config(provider: str, model: str = "model") -> dict: + return {"provider": {"name": provider}, "model": {"name": model}} + + +@pytest.mark.parametrize( + ("provider", "model", "expected"), + [ + ("Anthropic", "model", "anthropic/model"), + ("OpenAI", "model", "openai/model"), + ("Bedrock", "anthropic.claude-sonnet-4", "anthropic/claude-sonnet-4"), + ("Azure", "model", "openai/model"), + ("Gemini", "model", "google/model"), + ("Cohere", "model", "cohere/model"), + ("Cortex", "llama-4-scout", "meta/llama-4-scout"), + ("Cursor", "claude-sonnet-4", "anthropic/claude-sonnet-4"), + ("Databricks", "llama-4-maverick", "meta/llama-4-maverick"), + ("DeepSeek", "model", "deepseek/model"), + ("Fireworks AI", "qwen-3-235b", "alibaba/qwen-3-235b"), + ("Meta", "model", "meta/model"), + ("Mistral", "model", "mistral/model"), + ("Perplexity", "model", "perplexity/model"), + ("Vertex", "model", "google/model"), + ], +) +def test_maps_every_supported_launchdarkly_provider( + provider: str, model: str, expected: str +) -> None: + assert gateway_model_id(config(provider, model)) == expected + + +@pytest.mark.parametrize("provider", ["AI21 Labs", "IBM Watson"]) +def test_fails_locally_for_unsupported_creator(provider: str) -> None: + with pytest.raises(ValueError, match="currently exposes no models created by"): + gateway_model_id(config(provider)) + + +def test_preserves_explicit_gateway_id() -> None: + assert gateway_model_id(config("Bedrock", "amazon/nova-pro")) == "amazon/nova-pro" + + +def test_removes_bedrock_region_and_dotted_creator_prefix() -> None: + assert ( + gateway_model_id(config("Bedrock", "us.anthropic.claude-sonnet-4")) + == "anthropic/claude-sonnet-4" + ) + + +def test_fails_locally_for_unidentifiable_multi_model_host() -> None: + with pytest.raises(ValueError, match="hosts models from multiple creators"): + gateway_model_id(config("Databricks", "custom-deployment")) diff --git a/pyproject.toml b/pyproject.toml index 83e3de63..5e0417b7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dev = [ "ruff>=0.15.20", "pre-commit>=4.6.0", "brotlicffi>=1.0", + "ai[openai]>=0.7", ] [tool.pytest.ini_options] @@ -65,4 +66,6 @@ known-first-party = [ "launchdarkly_ai_langchain_messages", "launchdarkly_ai_openai_agents", "launchdarkly_ai_openai_messages", + "launchdarkly_ai_vercel_agents", + "launchdarkly_ai_vercel_messages", ] diff --git a/release-please-config.json b/release-please-config.json index 9d2a04e3..953ea37e 100644 --- a/release-please-config.json +++ b/release-please-config.json @@ -89,6 +89,28 @@ "src/launchdarkly_ai_langchain_messages/__init__.py" ], "component": "launchdarkly-ai-langchain-messages" + }, + "packages/vercel-agents": { + "release-type": "python", + "versioning": "default", + "bump-minor-pre-major": true, + "bump-patch-for-minor-pre-major": true, + "include-v-in-tag": false, + "extra-files": [ + "src/launchdarkly_ai_vercel_agents/__init__.py" + ], + "component": "launchdarkly-ai-vercel-agents" + }, + "packages/vercel-messages": { + "release-type": "python", + "versioning": "default", + "bump-minor-pre-major": true, + "bump-patch-for-minor-pre-major": true, + "include-v-in-tag": false, + "extra-files": [ + "src/launchdarkly_ai_vercel_messages/__init__.py" + ], + "component": "launchdarkly-ai-vercel-messages" } } } diff --git a/tests/test_cross_handler_parity.py b/tests/test_cross_handler_parity.py index 3ab7513a..7cd4e0d9 100644 --- a/tests/test_cross_handler_parity.py +++ b/tests/test_cross_handler_parity.py @@ -41,6 +41,8 @@ "openai-agents": "launchdarkly_ai_openai_agents.spans", "langchain-messages": "launchdarkly_ai_langchain_messages.spans", "langchain-agents": "launchdarkly_ai_langchain_agents.spans", + "vercel-messages": "launchdarkly_ai_vercel_messages.spans", + "vercel-agents": "launchdarkly_ai_vercel_agents.spans", } #: `claude-agents` builds its `chat` span inside an inference tracker rather than in a standalone diff --git a/uv.lock b/uv.lock index c1c96777..d43a1673 100644 --- a/uv.lock +++ b/uv.lock @@ -21,10 +21,13 @@ members = [ "launchdarkly-ai-openai-messages", "launchdarkly-ai-python", "launchdarkly-ai-server", + "launchdarkly-ai-vercel-agents", + "launchdarkly-ai-vercel-messages", ] [manifest.dependency-groups] dev = [ + { name = "ai", extras = ["openai"], specifier = ">=0.7" }, { name = "brotlicffi", specifier = ">=1.0" }, { name = "claude-agent-sdk", specifier = ">=0.2.110" }, { name = "langchain-anthropic", specifier = ">=1.4.8" }, @@ -41,6 +44,26 @@ dev = [ { name = "ruff", specifier = ">=0.15.20" }, ] +[[package]] +name = "ai" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "httpx2" }, + { name = "modelsdotdev" }, + { name = "pydantic" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3c/09/efc7051a582507474e5f8c9333b1f6bcfd0f43736a014428c3d6f65e69d3/ai-0.7.0.tar.gz", hash = "sha256:1b0b4b20430ab12de38e388b893b2774681f6a3b1ec79850f64db324b5ee55d1", size = 1226720, upload-time = "2026-09-18T16:19:27.569Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ed/42/7823fa36389f04c4d728d9ed58cdcf27fec4d93e54a722839015b1eeadf5/ai-0.7.0-py3-none-any.whl", hash = "sha256:9a94279254814e562cb65618aedd5c5f832e59816f5be25375913c737083365a", size = 210537, upload-time = "2026-09-18T16:19:25.878Z" }, +] + +[package.optional-dependencies] +openai = [ + { name = "openai" }, +] + [[package]] name = "annotated-types" version = "0.7.0" @@ -477,6 +500,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" }, ] +[[package]] +name = "httpcore2" +version = "2.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "h11" }, + { name = "truststore" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/15/8c/e925b1c92018abb3a1863ce1549d76d2381e334d21d65d4ac8f65dabd78a/httpcore2-2.13.0.tar.gz", hash = "sha256:2adc8be4fb285fbcd6d894298db3b52c177e74b6674eda3a76bd36be3292a3db", size = 67740, upload-time = "2026-09-14T14:18:04.717Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/0d/117a771a2bb91df334b66bf4da14cd02f21aefbcfe53180f336ce55e8f90/httpcore2-2.13.0-py3-none-any.whl", hash = "sha256:35ae5be347aa40467b4a5dc032ac67ebb6d27189fc97e8cebcf99616f6a1bb9e", size = 83162, upload-time = "2026-09-14T14:18:02.529Z" }, +] + [[package]] name = "httpx" version = "0.28.1" @@ -501,6 +537,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d2/fd/6668e5aec43ab844de6fc74927e155a3b37bf40d7c3790e49fc0406b6578/httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc", size = 8960, upload-time = "2025-10-10T21:48:21.158Z" }, ] +[[package]] +name = "httpx2" +version = "2.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio", marker = "sys_platform != 'emscripten'" }, + { name = "httpcore2", marker = "sys_platform != 'emscripten'" }, + { name = "httpx2-jsfetch", marker = "sys_platform == 'emscripten'" }, + { name = "idna" }, + { name = "truststore", marker = "sys_platform != 'emscripten'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b9/a0/e9deef4654132857b5a5dbe4eddd0ac59c2814500e11f2f5044cd81103ee/httpx2-2.13.0.tar.gz", hash = "sha256:81bd07dc67a3701729ef1f777a3c00c915d4539604fdb5afd327f8682f6b7b44", size = 100290, upload-time = "2026-09-14T14:18:05.486Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/d1/a0c72b0e006df654709fbc366cc5bcb53e5aee13e1e3395152c6dd293376/httpx2-2.13.0-py3-none-any.whl", hash = "sha256:fc12720cedf72faa26cca6b4ca394e05c894e7d7933fc45cafe767960804e49a", size = 95565, upload-time = "2026-09-14T14:18:03.553Z" }, +] + +[[package]] +name = "httpx2-jsfetch" +version = "1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cd/c4/0e5636363151a2a1795e0a77617168b9ca438e1748ec05fc9b5687f93d64/httpx2_jsfetch-1.0.tar.gz", hash = "sha256:70a0e3eabfef7cce5ad9c629f7d01ca05e418f586646f4ddf14782e4c1454c60", size = 6872, upload-time = "2026-08-07T00:13:07.492Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9b/43/832f631d32e4f1211caa2ba368317739fe71f0b8530e4c9d15dc454bac2a/httpx2_jsfetch-1.0-py3-none-any.whl", hash = "sha256:cb916b707601e69a07721aabc8f3f6659be3a6893bc1ff5c6f9e02241df2da32", size = 6382, upload-time = "2026-08-07T00:13:06.567Z" }, +] + [[package]] name = "identify" version = "2.6.19" @@ -938,6 +1000,40 @@ requires-dist = [ ] provides-extras = ["otel"] +[[package]] +name = "launchdarkly-ai-vercel-agents" +version = "0.1.0" +source = { editable = "packages/vercel-agents" } +dependencies = [ + { name = "ai" }, + { name = "launchdarkly-ai-server" }, + { name = "opentelemetry-api" }, +] + +[package.metadata] +requires-dist = [ + { name = "ai", specifier = ">=0.7" }, + { name = "launchdarkly-ai-server", editable = "packages/client" }, + { name = "opentelemetry-api", specifier = ">=1.25" }, +] + +[[package]] +name = "launchdarkly-ai-vercel-messages" +version = "0.1.0" +source = { editable = "packages/vercel-messages" } +dependencies = [ + { name = "ai" }, + { name = "launchdarkly-ai-server" }, + { name = "opentelemetry-api" }, +] + +[package.metadata] +requires-dist = [ + { name = "ai", specifier = ">=0.7" }, + { name = "launchdarkly-ai-server", editable = "packages/client" }, + { name = "opentelemetry-api", specifier = ">=1.25" }, +] + [[package]] name = "launchdarkly-eventsource" version = "1.7.1" @@ -1052,6 +1148,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e2/5e/d118fce19f87a2e7d8101c35c8ae0ec289098a4df0ff244cec23e415aca0/mcp-1.28.1-py3-none-any.whl", hash = "sha256:2726bca5e7193f61c5dde8b12500a6de2d9acf6d1a1c0be9e8c2e706437991df", size = 222620, upload-time = "2026-06-26T12:57:27.218Z" }, ] +[[package]] +name = "modelsdotdev" +version = "0.20260818.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/17/34/e5d0b54bf967bee3e37cb7d8261143670a79603f6e770fb8a68d8ffcd17a/modelsdotdev-0.20260818.0.tar.gz", hash = "sha256:375e6793dc5bc731075e7631ee32b7d0a3be6a4a6cade3e3d417664dae40d8fd", size = 1087820, upload-time = "2026-08-18T08:14:47.009Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6f/92/242f3c5a20d64e964895a02fe5c86ae66ed09e6317be600c67012bc1addc/modelsdotdev-0.20260818.0-py3-none-any.whl", hash = "sha256:dd55b13bfd029e2a8345f7326e4239dcb8b972d2bbd174f422aa0ca59da1ad9c", size = 1096608, upload-time = "2026-08-18T08:14:45.326Z" }, +] + [[package]] name = "mypy" version = "2.1.0" @@ -2049,6 +2154,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d8/8e/bb97bb0c71802080bfc8952937d174e49cfc50de5c951dd47b2496f0dcdb/tqdm-4.68.3-py3-none-any.whl", hash = "sha256:39832cc2def2789a6f29df83f172db7416cea70052c0907a57801c5f2fdccb03", size = 78337, upload-time = "2026-06-17T07:36:50.132Z" }, ] +[[package]] +name = "truststore" +version = "0.10.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/53/a3/1585216310e344e8102c22482f6060c7a6ea0322b63e026372e6dcefcfd6/truststore-0.10.4.tar.gz", hash = "sha256:9d91bd436463ad5e4ee4aba766628dd6cd7010cf3e2461756b3303710eebc301", size = 26169, upload-time = "2025-08-12T18:49:02.73Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/19/97/56608b2249fe206a67cd573bc93cd9896e1efb9e98bce9c163bcdc704b88/truststore-0.10.4-py3-none-any.whl", hash = "sha256:adaeaecf1cbb5f4de3b1959b42d41f6fab57b2b1666adb59e89cb0b53361d981", size = 18660, upload-time = "2025-08-12T18:49:01.46Z" }, +] + [[package]] name = "typing-extensions" version = "4.15.0" From e83fa1163dd69edcf7b123104e845e188913d480 Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Tue, 22 Sep 2026 12:08:38 -0800 Subject: [PATCH 2/5] fix(AIC-3408): tighten Vercel telemetry and structured output Use the shared span helpers in production, trace tool execution, and emit provider-compatible strict schemas for messages output. --- .sdk_metadata.json | 3 +- .../launchdarkly_ai_vercel_agents/handler.py | 148 +++++++------ .../launchdarkly_ai_vercel_agents/spans.py | 5 + packages/vercel-agents/tests/test_handler.py | 9 +- .../handler.py | 205 ++++++++++-------- .../launchdarkly_ai_vercel_messages/spans.py | 5 + .../vercel-messages/tests/test_handler.py | 52 ++++- 7 files changed, 254 insertions(+), 173 deletions(-) diff --git a/.sdk_metadata.json b/.sdk_metadata.json index 6d63f548..4ac14dad 100644 --- a/.sdk_metadata.json +++ b/.sdk_metadata.json @@ -111,8 +111,7 @@ "aiSdkNames": ["launchdarkly-ai-vercel-messages"], "releases": { "tag-prefix": "launchdarkly-ai-vercel-messages-" }, "features": { - "aiCompletionConfig": { "introduced": "0.1" }, - "aiTrackMetrics": { "introduced": "0.1" } + "aiCompletionConfig": { "introduced": "0.1" } } }, "python-ai-vercel-agents": { diff --git a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/handler.py b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/handler.py index fe43b862..c946874e 100644 --- a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/handler.py +++ b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/handler.py @@ -8,8 +8,6 @@ import ai from ai.types.tools import ToolSpec -from opentelemetry import trace -from opentelemetry.trace import StatusCode from pydantic import BaseModel, ConfigDict, create_model from launchdarkly_ai_server import ( @@ -17,14 +15,25 @@ LDContext, NativeTool, ProviderHandler, + SpanUsage, compose_history, config, create_handler, parse_template, - set_ld_span_attributes, ) from .model_id import gateway_model_id +from .spans import ( + fail_span, + finish_model_span, + finish_root_span, + mark_ok, + model_name, + parent_context_of, + start_model_span, + start_root_span, + start_tool_span, +) OWNED_PARAMETERS = { "model", @@ -184,6 +193,8 @@ def build_agent_tools( definitions: dict[str, Any] | None, handlers: dict[str, Any] | None, runtime: Any = None, + *, + parent: Any = None, ) -> list[Any]: runtime = runtime or ai result: list[Any] = [] @@ -192,8 +203,28 @@ def build_agent_tools( if not callable(handler) or isinstance(handler, NativeTool): continue - async def execute(_handler: Callable[..., Any] = handler, **kwargs: Any) -> Any: - return await _call_tool(_handler, kwargs) + async def execute( + _handler: Callable[..., Any] = handler, + _name: str = name, + **kwargs: Any, + ) -> Any: + tool_span = ( + start_tool_span(_name, "", parent) if parent is not None else None + ) + try: + value = await _call_tool(_handler, kwargs) + except BaseException as exc: + if tool_span is not None: + if isinstance(exc, Exception): + fail_span(tool_span, exc) + else: + tool_span.set_attribute("launchdarkly.run.cancelled", True) + tool_span.end() + raise + if tool_span is not None: + mark_ok(tool_span) + tool_span.end() + return value spec = ToolSpec( description=definition.get("description"), @@ -304,43 +335,6 @@ def text_delta(event: Any) -> str | None: return None -def _start_span(cfg: AiConfigRep, variables: dict[str, Any]) -> Any: - span = trace.get_tracer("@launchdarkly/ai-vercel-agents").start_span("invoke_agent") - provider = str(cfg.get("provider", {}).get("name") or "vercel").lower() - model = str(cfg.get("model", {}).get("name") or "") - span.set_attribute("gen_ai.operation.name", "invoke_agent") - span.set_attribute("gen_ai.system", provider) - span.set_attribute("gen_ai.provider.name", provider) - span.set_attribute("gen_ai.request.model", model) - set_ld_span_attributes(span, variables) - return span - - -def _start_model_span(cfg: AiConfigRep, root: Any) -> Any: - model = str(cfg.get("model", {}).get("name") or "") - provider = str(cfg.get("provider", {}).get("name") or "vercel").lower() - span = trace.get_tracer("@launchdarkly/ai-vercel-agents").start_span( - f"chat {model}", context=trace.set_span_in_context(root) - ) - span.set_attribute("gen_ai.operation.name", "chat") - span.set_attribute("gen_ai.system", provider) - span.set_attribute("gen_ai.provider.name", provider) - span.set_attribute("gen_ai.request.model", model) - return span - - -def _set_usage(span: Any, usage: dict[str, int]) -> None: - input_tokens = usage["input_tokens"] - output_tokens = usage["output_tokens"] - span.set_attribute("gen_ai.usage.input_tokens", input_tokens) - span.set_attribute("gen_ai.usage.output_tokens", output_tokens) - span.set_attribute("gen_ai.usage.total_tokens", input_tokens + output_tokens) - span.set_attribute("gen_ai.usage.prompt_tokens", input_tokens) - span.set_attribute("gen_ai.usage.completion_tokens", output_tokens) - span.set_attribute("gen_ai.usage.cache_read.input_tokens", 0) - span.set_attribute("gen_ai.usage.cache_creation.input_tokens", 0) - - def create_vercel_agents_handler( model: ModelSource | None = None, *, @@ -354,10 +348,14 @@ async def run( history: list[dict[str, Any]] | None = None, ) -> dict[str, Any]: vs = variables or {} - span = _start_span(cfg, vs) - model_span = _start_model_span(cfg, span) + span = start_root_span(cfg, vs) + parent = parent_context_of(span) + model_span = start_model_span(cfg, parent) + failed = False try: - agent = ai.Agent(tools=build_agent_tools(cfg.get("tools"), tool_handlers)) + agent = ai.Agent( + tools=build_agent_tools(cfg.get("tools"), tool_handlers, parent=parent) + ) run_kwargs: dict[str, Any] = { "model": await resolve_model(model, cfg), "messages": build_messages(cfg, user_input, vs, history), @@ -370,21 +368,25 @@ async def run( pass usage = usage_of(provider_stream) output = output_of(provider_stream) - _set_usage(model_span, usage) - _set_usage(span, usage) - model_span.set_status(StatusCode.OK) - span.set_status(StatusCode.OK) + span_usage = SpanUsage( + input=usage["input_tokens"], output=usage["output_tokens"] + ) + response_model = model_name(cfg) + finish_model_span(model_span, response_model, span_usage) + finish_root_span(span, response_model, span_usage) + mark_ok(model_span) + mark_ok(span) return {"output": output, "usage": usage} except BaseException as exc: if isinstance(exc, Exception): - model_span.record_exception(exc) - model_span.set_status(StatusCode.ERROR, str(exc)) - span.record_exception(exc) - span.set_status(StatusCode.ERROR, str(exc)) + failed = True + fail_span(model_span, exc) + fail_span(span, exc) raise finally: - model_span.end() - span.end() + if not failed: + model_span.end() + span.end() async def stream( cfg: AiConfigRep, @@ -394,12 +396,16 @@ async def stream( history: list[dict[str, Any]] | None = None, ) -> AsyncGenerator[dict[str, Any], None]: vs = variables or {} - span = _start_span(cfg, vs) - model_span = _start_model_span(cfg, span) + span = start_root_span(cfg, vs) + parent = parent_context_of(span) + model_span = start_model_span(cfg, parent) completed = False + failed = False provider_stream: Any = None try: - agent = ai.Agent(tools=build_agent_tools(cfg.get("tools"), tool_handlers)) + agent = ai.Agent( + tools=build_agent_tools(cfg.get("tools"), tool_handlers, parent=parent) + ) async with agent.run( model=await resolve_model(model, cfg), messages=build_messages(cfg, user_input, vs, history), @@ -411,10 +417,14 @@ async def stream( yield {"type": "chunk", "text": text} completed = True usage = usage_of(provider_stream) - _set_usage(model_span, usage) - _set_usage(span, usage) - model_span.set_status(StatusCode.OK) - span.set_status(StatusCode.OK) + span_usage = SpanUsage( + input=usage["input_tokens"], output=usage["output_tokens"] + ) + response_model = model_name(cfg) + finish_model_span(model_span, response_model, span_usage) + finish_root_span(span, response_model, span_usage) + mark_ok(model_span) + mark_ok(span) yield { "type": "done", "output": output_of(provider_stream), @@ -422,17 +432,17 @@ async def stream( } except BaseException as exc: if isinstance(exc, Exception): - model_span.record_exception(exc) - model_span.set_status(StatusCode.ERROR, str(exc)) - span.record_exception(exc) - span.set_status(StatusCode.ERROR, str(exc)) + failed = True + fail_span(model_span, exc) + fail_span(span, exc) raise finally: - if not completed: + if not completed and not failed: model_span.set_attribute("launchdarkly.stream.abandoned", True) span.set_attribute("launchdarkly.stream.abandoned", True) - model_span.end() - span.end() + if not failed: + model_span.end() + span.end() return create_handler(("*", "agent"), run, stream, capture_content=capture_content) diff --git a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/spans.py b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/spans.py index d782f2b2..0762842f 100644 --- a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/spans.py +++ b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/spans.py @@ -60,6 +60,11 @@ def finish_root_span(span: Any, response_model: str, usage: SpanUsage) -> None: set_usage_span_attributes(span, usage) +def finish_model_span(span: Any, response_model: str, usage: SpanUsage) -> None: + span.set_attribute("gen_ai.response.model", response_model) + set_usage_span_attributes(span, usage) + + def succeed_span(span: Any) -> None: span.set_status(StatusCode.OK) span.end() diff --git a/packages/vercel-agents/tests/test_handler.py b/packages/vercel-agents/tests/test_handler.py index 12aa97db..5453964e 100644 --- a/packages/vercel-agents/tests/test_handler.py +++ b/packages/vercel-agents/tests/test_handler.py @@ -235,7 +235,14 @@ async def test_only_callable_tools_reach_agent(self, ai_runtime: MagicMock) -> N await create_vercel_agents_handler()(config, "hello", {"lookup": lookup}) tools = ai_runtime.Agent.call_args.kwargs["tools"] assert [tool.name for tool in tools] == ["lookup"] - assert await tools[0].execute(key="x") == "found" + tool_span = MagicMock() + with patch.object( + handler_mod, "start_tool_span", return_value=tool_span + ) as start_tool_span: + assert await tools[0].execute(key="x") == "found" + start_tool_span.assert_called_once() + assert start_tool_span.call_args.args[0] == "lookup" + tool_span.end.assert_called_once() lookup.assert_awaited_once_with({"key": "x"}) @pytest.mark.asyncio diff --git a/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/handler.py b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/handler.py index 22f3a7ad..dcb3cfa2 100644 --- a/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/handler.py +++ b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/handler.py @@ -9,8 +9,6 @@ import ai from ai.types.tools import ToolSpec -from opentelemetry import trace -from opentelemetry.trace import StatusCode from pydantic import BaseModel, ConfigDict, create_model from launchdarkly_ai_server import ( @@ -18,14 +16,25 @@ LDContext, NativeTool, ProviderHandler, + SpanUsage, compose_history, config, create_handler, parse_template, - set_ld_span_attributes, ) from .model_id import gateway_model_id +from .spans import ( + fail_span, + finish_model_span, + finish_root_span, + mark_ok, + model_name, + parent_context_of, + start_model_span, + start_root_span, + start_tool_span, +) OWNED_PARAMETERS = { "model", @@ -191,7 +200,11 @@ def _tool_args(call: Any) -> dict[str, Any]: return args if isinstance(args, dict) else {} -async def _tool_result(call: Any, executors: dict[str, Callable[..., Any]]) -> Any: +async def _tool_result( + call: Any, + executors: dict[str, Callable[..., Any]], + parent: Any, +) -> Any: handler = executors.get(call.tool_name) if handler is None: return ai.tool_result_part( @@ -200,15 +213,24 @@ async def _tool_result(call: Any, executors: dict[str, Callable[..., Any]]) -> A result=f"No handler registered for tool {call.tool_name!r}", is_error=True, ) + tool_span = start_tool_span(call.tool_name, call.tool_call_id, parent) try: result = await _call_tool(handler, _tool_args(call)) - except Exception as exc: # surfaced to the model so it can recover or explain - return ai.tool_result_part( - call.tool_call_id, - tool_name=call.tool_name, - result=str(exc), - is_error=True, - ) + except BaseException as exc: + if isinstance(exc, Exception): + # Surface ordinary tool failures to the model so it can recover. + fail_span(tool_span, exc) + return ai.tool_result_part( + call.tool_call_id, + tool_name=call.tool_name, + result=str(exc), + is_error=True, + ) + tool_span.set_attribute("launchdarkly.run.cancelled", True) + tool_span.end() + raise + mark_ok(tool_span) + tool_span.end() return ai.tool_result_part( call.tool_call_id, tool_name=call.tool_name, result=result ) @@ -251,10 +273,8 @@ async def execute(_handler: Callable[..., Any] = handler, **kwargs: Any) -> Any: return result -def _python_type(schema: dict[str, Any]) -> Any: +def _python_type(schema: dict[str, Any], name: str) -> Any: kind = schema.get("type") - if kind == "string": - return str if kind == "integer": return int if kind == "number": @@ -262,21 +282,26 @@ def _python_type(schema: dict[str, Any]) -> Any: if kind == "boolean": return bool if kind == "array": - return list[Any] - return Any - - -def _output_type(schema: dict[str, Any]) -> type[BaseModel]: - required = set(schema.get("required") or []) - fields: dict[str, Any] = {} - for name, field_schema in (schema.get("properties") or {}).items(): - typ = _python_type(field_schema) - fields[name] = (typ, ... if name in required else None) - return create_model( - "VercelStructuredOutput", - __config__=ConfigDict(extra="allow"), - **fields, - ) + items = schema.get("items") + item_type = ( + _python_type(items, f"{name}Item") if isinstance(items, dict) else str + ) + return list[item_type] # type: ignore[valid-type] + if kind == "object" or schema.get("properties"): + return _output_type(schema, name=f"{name}Object") + # Strict structured output rejects untyped members. + return str + + +def _output_type( + schema: dict[str, Any], *, name: str = "VercelStructuredOutput" +) -> type[BaseModel]: + """Build the strict Pydantic model required by structured-output providers.""" + fields: dict[str, Any] = { + field: (_python_type(field_schema, field.title().replace("_", "")), ...) + for field, field_schema in (schema.get("properties") or {}).items() + } + return create_model(name, __config__=ConfigDict(extra="forbid"), **fields) def _usage(stream: Any) -> dict[str, int]: @@ -298,45 +323,6 @@ def _text_delta(event: Any) -> str | None: return None -def _start_span(cfg: AiConfigRep, variables: dict[str, Any]) -> Any: - span = trace.get_tracer("@launchdarkly/ai-vercel-messages").start_span( - "invoke_agent" - ) - provider = str(cfg.get("provider", {}).get("name") or "vercel").lower() - model = str(cfg.get("model", {}).get("name") or "") - span.set_attribute("gen_ai.operation.name", "invoke_agent") - span.set_attribute("gen_ai.system", provider) - span.set_attribute("gen_ai.provider.name", provider) - span.set_attribute("gen_ai.request.model", model) - set_ld_span_attributes(span, variables) - return span - - -def _start_model_span(cfg: AiConfigRep, root: Any) -> Any: - model = str(cfg.get("model", {}).get("name") or "") - provider = str(cfg.get("provider", {}).get("name") or "vercel").lower() - span = trace.get_tracer("@launchdarkly/ai-vercel-messages").start_span( - f"chat {model}", context=trace.set_span_in_context(root) - ) - span.set_attribute("gen_ai.operation.name", "chat") - span.set_attribute("gen_ai.system", provider) - span.set_attribute("gen_ai.provider.name", provider) - span.set_attribute("gen_ai.request.model", model) - return span - - -def _set_usage(span: Any, usage: dict[str, int]) -> None: - input_tokens = usage["input_tokens"] - output_tokens = usage["output_tokens"] - span.set_attribute("gen_ai.usage.input_tokens", input_tokens) - span.set_attribute("gen_ai.usage.output_tokens", output_tokens) - span.set_attribute("gen_ai.usage.total_tokens", input_tokens + output_tokens) - span.set_attribute("gen_ai.usage.prompt_tokens", input_tokens) - span.set_attribute("gen_ai.usage.completion_tokens", output_tokens) - span.set_attribute("gen_ai.usage.cache_read.input_tokens", 0) - span.set_attribute("gen_ai.usage.cache_creation.input_tokens", 0) - - async def _run_conversation( cfg: AiConfigRep, model: ModelSource | None, @@ -344,6 +330,7 @@ async def _run_conversation( variables: dict[str, Any], history: list[dict[str, Any]] | None, tool_handlers: dict[str, Any] | None, + parent: Any, *, structured: bool, ) -> AsyncGenerator[dict[str, Any], None]: @@ -395,7 +382,7 @@ async def _run_conversation( }, } return - results = [await _tool_result(call, executors) for call in calls] + results = [await _tool_result(call, executors, parent) for call in calls] messages = [*messages, message, ai.tool_message(*results)] raise RuntimeError( f"Vercel messages run did not reach a final response within {MAX_STEPS} steps" @@ -415,35 +402,48 @@ async def invoke( history: list[dict[str, Any]] | None = None, ) -> dict[str, Any]: vs = variables or {} - span = _start_span(cfg, vs) - model_span = _start_model_span(cfg, span) + span = start_root_span(cfg, vs) + parent = parent_context_of(span) + model_span = start_model_span(cfg, parent) + failed = False try: text = "" usage = {"input_tokens": 0, "output_tokens": 0} async with aclosing( _run_conversation( - cfg, model, user_input, vs, history, tool_handlers, structured=True + cfg, + model, + user_input, + vs, + history, + tool_handlers, + parent, + structured=True, ) ) as events: async for event in events: if event["type"] == "done": text = event["output"] usage = event["usage"] - _set_usage(model_span, usage) - _set_usage(span, usage) - model_span.set_status(StatusCode.OK) - span.set_status(StatusCode.OK) + span_usage = SpanUsage( + input=usage["input_tokens"], output=usage["output_tokens"] + ) + response_model = model_name(cfg) + finish_model_span(model_span, response_model, span_usage) + finish_root_span(span, response_model, span_usage) + mark_ok(model_span) + mark_ok(span) return {"output": text, "usage": usage} except BaseException as exc: if isinstance(exc, Exception): - model_span.record_exception(exc) - model_span.set_status(StatusCode.ERROR, str(exc)) - span.record_exception(exc) - span.set_status(StatusCode.ERROR, str(exc)) + failed = True + fail_span(model_span, exc) + fail_span(span, exc) raise finally: - model_span.end() - span.end() + if not failed: + model_span.end() + span.end() async def stream( cfg: AiConfigRep, @@ -453,9 +453,11 @@ async def stream( history: list[dict[str, Any]] | None = None, ) -> AsyncGenerator[dict[str, Any], None]: vs = variables or {} - span = _start_span(cfg, vs) - model_span = _start_model_span(cfg, span) + span = start_root_span(cfg, vs) + parent = parent_context_of(span) + model_span = start_model_span(cfg, parent) completed = False + failed = False try: output = "" usage = {"input_tokens": 0, "output_tokens": 0} @@ -463,7 +465,14 @@ async def stream( # stream's context rather than leaving it open until finalization. async with aclosing( _run_conversation( - cfg, model, user_input, vs, history, tool_handlers, structured=False + cfg, + model, + user_input, + vs, + history, + tool_handlers, + parent, + structured=False, ) ) as events: async for event in events: @@ -473,24 +482,28 @@ async def stream( completed = True output = event["output"] usage = event["usage"] - _set_usage(model_span, usage) - _set_usage(span, usage) - model_span.set_status(StatusCode.OK) - span.set_status(StatusCode.OK) + span_usage = SpanUsage( + input=usage["input_tokens"], output=usage["output_tokens"] + ) + response_model = model_name(cfg) + finish_model_span(model_span, response_model, span_usage) + finish_root_span(span, response_model, span_usage) + mark_ok(model_span) + mark_ok(span) yield {"type": "done", "output": output, "usage": usage} except BaseException as exc: if isinstance(exc, Exception): - model_span.record_exception(exc) - model_span.set_status(StatusCode.ERROR, str(exc)) - span.record_exception(exc) - span.set_status(StatusCode.ERROR, str(exc)) + failed = True + fail_span(model_span, exc) + fail_span(span, exc) raise finally: - if not completed: + if not completed and not failed: model_span.set_attribute("launchdarkly.stream.abandoned", True) span.set_attribute("launchdarkly.stream.abandoned", True) - model_span.end() - span.end() + if not failed: + model_span.end() + span.end() return create_handler( ("*", "messages"), diff --git a/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/spans.py b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/spans.py index 648dc88d..a4a7cef4 100644 --- a/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/spans.py +++ b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/spans.py @@ -60,6 +60,11 @@ def finish_root_span(span: Any, response_model: str, usage: SpanUsage) -> None: set_usage_span_attributes(span, usage) +def finish_model_span(span: Any, response_model: str, usage: SpanUsage) -> None: + span.set_attribute("gen_ai.response.model", response_model) + set_usage_span_attributes(span, usage) + + def succeed_span(span: Any) -> None: span.set_status(StatusCode.OK) span.end() diff --git a/packages/vercel-messages/tests/test_handler.py b/packages/vercel-messages/tests/test_handler.py index 577906b8..6666db4a 100644 --- a/packages/vercel-messages/tests/test_handler.py +++ b/packages/vercel-messages/tests/test_handler.py @@ -14,6 +14,7 @@ import pytest import launchdarkly_ai_vercel_messages.handler as handler_mod +import launchdarkly_ai_vercel_messages.spans as spans_mod from launchdarkly_ai_vercel_messages import create_vercel_messages_handler CONFIG: dict[str, Any] = { @@ -353,12 +354,19 @@ async def test_executes_tool_calls_and_continues_until_a_final_answer( FakeStream(text="sunny in Oakland", input_tokens=7, output_tokens=2), ] weather = AsyncMock(return_value="sunny") + tool_span = MagicMock() - result = await create_vercel_messages_handler()( - config, "hello", {"weather": weather} - ) + with patch.object( + handler_mod, "start_tool_span", return_value=tool_span + ) as start_tool_span: + result = await create_vercel_messages_handler()( + config, "hello", {"weather": weather} + ) weather.assert_awaited_once_with({"city": "Oakland"}) + start_tool_span.assert_called_once() + assert start_tool_span.call_args.args[:2] == ("weather", "tc-1") + tool_span.end.assert_called_once() assert result["output"] == "sunny in Oakland" assert result["usage"] == {"input_tokens": 17, "output_tokens": 6} assert ai_runtime.stream.call_count == 2 @@ -419,9 +427,43 @@ async def test_structured_output_uses_output_type_and_serializes_result( result = await create_vercel_messages_handler()( {**CONFIG, "outputFormat": schema}, "question" ) - assert _stream_kwargs(ai_runtime)["output_type"] is not None + output_type = _stream_kwargs(ai_runtime)["output_type"] + output_schema = output_type.model_json_schema() + assert output_schema["required"] == ["answer"] + assert output_schema["additionalProperties"] is False assert json.loads(result["output"]) == {"answer": "yes"} + @pytest.mark.asyncio + async def test_structured_output_types_nested_objects_and_array_items( + self, ai_runtime: MagicMock + ) -> None: + schema = { + "type": "object", + "properties": { + "sources": {"type": "array", "items": {"type": "string"}}, + "details": { + "type": "object", + "properties": {"summary": {"type": "string"}}, + }, + }, + } + ai_runtime.stream.return_value = FakeStream( + text="", + output={"sources": ["docs"], "details": {"summary": "answer"}}, + ) + + await create_vercel_messages_handler()( + {**CONFIG, "outputFormat": schema}, "question" + ) + + output_schema = _stream_kwargs(ai_runtime)["output_type"].model_json_schema() + assert output_schema["required"] == ["sources", "details"] + assert output_schema["additionalProperties"] is False + assert output_schema["properties"]["sources"]["items"]["type"] == "string" + nested = output_schema["$defs"]["DetailsObject"] + assert nested["required"] == ["summary"] + assert nested["additionalProperties"] is False + @pytest.mark.asyncio async def test_streaming_ignores_output_format(self, ai_runtime: MagicMock) -> None: handler = create_vercel_messages_handler() @@ -493,7 +535,7 @@ async def test_span_identity_uses_provider_and_preserves_gateway_model( span = MagicMock() tracer = MagicMock() tracer.start_span.return_value = span - with patch.object(handler_mod.trace, "get_tracer", return_value=tracer): + with patch.object(spans_mod.trace, "get_tracer", return_value=tracer): await create_vercel_messages_handler()(CONFIG, "hello") attributes = { call.args[0]: call.args[1] for call in span.set_attribute.call_args_list From c4ef4b88b12f6cf19cb797e08bab8516da384c3a Mon Sep 17 00:00:00 2001 From: Alexis Georges Date: Fri, 25 Sep 2026 12:38:37 -0400 Subject: [PATCH 3/5] refactor(vercel-agents): rename the graph span to launchdarkly.graph --- .../src/launchdarkly_ai_vercel_agents/native_graph.py | 6 +++--- packages/vercel-agents/tests/test_native_graph.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/native_graph.py b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/native_graph.py index b7f47313..b567b17c 100644 --- a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/native_graph.py +++ b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/native_graph.py @@ -127,9 +127,9 @@ async def build(node: GraphNode) -> None: await build(definition.root) span = trace.get_tracer("@launchdarkly/ai-vercel-agents").start_span( - "ld.ai.graph" + "launchdarkly.graph" ) - span.set_attribute("ld.ai.graph.key", definition.key) + span.set_attribute("launchdarkly.graph.key", definition.key) context = options.get("context") run_id = str(uuid.uuid4()) start = time.monotonic() @@ -190,7 +190,7 @@ async def build(node: GraphNode) -> None: current_input = final_text total = input_tokens + output_tokens - span.set_attribute("ld.ai.graph.path", "->".join(path)) + span.set_attribute("launchdarkly.graph.path", "->".join(path)) span.set_status(StatusCode.OK) if context is not None: _track( diff --git a/packages/vercel-agents/tests/test_native_graph.py b/packages/vercel-agents/tests/test_native_graph.py index d0558afd..734a9696 100644 --- a/packages/vercel-agents/tests/test_native_graph.py +++ b/packages/vercel-agents/tests/test_native_graph.py @@ -347,7 +347,7 @@ async def test_handoff_path_has_no_duplicate_entries( path_values = [ call.args[1] for call in span.set_attribute.call_args_list - if call.args[0] == "ld.ai.graph.path" + if call.args[0] == "launchdarkly.graph.path" ] assert path_values path = path_values[-1].split("->") @@ -389,7 +389,7 @@ async def test_graph_span_ends_on_success(self, ai_runtime: MagicMock) -> None: tracer.start_span.return_value = span with patch.object(native_graph_mod.trace, "get_tracer", return_value=tracer): await to_vercel_agents(_definition(_graph())).invoke("hello") - tracer.start_span.assert_called_once_with("ld.ai.graph") + tracer.start_span.assert_called_once_with("launchdarkly.graph") span.end.assert_called_once() @pytest.mark.asyncio From 46eb9722766c96f9e7ac9166dad567ce04b8be73 Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Mon, 28 Sep 2026 11:02:28 -0800 Subject: [PATCH 4/5] fix(AIC-3408): mark cancelled Vercel runs and accept handoff arguments CancelledError is not an Exception, so those spans were ending unmarked or as abandoned. Handoff tools now ignore extra model arguments instead of raising TypeError. Co-authored-by: Cursor --- .../launchdarkly_ai_vercel_agents/handler.py | 37 ++++++++++----- .../native_graph.py | 2 +- packages/vercel-agents/tests/test_handler.py | 47 +++++++++++++++++++ .../vercel-agents/tests/test_native_graph.py | 2 + .../handler.py | 37 ++++++++++----- .../vercel-messages/tests/test_handler.py | 47 +++++++++++++++++++ 6 files changed, 149 insertions(+), 23 deletions(-) diff --git a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/handler.py b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/handler.py index c946874e..ef0598c0 100644 --- a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/handler.py +++ b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/handler.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import base64 import inspect import json @@ -352,6 +353,7 @@ async def run( parent = parent_context_of(span) model_span = start_model_span(cfg, parent) failed = False + cancelled = False try: agent = ai.Agent( tools=build_agent_tools(cfg.get("tools"), tool_handlers, parent=parent) @@ -377,13 +379,20 @@ async def run( mark_ok(model_span) mark_ok(span) return {"output": output, "usage": usage} - except BaseException as exc: - if isinstance(exc, Exception): - failed = True - fail_span(model_span, exc) - fail_span(span, exc) + except asyncio.CancelledError: + # Not an Exception: a timeout or task.cancel() must still mark the run + # cancelled, matching the other handlers and the native graph path. + cancelled = True + raise + except Exception as exc: + failed = True + fail_span(model_span, exc) + fail_span(span, exc) raise finally: + if cancelled: + model_span.set_attribute("launchdarkly.run.cancelled", True) + span.set_attribute("launchdarkly.run.cancelled", True) if not failed: model_span.end() span.end() @@ -401,6 +410,7 @@ async def stream( model_span = start_model_span(cfg, parent) completed = False failed = False + cancelled = False provider_stream: Any = None try: agent = ai.Agent( @@ -430,14 +440,19 @@ async def stream( "output": output_of(provider_stream), "usage": usage, } - except BaseException as exc: - if isinstance(exc, Exception): - failed = True - fail_span(model_span, exc) - fail_span(span, exc) + except asyncio.CancelledError: + cancelled = True + raise + except Exception as exc: + failed = True + fail_span(model_span, exc) + fail_span(span, exc) raise finally: - if not completed and not failed: + if cancelled: + model_span.set_attribute("launchdarkly.run.cancelled", True) + span.set_attribute("launchdarkly.run.cancelled", True) + elif not completed and not failed: model_span.set_attribute("launchdarkly.stream.abandoned", True) span.set_attribute("launchdarkly.stream.abandoned", True) if not failed: diff --git a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/native_graph.py b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/native_graph.py index b567b17c..c7b3d071 100644 --- a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/native_graph.py +++ b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/native_graph.py @@ -35,7 +35,7 @@ def _handoff_tool( description: str | None, selected: dict[str, str | None], ) -> Any: - async def transfer() -> str: + async def transfer(**_kwargs: Any) -> str: selected[source] = target return target diff --git a/packages/vercel-agents/tests/test_handler.py b/packages/vercel-agents/tests/test_handler.py index 5453964e..183e825e 100644 --- a/packages/vercel-agents/tests/test_handler.py +++ b/packages/vercel-agents/tests/test_handler.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -331,6 +332,52 @@ async def test_early_stream_exit_closes_context( assert provider_stream.exited == 1 assert provider_stream.consumed == 1 + @pytest.mark.asyncio + async def test_blocking_cancellation_marks_run_cancelled( + self, ai_runtime: MagicMock + ) -> None: + ai_runtime.Agent.return_value.run.return_value = _CancelledContext() + root, model = MagicMock(), MagicMock() + with ( + patch.object(handler_mod, "start_root_span", return_value=root), + patch.object(handler_mod, "start_model_span", return_value=model), + pytest.raises(asyncio.CancelledError), + ): + await create_vercel_agents_handler()(CONFIG, "hello") + for span in (root, model): + span.set_attribute.assert_any_call("launchdarkly.run.cancelled", True) + span.end.assert_called_once() + span.record_exception.assert_not_called() + + @pytest.mark.asyncio + async def test_stream_cancellation_is_not_reported_as_abandoned( + self, ai_runtime: MagicMock + ) -> None: + ai_runtime.Agent.return_value.run.return_value = _CancelledContext() + root, model = MagicMock(), MagicMock() + stream = await create_vercel_agents_handler().stream(CONFIG, "hello") + with ( + patch.object(handler_mod, "start_root_span", return_value=root), + patch.object(handler_mod, "start_model_span", return_value=model), + pytest.raises(asyncio.CancelledError), + ): + async for _event in stream: + pass + for span in (root, model): + span.set_attribute.assert_any_call("launchdarkly.run.cancelled", True) + assert "launchdarkly.stream.abandoned" not in [ + call.args[0] for call in span.set_attribute.call_args_list + ] + span.end.assert_called_once() + + +class _CancelledContext: + async def __aenter__(self) -> None: + raise asyncio.CancelledError + + async def __aexit__(self, *_args: Any) -> None: + return None + class TestGraphWrapper: def test_prewires_exactly_one_wildcard_handler(self) -> None: diff --git a/packages/vercel-agents/tests/test_native_graph.py b/packages/vercel-agents/tests/test_native_graph.py index 734a9696..687751e5 100644 --- a/packages/vercel-agents/tests/test_native_graph.py +++ b/packages/vercel-agents/tests/test_native_graph.py @@ -246,6 +246,8 @@ async def test_non_terminal_node_has_transfer_tool_per_edge( assert [ tool.name for tool in root_tools if tool.name.startswith("transfer_to_") ] == ["transfer_to_leaf"] + transfer = next(tool for tool in root_tools if tool.name == "transfer_to_leaf") + assert await transfer.execute(reason="the leaf owns this") == "leaf" @pytest.mark.asyncio async def test_terminal_node_has_no_handoff_tools( diff --git a/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/handler.py b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/handler.py index dcb3cfa2..f46454f9 100644 --- a/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/handler.py +++ b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/handler.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import base64 import inspect import json @@ -406,6 +407,7 @@ async def invoke( parent = parent_context_of(span) model_span = start_model_span(cfg, parent) failed = False + cancelled = False try: text = "" usage = {"input_tokens": 0, "output_tokens": 0} @@ -434,13 +436,20 @@ async def invoke( mark_ok(model_span) mark_ok(span) return {"output": text, "usage": usage} - except BaseException as exc: - if isinstance(exc, Exception): - failed = True - fail_span(model_span, exc) - fail_span(span, exc) + except asyncio.CancelledError: + # Not an Exception: a timeout or task.cancel() must still mark the run + # cancelled, matching the other handlers. + cancelled = True + raise + except Exception as exc: + failed = True + fail_span(model_span, exc) + fail_span(span, exc) raise finally: + if cancelled: + model_span.set_attribute("launchdarkly.run.cancelled", True) + span.set_attribute("launchdarkly.run.cancelled", True) if not failed: model_span.end() span.end() @@ -458,6 +467,7 @@ async def stream( model_span = start_model_span(cfg, parent) completed = False failed = False + cancelled = False try: output = "" usage = {"input_tokens": 0, "output_tokens": 0} @@ -491,14 +501,19 @@ async def stream( mark_ok(model_span) mark_ok(span) yield {"type": "done", "output": output, "usage": usage} - except BaseException as exc: - if isinstance(exc, Exception): - failed = True - fail_span(model_span, exc) - fail_span(span, exc) + except asyncio.CancelledError: + cancelled = True + raise + except Exception as exc: + failed = True + fail_span(model_span, exc) + fail_span(span, exc) raise finally: - if not completed and not failed: + if cancelled: + model_span.set_attribute("launchdarkly.run.cancelled", True) + span.set_attribute("launchdarkly.run.cancelled", True) + elif not completed and not failed: model_span.set_attribute("launchdarkly.stream.abandoned", True) span.set_attribute("launchdarkly.stream.abandoned", True) if not failed: diff --git a/packages/vercel-messages/tests/test_handler.py b/packages/vercel-messages/tests/test_handler.py index 6666db4a..0ea215c6 100644 --- a/packages/vercel-messages/tests/test_handler.py +++ b/packages/vercel-messages/tests/test_handler.py @@ -6,6 +6,7 @@ from __future__ import annotations +import asyncio import json from types import SimpleNamespace from typing import Any @@ -519,6 +520,44 @@ async def test_early_exit_closes_async_stream_context( assert provider_stream.exited == 1 assert provider_stream.consumed == 1 + @pytest.mark.asyncio + async def test_blocking_cancellation_marks_run_cancelled( + self, ai_runtime: MagicMock + ) -> None: + ai_runtime.stream.return_value = _CancelledContext() + root, model = MagicMock(), MagicMock() + with ( + patch.object(handler_mod, "start_root_span", return_value=root), + patch.object(handler_mod, "start_model_span", return_value=model), + pytest.raises(asyncio.CancelledError), + ): + await create_vercel_messages_handler()(CONFIG, "hello") + for span in (root, model): + span.set_attribute.assert_any_call("launchdarkly.run.cancelled", True) + span.end.assert_called_once() + span.record_exception.assert_not_called() + + @pytest.mark.asyncio + async def test_stream_cancellation_is_not_reported_as_abandoned( + self, ai_runtime: MagicMock + ) -> None: + ai_runtime.stream.return_value = _CancelledContext() + root, model = MagicMock(), MagicMock() + stream = await create_vercel_messages_handler().stream(CONFIG, "hello") + with ( + patch.object(handler_mod, "start_root_span", return_value=root), + patch.object(handler_mod, "start_model_span", return_value=model), + pytest.raises(asyncio.CancelledError), + ): + async for _event in stream: + pass + for span in (root, model): + span.set_attribute.assert_any_call("launchdarkly.run.cancelled", True) + assert "launchdarkly.stream.abandoned" not in [ + call.args[0] for call in span.set_attribute.call_args_list + ] + span.end.assert_called_once() + @pytest.mark.asyncio async def test_runtime_telemetry_is_not_enabled( self, ai_runtime: MagicMock @@ -543,3 +582,11 @@ async def test_span_identity_uses_provider_and_preserves_gateway_model( assert attributes["gen_ai.system"] == "anthropic" assert attributes["gen_ai.provider.name"] == "anthropic" assert attributes["gen_ai.request.model"] == "anthropic/claude-sonnet-4" + + +class _CancelledContext: + async def __aenter__(self) -> None: + raise asyncio.CancelledError + + async def __aexit__(self, *_args: Any) -> None: + return None From 89c52d5a7c0251ae7eead438e301dd0015cd2125 Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Wed, 30 Sep 2026 10:45:42 -0800 Subject: [PATCH 5/5] fix(AIC-3408): split Vercel chat spans by turn and record graph usage Tool loops were collapsing every inference into one span and forcing the final schema onto intermediate tool calls, and the native graph span omitted the token totals the other adapters write. --- .../launchdarkly_ai_vercel_agents/handler.py | 107 ++++++++++++--- .../native_graph.py | 3 + packages/vercel-agents/tests/test_handler.py | 45 ++++++ .../vercel-agents/tests/test_native_graph.py | 6 + .../handler.py | 129 +++++++++++------- .../vercel-messages/tests/test_handler.py | 69 ++++++++++ 6 files changed, 287 insertions(+), 72 deletions(-) diff --git a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/handler.py b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/handler.py index ef0598c0..75b80a57 100644 --- a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/handler.py +++ b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/handler.py @@ -5,6 +5,7 @@ import inspect import json from collections.abc import AsyncGenerator, Callable +from types import SimpleNamespace from typing import Any import ai @@ -326,6 +327,78 @@ def build_output_type( return create_model(name, __config__=ConfigDict(extra="forbid"), **fields) +class _ChatTurns: + """One ``chat`` span per model turn. + + ``Agent.run`` loops inferences inside a single call and brackets each one + with ``stream_start`` / ``stream_end``. A run that never emits those + boundaries — the in-memory test double — keeps the span opened up front + and records the run's usage on it. + """ + + def __init__(self, cfg: AiConfigRep, parent: Any) -> None: + self._cfg = cfg + self._parent = parent + self._open: Any = start_model_span(cfg, parent) + self._closed_a_turn = False + + def observe(self, event: Any) -> str | None: + kind = getattr(event, "kind", None) + if kind == "stream_end": + if self._open is None: + self._open = start_model_span(self._cfg, self._parent) + self._close(getattr(event, "usage", None)) + self._closed_a_turn = True + elif kind == "stream_start" and self._open is None: + self._open = start_model_span(self._cfg, self._parent) + return text_delta(event) + + def finish(self, usage: dict[str, int]) -> None: + """Close a turn that never emitted ``stream_end``.""" + if self._open is None: + return + if self._closed_a_turn: + self._close(None) + return + self._close( + SimpleNamespace( + input_tokens=usage["input_tokens"], + output_tokens=usage["output_tokens"], + ) + ) + + def cancel(self) -> None: + if self._open is None: + return + self._open.set_attribute("launchdarkly.run.cancelled", True) + self._open.end() + self._open = None + + def fail(self, exc: BaseException) -> None: + if self._open is None: + return + fail_span(self._open, exc) + self._open = None + + def abandon(self) -> None: + if self._open is None: + return + self._open.set_attribute("launchdarkly.stream.abandoned", True) + self._open.end() + self._open = None + + def _close(self, usage: Any) -> None: + incoming, outgoing = _usage_values(usage) + finish_model_span( + self._open, + model_name(self._cfg), + SpanUsage(input=incoming, output=outgoing), + ) + mark_ok(self._open) + self._open.end() + self._open = None + + def text_delta(event: Any) -> str | None: if getattr(event, "kind", None) == "text_delta": return str(getattr(event, "chunk", "")) @@ -351,7 +424,7 @@ async def run( vs = variables or {} span = start_root_span(cfg, vs) parent = parent_context_of(span) - model_span = start_model_span(cfg, parent) + turns = _ChatTurns(cfg, parent) failed = False cancelled = False try: @@ -366,35 +439,32 @@ async def run( if cfg.get("outputFormat"): run_kwargs["output_type"] = build_output_type(cfg["outputFormat"]) async with agent.run(**run_kwargs) as provider_stream: - async for _ in provider_stream: - pass + async for event in provider_stream: + turns.observe(event) usage = usage_of(provider_stream) output = output_of(provider_stream) + turns.finish(usage) span_usage = SpanUsage( input=usage["input_tokens"], output=usage["output_tokens"] ) - response_model = model_name(cfg) - finish_model_span(model_span, response_model, span_usage) - finish_root_span(span, response_model, span_usage) - mark_ok(model_span) + finish_root_span(span, model_name(cfg), span_usage) mark_ok(span) return {"output": output, "usage": usage} except asyncio.CancelledError: # Not an Exception: a timeout or task.cancel() must still mark the run # cancelled, matching the other handlers and the native graph path. cancelled = True + turns.cancel() raise except Exception as exc: failed = True - fail_span(model_span, exc) + turns.fail(exc) fail_span(span, exc) raise finally: if cancelled: - model_span.set_attribute("launchdarkly.run.cancelled", True) span.set_attribute("launchdarkly.run.cancelled", True) if not failed: - model_span.end() span.end() async def stream( @@ -407,7 +477,7 @@ async def stream( vs = variables or {} span = start_root_span(cfg, vs) parent = parent_context_of(span) - model_span = start_model_span(cfg, parent) + turns = _ChatTurns(cfg, parent) completed = False failed = False cancelled = False @@ -422,18 +492,16 @@ async def stream( params=build_request_params(cfg), ) as provider_stream: async for event in provider_stream: - text = text_delta(event) + text = turns.observe(event) if text is not None: yield {"type": "chunk", "text": text} completed = True usage = usage_of(provider_stream) + turns.finish(usage) span_usage = SpanUsage( input=usage["input_tokens"], output=usage["output_tokens"] ) - response_model = model_name(cfg) - finish_model_span(model_span, response_model, span_usage) - finish_root_span(span, response_model, span_usage) - mark_ok(model_span) + finish_root_span(span, model_name(cfg), span_usage) mark_ok(span) yield { "type": "done", @@ -442,21 +510,20 @@ async def stream( } except asyncio.CancelledError: cancelled = True + turns.cancel() raise except Exception as exc: failed = True - fail_span(model_span, exc) + turns.fail(exc) fail_span(span, exc) raise finally: if cancelled: - model_span.set_attribute("launchdarkly.run.cancelled", True) span.set_attribute("launchdarkly.run.cancelled", True) elif not completed and not failed: - model_span.set_attribute("launchdarkly.stream.abandoned", True) + turns.abandon() span.set_attribute("launchdarkly.stream.abandoned", True) if not failed: - model_span.end() span.end() return create_handler(("*", "agent"), run, stream, capture_content=capture_content) diff --git a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/native_graph.py b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/native_graph.py index c7b3d071..5b0784cc 100644 --- a/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/native_graph.py +++ b/packages/vercel-agents/src/launchdarkly_ai_vercel_agents/native_graph.py @@ -191,6 +191,9 @@ async def build(node: GraphNode) -> None: total = input_tokens + output_tokens span.set_attribute("launchdarkly.graph.path", "->".join(path)) + span.set_attribute("gen_ai.usage.input_tokens", input_tokens) + span.set_attribute("gen_ai.usage.output_tokens", output_tokens) + span.set_attribute("gen_ai.usage.total_tokens", total) span.set_status(StatusCode.OK) if context is not None: _track( diff --git a/packages/vercel-agents/tests/test_handler.py b/packages/vercel-agents/tests/test_handler.py index 183e825e..da2ad7b9 100644 --- a/packages/vercel-agents/tests/test_handler.py +++ b/packages/vercel-agents/tests/test_handler.py @@ -278,6 +278,51 @@ async def test_sums_usage_across_every_turn_of_a_tool_loop( result = await create_vercel_agents_handler()(CONFIG, "hello") assert result["usage"] == {"input_tokens": 63, "output_tokens": 15} + @pytest.mark.asyncio + async def test_opens_one_chat_span_per_model_turn( + self, ai_runtime: MagicMock + ) -> None: + class Boundary: + def __init__( + self, kind: str, input_tokens: int = 0, output_tokens: int = 0 + ) -> None: + self.kind = kind + self.usage = SimpleNamespace( + input_tokens=input_tokens, output_tokens=output_tokens + ) + + ai_runtime.Agent.return_value.run.return_value = AgentStream( + [ + Boundary("stream_start"), + TextDelta("a"), + Boundary("stream_end", 3, 1), + Boundary("stream_start"), + TextDelta("b"), + Boundary("stream_end", 4, 2), + ], + text="ab", + input_tokens=7, + output_tokens=3, + ) + first, second = MagicMock(), MagicMock() + with patch.object(handler_mod, "start_model_span", side_effect=[first, second]): + result = await create_vercel_agents_handler()(CONFIG, "hello") + + def usage(span: MagicMock) -> dict[str, int]: + return { + call.args[0]: call.args[1] + for call in span.set_attribute.call_args_list + if str(call.args[0]).startswith("gen_ai.usage.") + } + + assert result["usage"] == {"input_tokens": 7, "output_tokens": 3} + assert usage(first)["gen_ai.usage.input_tokens"] == 3 + assert usage(first)["gen_ai.usage.output_tokens"] == 1 + assert usage(second)["gen_ai.usage.input_tokens"] == 4 + assert usage(second)["gen_ai.usage.output_tokens"] == 2 + first.end.assert_called_once() + second.end.assert_called_once() + @pytest.mark.asyncio async def test_structured_output_is_requested_and_serialized( self, ai_runtime: MagicMock diff --git a/packages/vercel-agents/tests/test_native_graph.py b/packages/vercel-agents/tests/test_native_graph.py index 687751e5..341cbb6b 100644 --- a/packages/vercel-agents/tests/test_native_graph.py +++ b/packages/vercel-agents/tests/test_native_graph.py @@ -355,6 +355,12 @@ async def test_handoff_path_has_no_duplicate_entries( path = path_values[-1].split("->") assert path == ["root", "leaf"] assert len(path) == len(set(path)) + attributes = { + call.args[0]: call.args[1] for call in span.set_attribute.call_args_list + } + assert attributes["gen_ai.usage.input_tokens"] == 14 + assert attributes["gen_ai.usage.output_tokens"] == 5 + assert attributes["gen_ai.usage.total_tokens"] == 19 class TestTelemetryAndCleanup: diff --git a/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/handler.py b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/handler.py index f46454f9..6b6a25e1 100644 --- a/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/handler.py +++ b/packages/vercel-messages/src/launchdarkly_ai_vercel_messages/handler.py @@ -354,40 +354,79 @@ async def _run_conversation( input_tokens = 0 output_tokens = 0 - for _ in range(MAX_STEPS): - kwargs: dict[str, Any] = { - "model": resolved_model, - "messages": messages, - "tools": tools, - "params": params, - } - if output_type is not None: - kwargs["output_type"] = output_type - async with ai.stream(**kwargs) as provider_stream: - async for event in provider_stream: - text = _text_delta(event) - if text is not None: - yield {"type": "chunk", "text": text} - usage = _usage(provider_stream) - input_tokens += usage["input_tokens"] - output_tokens += usage["output_tokens"] - message = provider_stream.message - calls = list(message.tool_calls or []) - if not calls: - yield { - "type": "done", - "output": _final_output(provider_stream), - "usage": { - "input_tokens": input_tokens, - "output_tokens": output_tokens, - }, - } - return - results = [await _tool_result(call, executors, parent) for call in calls] - messages = [*messages, message, ai.tool_message(*results)] - raise RuntimeError( - f"Vercel messages run did not reach a final response within {MAX_STEPS} steps" - ) + # A schema on a turn that can still call tools stops the model from emitting + # those calls, or fails validation before a final answer exists. Tool turns + # stay unconstrained; the schema is applied only once the model has stopped + # calling tools (or immediately, when the config has no tools). + constrain_output = output_type is not None and not tools + open_model_span: Any = None + try: + for _ in range(MAX_STEPS): + model_span = start_model_span(cfg, parent) + open_model_span = model_span + try: + kwargs: dict[str, Any] = { + "model": resolved_model, + "messages": messages, + "tools": [] if constrain_output else tools, + "params": params, + } + if constrain_output and output_type is not None: + kwargs["output_type"] = output_type + async with ai.stream(**kwargs) as provider_stream: + async for event in provider_stream: + text = _text_delta(event) + if text is not None: + yield {"type": "chunk", "text": text} + usage = _usage(provider_stream) + input_tokens += usage["input_tokens"] + output_tokens += usage["output_tokens"] + # This turn's counts, not the run total. The root span carries the sum. + finish_model_span( + model_span, + model_name(cfg), + SpanUsage( + input=usage["input_tokens"], output=usage["output_tokens"] + ), + ) + mark_ok(model_span) + model_span.end() + open_model_span = None + except asyncio.CancelledError: + model_span.set_attribute("launchdarkly.run.cancelled", True) + model_span.end() + open_model_span = None + raise + except Exception as exc: + fail_span(model_span, exc) + open_model_span = None + raise + message = provider_stream.message + calls = list(message.tool_calls or []) + if not calls or constrain_output: + if output_type is not None and tools and not constrain_output: + constrain_output = True + continue + yield { + "type": "done", + "output": _final_output(provider_stream), + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + }, + } + return + results = [await _tool_result(call, executors, parent) for call in calls] + messages = [*messages, message, ai.tool_message(*results)] + raise RuntimeError( + f"Vercel messages run did not reach a final response within {MAX_STEPS} steps" + ) + finally: + # A consumer that stops reading injects GeneratorExit, which neither + # except above sees. The chat span for the in-flight turn is still open. + if open_model_span is not None: + open_model_span.set_attribute("launchdarkly.stream.abandoned", True) + open_model_span.end() def create_vercel_messages_handler( @@ -405,7 +444,6 @@ async def invoke( vs = variables or {} span = start_root_span(cfg, vs) parent = parent_context_of(span) - model_span = start_model_span(cfg, parent) failed = False cancelled = False try: @@ -430,28 +468,23 @@ async def invoke( span_usage = SpanUsage( input=usage["input_tokens"], output=usage["output_tokens"] ) - response_model = model_name(cfg) - finish_model_span(model_span, response_model, span_usage) - finish_root_span(span, response_model, span_usage) - mark_ok(model_span) + finish_root_span(span, model_name(cfg), span_usage) mark_ok(span) return {"output": text, "usage": usage} except asyncio.CancelledError: # Not an Exception: a timeout or task.cancel() must still mark the run - # cancelled, matching the other handlers. + # cancelled, matching the other handlers. The in-flight chat span is + # marked inside _run_conversation, where that turn was opened. cancelled = True raise except Exception as exc: failed = True - fail_span(model_span, exc) fail_span(span, exc) raise finally: if cancelled: - model_span.set_attribute("launchdarkly.run.cancelled", True) span.set_attribute("launchdarkly.run.cancelled", True) if not failed: - model_span.end() span.end() async def stream( @@ -464,7 +497,6 @@ async def stream( vs = variables or {} span = start_root_span(cfg, vs) parent = parent_context_of(span) - model_span = start_model_span(cfg, parent) completed = False failed = False cancelled = False @@ -495,10 +527,7 @@ async def stream( span_usage = SpanUsage( input=usage["input_tokens"], output=usage["output_tokens"] ) - response_model = model_name(cfg) - finish_model_span(model_span, response_model, span_usage) - finish_root_span(span, response_model, span_usage) - mark_ok(model_span) + finish_root_span(span, model_name(cfg), span_usage) mark_ok(span) yield {"type": "done", "output": output, "usage": usage} except asyncio.CancelledError: @@ -506,18 +535,14 @@ async def stream( raise except Exception as exc: failed = True - fail_span(model_span, exc) fail_span(span, exc) raise finally: if cancelled: - model_span.set_attribute("launchdarkly.run.cancelled", True) span.set_attribute("launchdarkly.run.cancelled", True) elif not completed and not failed: - model_span.set_attribute("launchdarkly.stream.abandoned", True) span.set_attribute("launchdarkly.stream.abandoned", True) if not failed: - model_span.end() span.end() return create_handler( diff --git a/packages/vercel-messages/tests/test_handler.py b/packages/vercel-messages/tests/test_handler.py index 0ea215c6..62bec864 100644 --- a/packages/vercel-messages/tests/test_handler.py +++ b/packages/vercel-messages/tests/test_handler.py @@ -375,6 +375,41 @@ async def test_executes_tool_calls_and_continues_until_a_final_answer( "tc-1", tool_name="weather", result="sunny" ) + @pytest.mark.asyncio + async def test_tool_loop_records_each_turn_on_its_own_chat_span( + self, ai_runtime: MagicMock + ) -> None: + config = { + **CONFIG, + "tools": {"weather": {"description": "Get weather", "parameters": {}}}, + } + call = SimpleNamespace( + tool_call_id="tc-1", tool_name="weather", tool_args={"city": "Oakland"} + ) + ai_runtime.stream.side_effect = [ + FakeStream(text="", tool_calls=[call], input_tokens=10, output_tokens=4), + FakeStream(text="sunny in Oakland", input_tokens=7, output_tokens=2), + ] + first, second = MagicMock(), MagicMock() + with patch.object(handler_mod, "start_model_span", side_effect=[first, second]): + await create_vercel_messages_handler()( + config, "hello", {"weather": AsyncMock(return_value="sunny")} + ) + + def usage(span: MagicMock) -> dict[str, int]: + return { + call.args[0]: call.args[1] + for call in span.set_attribute.call_args_list + if call.args[0].startswith("gen_ai.usage.") + } + + assert usage(first)["gen_ai.usage.input_tokens"] == 10 + assert usage(first)["gen_ai.usage.output_tokens"] == 4 + assert usage(second)["gen_ai.usage.input_tokens"] == 7 + assert usage(second)["gen_ai.usage.output_tokens"] == 2 + first.end.assert_called_once() + second.end.assert_called_once() + @pytest.mark.asyncio async def test_reports_a_failing_tool_back_to_the_model( self, ai_runtime: MagicMock @@ -434,6 +469,40 @@ async def test_structured_output_uses_output_type_and_serializes_result( assert output_schema["additionalProperties"] is False assert json.loads(result["output"]) == {"answer": "yes"} + @pytest.mark.asyncio + async def test_output_type_waits_until_the_model_stops_calling_tools( + self, ai_runtime: MagicMock + ) -> None: + schema = { + "type": "object", + "properties": {"answer": {"type": "string"}}, + "required": ["answer"], + } + config = { + **CONFIG, + "outputFormat": schema, + "tools": {"weather": {"description": "Get weather", "parameters": {}}}, + } + call = SimpleNamespace( + tool_call_id="tc-1", tool_name="weather", tool_args={"city": "Oakland"} + ) + ai_runtime.stream.side_effect = [ + FakeStream(text="", tool_calls=[call]), + FakeStream(text="sunny"), + FakeStream(text="", output='{"answer":"yes"}'), + ] + result = await create_vercel_messages_handler()( + config, "question", {"weather": AsyncMock(return_value="sunny")} + ) + rounds = ai_runtime.stream.call_args_list + assert len(rounds) == 3 + assert "output_type" not in rounds[0].kwargs + assert "output_type" not in rounds[1].kwargs + assert rounds[1].kwargs["tools"] + assert "output_type" in rounds[2].kwargs + assert rounds[2].kwargs["tools"] == [] + assert result["output"] == '{"answer":"yes"}' + @pytest.mark.asyncio async def test_structured_output_types_nested_objects_and_array_items( self, ai_runtime: MagicMock