chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 12:58:18 +08:00
commit 6d5d58c1a9
18293 changed files with 3502153 additions and 0 deletions
+124
View File
@@ -0,0 +1,124 @@
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
cover/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
.python-version
# celery beat schedule file
celerybeat-schedule
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyderworkspace
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# pytype static type analyzer
.pytype/
# Cython debug symbols
cython_debug/
+63
View File
@@ -0,0 +1,63 @@
# CopilotKit Python SDK
[![PyPI version](https://badge.fury.io/py/copilotkit.svg)](https://badge.fury.io/py/copilotkit)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
The official Python SDK for CopilotKit - build AI copilots and agents into your applications.
## Features
- 🚀 Easy integration with LangGraph and LangChain
- 🔄 Built-in support for stateful conversations
- 🛠 Extensible agent framework
- 🔌 FastAPI-ready endpoints
- 🤝 Optional CrewAI integration
## Installation
```bash
pip install copilotkit
```
With CrewAI support:
```bash
pip install "copilotkit[crewai]"
```
## Quick Start
```python
from copilotkit import Copilot
# Initialize a copilot
copilot = Copilot()
# Add your tools and configure the copilot
copilot.add_tool(my_custom_tool)
# Run the copilot
response = copilot.run("Your task description here")
```
## Documentation
For detailed documentation and examples, visit [copilotkit.ai](https://copilotkit.ai)
## Contributing
We welcome contributions! Please see our [Contributing Guidelines](https://github.com/CopilotKit/CopilotKit/blob/main/CONTRIBUTING.md) for details.
## License
This project is licensed under the MIT License - see the [LICENSE](https://github.com/CopilotKit/CopilotKit/blob/main/LICENSE) file for details.
## Support
- 📚 [Documentation](https://docs.copilotkit.ai)
- 💬 [Discord Community](https://discord.gg/6dffbvGU)
- 🐛 [Issue Tracker](https://github.com/CopilotKit/CopilotKit/issues)
---
Built with ❤️ by the CopilotKit team
+57
View File
@@ -0,0 +1,57 @@
"""CopilotKit SDK"""
from .sdk import (
CopilotKitRemoteEndpoint,
CopilotKitContext,
CopilotKitSDK,
CopilotKitSDKContext,
)
from .action import Action
from .exc import (
CopilotKitError,
CopilotKitMisuseError,
ActionNotFoundException,
AgentNotFoundException,
ActionExecutionException,
AgentExecutionException,
)
from .langgraph import CopilotKitState
from .parameter import Parameter
from .agent import Agent
from .langgraph_agui_agent import LangGraphAGUIAgent
from .copilotkit_lg_middleware import CopilotKitMiddleware
from ag_ui_langgraph.middlewares.state_streaming import (
StateStreamingMiddleware,
StateItem,
)
from .header_propagation import (
set_forwarded_headers,
get_forwarded_headers,
install_httpx_hook,
)
__all__ = [
"CopilotKitRemoteEndpoint",
"CopilotKitSDK",
"Action",
"CopilotKitState",
"Parameter",
"Agent",
"CopilotKitContext",
"CopilotKitSDKContext",
"CopilotKitError",
"CopilotKitMisuseError",
"ActionNotFoundException",
"AgentNotFoundException",
"ActionExecutionException",
"AgentExecutionException",
"CrewAIAgent", # pyright: ignore[reportUnsupportedDunderAll] pylint: disable=undefined-all-variable
"LangGraphAGUIAgent",
"CopilotKitMiddleware",
"StateStreamingMiddleware",
"StateItem",
"set_forwarded_headers",
"get_forwarded_headers",
"install_httpx_hook",
]
+240
View File
@@ -0,0 +1,240 @@
"""
A2UI helpers — build v0.9 A2UI operations from schema + data.
Usage:
from copilotkit import a2ui
schema = a2ui.load_schema("flight_card.json")
@tool
def search_flights(flights: list[Flight]) -> str:
return a2ui.render([
a2ui.create_surface("my-surface"),
a2ui.update_components("my-surface", schema),
a2ui.update_data_model("my-surface", {"flights": flights}),
])
"""
from __future__ import annotations
import json
from pathlib import Path
from typing import Any
def load_schema(path: str | Path) -> list[dict[str, Any]]:
"""Load an A2UI component schema from a JSON file."""
with open(path) as f:
return json.load(f)
def update_components(
surface_id: str,
components: list[dict[str, Any]],
) -> dict[str, Any]:
"""Build a v0.9 updateComponents operation."""
return {
"version": "v0.9",
"updateComponents": {
"surfaceId": surface_id,
"components": components,
},
}
def update_data_model(
surface_id: str,
data: Any,
path: str = "/",
) -> dict[str, Any]:
"""Build a v0.9 updateDataModel operation with plain JSON value."""
return {
"version": "v0.9",
"updateDataModel": {
"surfaceId": surface_id,
"path": path,
"value": data,
},
}
BASIC_CATALOG_ID = "https://a2ui.org/specification/v0_9/basic_catalog.json"
"""The catalog ID for the standard v0.9 basic catalog."""
def create_surface(
surface_id: str,
catalog_id: str = BASIC_CATALOG_ID,
) -> dict[str, Any]:
"""Build a v0.9 createSurface operation."""
return {
"version": "v0.9",
"createSurface": {
"surfaceId": surface_id,
"catalogId": catalog_id,
},
}
A2UI_OPERATIONS_KEY = "a2ui_operations"
"""The container key used to wrap A2UI operations for explicit detection."""
def render(operations: list[dict[str, Any]]) -> str:
"""Wrap operations in the a2ui_operations container and serialize to JSON.
Args:
operations: The A2UI v0.9 operations (createSurface, updateComponents, updateDataModel).
Example::
render(
operations=[...],
)
"""
result: dict[str, Any] = {A2UI_OPERATIONS_KEY: operations}
return json.dumps(result)
# ---------------------------------------------------------------------------
# Dynamic A2UI prompt builder
# ---------------------------------------------------------------------------
DEFAULT_GENERATION_GUIDELINES = """\
Generate A2UI v0.9 JSON.
## A2UI Protocol Instructions
A2UI (Agent to UI) is a protocol for rendering rich UI surfaces from agent responses.
CRITICAL: You MUST call the render_a2ui tool with ALL of these arguments:
- surfaceId: A unique ID for the surface (e.g. "product-comparison")
- components: REQUIRED — the A2UI component array. NEVER omit this. Use a List with
children: { componentId: "card-id", path: "/items" } for repeating cards.
- data: OPTIONAL — a JSON object written to the root of the surface data model.
Use for pre-filling form values or providing data for path-bound components.
- every component must have the "component" field specifying the component type (e.g. "Text", "Image", "Row", "Column", "List", "Button", etc.)
COMPONENT ID RULES:
- Exactly one component MUST have id="root". This is the surface's entry
point — the renderer begins at "root" and walks the child/children tree
from there. Every other component must be reachable from "root". If no
component has id="root", the surface renders an empty loading placeholder
and none of your components will be shown.
- Every component ID must be unique within the surface.
- A component MUST NOT reference itself as child/children. This causes a
circular dependency error. For example, if a component has id="avatar",
its child must be a DIFFERENT id (e.g. "avatar-img"), never "avatar".
- The child/children tree must be a DAG — no cycles allowed.
PATH RULES FOR TEMPLATES:
Components inside a repeating List use RELATIVE paths (no leading slash).
The path is resolved relative to each array item automatically.
If List has children: { componentId: "card", path: "/items" } and item has key "name",
use { "path": "name" } (NO leading slash — relative to item).
CRITICAL: Do NOT use "/name" (absolute) inside templates — use "name" (relative).
The List's own path ("/items") uses a leading slash (absolute), but all
components INSIDE the template card use paths WITHOUT leading slash.
Do NOT use "/items/0/name" or "/items/{@key}/name" — just "name".
COMPONENT VALUES — DEFAULT RULE:
Use inline literal values for ALL component properties. Pass strings, numbers,
arrays, and objects directly on the component. Do NOT use { "path": "..." }
objects unless the property's schema explicitly allows it (see exception below).
CRITICAL: USING { "path": "..." } ON A PROPERTY THAT DOES NOT DECLARE PATH
SUPPORT IN ITS SCHEMA WILL CAUSE A RUNTIME CRASH AND BREAK THE ENTIRE UI.
ALWAYS CHECK THE COMPONENT SCHEMA FIRST — IF THE PROPERTY ONLY ACCEPTS A
PLAIN TYPE, YOU MUST USE A LITERAL VALUE.
VERY IMPORTANT: THE APPLICATION WILL BREAK IF YOU DO NOT FOLLOW THIS RULE!
For example, a chart's "data" must always be an inline array:
"data": [{"label": "Jan", "value": 100}, {"label": "Feb", "value": 200}]
A metric's "value" must always be an inline string:
"value": "$1,200"
PATH BINDING EXCEPTION — SCHEMA-DRIVEN:
A few properties accept { "path": "/some/path" } as an alternative to a literal
value. You can identify these in the Available Components schema: the property
will list BOTH a literal type AND an object-with-path option. If a property only
shows a single type (string, number, array, etc.), it does NOT support path
binding — use a literal value only.
Path binding is typically used for editable form inputs so the client can write
user input back to the data model. When building forms:
- Bind input "value" to a data model path: "value": { "path": "/form/name" }
- Pre-fill via the "data" tool argument: "data": { "form": { "name": "Alice" } }
- Capture values on submit via button action context:
"action": { "event": { "name": "submit", "context": { "name": { "path": "/form/name" } } } }
REPEATING CONTENT uses a structural children format (not the same as value binding):
children: { componentId: "card-id", path: "/items" }
Components inside templates use RELATIVE paths (no leading slash): { "path": "name" }."""
DEFAULT_DESIGN_GUIDELINES = """\
Create polished, visually appealing interfaces:
- Always include a title heading (h2) for the surface, outside the List.
Wrap in a Column: [title, list] as root.
- For card templates, create clear visual hierarchy:
- h3 for primary text (names, titles)
- h2 for featured numbers (prices, scores) — makes them stand out
- caption for secondary info (ratings, categories, metadata)
- body for descriptions
- Use Divider between logical sections within cards.
- Use Row with justify="spaceBetween" for label-value pairs
(e.g. "Rating" on left, "4.5/5" on right).
- Include images when relevant (logos, icons, product photos):
- Use Image component with variant="smallFeature" or "avatar"
- Prefer company logos for branded products — Google favicons are reliable:
https://www.google.com/s2/favicons?domain=sony.com&sz=128
https://www.google.com/s2/favicons?domain=bose.com&sz=128
- For generic icons: https://placehold.co/128x128/EEE/999?text=🎧
- Do NOT invent Unsplash photo-IDs — they will 404. Only use real, known URLs.
- Use horizontal List direction for side-by-side comparison cards.
- Keep cards clean — avoid clutter. Whitespace is good.
- Use consistent surfaceIds (lowercase, hyphenated).
- NEVER use the same ID for a component and its child — this creates a
circular dependency. E.g. if id="avatar", child must NOT be "avatar".
- Both Row and Column support "justify" and "align".
- Add Button for interactivity. Button needs child (Text ID) + action.
Action MUST use this exact nested format:
"action": { "event": { "name": "myAction", "context": { "key": "value" } } }
The "event" key holds an OBJECT with "name" (required) and "context" (optional).
Do NOT use a flat format like {"event": "name"} — "event" must be an object.
Use variant="primary" for main action buttons, variant="borderless" for links.
- For forms: wrap fields in a Card with a Column. Place the submit button in a
Row with justify="end". Every input MUST use path binding on the "value" property
(e.g. "value": { "path": "/form/name" }) to be editable. The submit button's action
context MUST reference the same paths to capture the user's input.
Use the SAME surfaceId as the main surface. Match action names to Button action event names."""
def a2ui_prompt(
component_schema: str,
generation_guidelines: str = DEFAULT_GENERATION_GUIDELINES,
design_guidelines: str = DEFAULT_DESIGN_GUIDELINES,
) -> str:
"""Build the system prompt for dynamic A2UI generation.
Args:
component_schema: JSON string of available components and their props.
Read from state["ag-ui"]["a2ui_schema"].
generation_guidelines: Instructions for how to call the render_a2ui
tool, path rules, and data format.
design_guidelines: Visual design rules, component hierarchy tips,
and action handler patterns.
Returns:
Complete system prompt string.
"""
return f"""\
{generation_guidelines}
## DESIGN GUIDELINES:
{design_guidelines}
## AVAILABLE COMPONENTS:
The following components are available for building UI surfaces.
Use ONLY these components with the specified props.
{component_schema}
"""
+57
View File
@@ -0,0 +1,57 @@
"""Actions"""
import re
from inspect import iscoroutinefunction
from typing import Optional, List, Callable, TypedDict, Any, cast
from .parameter import Parameter, normalize_parameters
class ActionDict(TypedDict):
"""Dict representation of an action"""
name: str
description: str
parameters: List[Parameter]
class ActionResultDict(TypedDict):
"""Dict representation of an action result"""
result: Any
class Action: # pylint: disable=too-few-public-methods
"""Action class for CopilotKit"""
def __init__(
self,
*,
name: str,
handler: Callable,
description: Optional[str] = None,
parameters: Optional[List[Parameter]] = None,
):
self.name = name
self.description = description
self.parameters = parameters
self.handler = handler
if not re.match(r"^[a-zA-Z0-9_-]+$", name):
raise ValueError(
f"Invalid action name '{name}': "
+ "must consist of alphanumeric characters, underscores, and hyphens only"
)
async def execute(self, *, arguments: dict) -> ActionResultDict:
"""Execute the action"""
result = self.handler(**arguments)
return {"result": await result if iscoroutinefunction(self.handler) else result}
def dict_repr(self) -> ActionDict:
"""Dict representation of the action"""
return {
"name": self.name,
"description": self.description or "",
"parameters": normalize_parameters(cast(Any, self.parameters)),
}
+66
View File
@@ -0,0 +1,66 @@
"""Agents"""
import re
from typing import Optional, List, TypedDict
from abc import ABC, abstractmethod
from .types import Message
from .action import ActionDict
from .types import MetaEvent
class AgentDict(TypedDict):
"""Agent dictionary"""
name: str
description: Optional[str]
class Agent(ABC):
"""Agent class for CopilotKit"""
def __init__(
self,
*,
name: str,
description: Optional[str] = None,
):
self.name = name
self.description = description
if not re.match(r"^[a-zA-Z0-9_-]+$", name):
raise ValueError(
f"Invalid agent name '{name}': "
+ "must consist of alphanumeric characters, underscores, and hyphens only"
)
@abstractmethod
def execute( # pylint: disable=too-many-arguments
self,
*,
state: dict,
config: Optional[dict] = None,
messages: List[Message],
thread_id: str,
actions: Optional[List[ActionDict]] = None,
meta_events: Optional[List[MetaEvent]] = None,
**kwargs,
):
"""Execute the agent"""
@abstractmethod
async def get_state(
self,
*,
thread_id: str,
):
"""Default get_state implementation"""
return {
"threadId": thread_id or "",
"threadExists": False,
"state": {},
"messages": [],
}
def dict_repr(self) -> AgentDict:
"""Dict representation of the action"""
return {"name": self.name, "description": self.description or ""}
@@ -0,0 +1,979 @@
"""
CopilotKit Middleware for LangGraph agents.
Works with any agent (prebuilt or custom).
Example:
from langgraph.prebuilt import create_agent
from copilotkit import CopilotKitMiddleware
agent = create_agent(
model="openai:gpt-4o",
tools=[backend_tool],
middleware=[CopilotKitMiddleware()],
)
"""
import json
import re
from typing import Any, Callable, Awaitable, ClassVar, Iterable, Optional, Union
from langchain_core.messages import AIMessage, SystemMessage, ToolMessage
from langchain.agents.middleware import (
AgentMiddleware,
AgentState,
ModelRequest,
ModelResponse,
)
from langgraph.runtime import Runtime
from .header_propagation import install_httpx_hook, set_forwarded_headers
from .langgraph import CopilotKitProperties
# Optional dependency: the A2UI subagent-tool factory ships in ag-ui-langgraph.
# Guarded so an older/skewed version without the factory degrades to
# "no auto-A2UI" instead of breaking the whole middleware import.
try: # pragma: no cover - exercised indirectly via the a2ui injection path
from ag_ui_langgraph import get_a2ui_tools, A2UIToolParams
except Exception: # noqa: BLE001 - any import failure means the feature is off
get_a2ui_tools = None
A2UIToolParams = None
# Track which httpx clients already have the header-propagation hook installed
# (by object id) so we never double-install on repeated model calls.
_hooked_clients: set[int] = set()
# ---------------------------------------------------------------------------
# Auto-A2UI: bridge the inferred model from the model-call hook to the
# tool-call hook
# ---------------------------------------------------------------------------
# The generate_a2ui tool drives a structured-output subagent and so needs a
# chat model. We "infer" that model from ``request.model`` in
# ``wrap_model_call`` (the only hook that exposes the bound model) and reuse it.
# But the tool actually *executes* later in ``wrap_tool_call``, whose request
# does NOT carry the model. ContextVars do not reliably survive LangGraph node
# boundaries, so we bridge the built tool across nodes via a module-level map
# keyed by the run's thread id.
_a2ui_tools_by_thread: dict[str, Any] = {}
# Fallback key for runs without a thread id (e.g. an in-memory invoke with no
# checkpointer). Collisions across concurrent context-less runs are an
# acceptable edge — the deployed path always carries a thread id.
_DEFAULT_THREAD_KEY = "__copilotkit_a2ui_default__"
def _current_thread_id() -> "str | None":
"""Best-effort read of the active run's thread id from the LangGraph config.
Returns ``None`` outside a runnable context (e.g. unit tests); callers then
fall back to ``_DEFAULT_THREAD_KEY``.
"""
try:
from langgraph.config import get_config
cfg = get_config() or {}
return (cfg.get("configurable") or {}).get("thread_id")
except Exception: # noqa: BLE001 - no active context / older langgraph
return None
def _extract_forwarded_headers_from_config() -> None:
"""Extract raw ``x-*`` headers from the current LangGraph RunnableConfig and
push them into the header-propagation ContextVar so the httpx hook can
forward them on outgoing LLM requests.
When an agent runs inside **langgraph-api** with
``LANGGRAPH_HTTP={"configurable_headers":{"include":["x-*"]}}``,
the server copies inbound HTTP ``x-*`` headers into
``config["configurable"]`` as individual keys (e.g.
``configurable["x-aimock-context"] = "value"``). This function reads those
keys and calls :func:`set_forwarded_headers` so they propagate to the
underlying LLM provider SDK via the httpx event hook.
Precedence: the wrapper dict ``copilotkit_forwarded_headers`` (if present)
takes priority over raw ``x-*`` keys. Raw keys are only used when the
wrapper dict is absent or does not contain a given header.
Safe to call outside a runnable context (e.g. in unit tests) — silently
returns without doing anything if ``get_config()`` raises.
"""
try:
from langgraph.config import (
get_config,
) # local import to avoid hard dep at module level
config = get_config()
except ImportError:
return
except RuntimeError:
# No active runnable context — clear the ContextVar so stale headers
# from a prior request in the same async context do not leak through.
set_forwarded_headers({})
return
try:
headers: dict[str, str] = {}
# Sources to scan: config["context"] (LangGraph >=0.6.0) and
# config["configurable"] (all versions).
context = config.get("context") or {}
configurable = config.get("configurable") or {}
# 1) Wrapper-dict path (highest priority): these are headers that
# CopilotKit explicitly bundled under a known key. Process context
# first with first-write-wins so context takes precedence over
# configurable (LangGraph >=0.6.0 introduced context as the newer
# preferred mechanism).
for src in (context, configurable):
if not isinstance(src, dict):
continue
wrapper = src.get("copilotkit_forwarded_headers")
if isinstance(wrapper, dict):
for k, v in wrapper.items():
lk = k.lower() if isinstance(k, str) else k
if isinstance(k, str) and isinstance(v, str) and lk not in headers:
headers[lk] = v
# 2) Raw x-* keys directly on context and configurable. These appear
# when langgraph-api's configurable_headers mechanism forwards inbound
# HTTP headers as individual configurable entries.
for src in (context, configurable):
if not isinstance(src, dict):
continue
for k, v in src.items():
if (
isinstance(k, str)
and k.lower().startswith("x-")
and isinstance(v, str)
):
# Don't overwrite wrapper-dict values (wrapper > raw).
# Lowercase at insertion so precedence checks are
# deterministic regardless of source casing.
lk = k.lower()
if lk not in headers:
headers[lk] = v
# Always set the ContextVar — even with an empty dict — so stale
# headers from previous calls in the same async context do not leak
# into this one.
set_forwarded_headers(headers)
except Exception as e:
# Header forwarding is best-effort. Never block the LLM call.
# Clear the ContextVar so stale headers from a prior request do not
# leak through on failure.
set_forwarded_headers({})
import logging
logging.getLogger(__name__).debug(
"Header forwarding extraction failed; continuing without forwarded headers: %s",
e,
)
def _ensure_httpx_hook(model: Any) -> None:
"""Install the header-propagation httpx hook on a LangChain chat model's
underlying HTTP client(s), if present. No-op for models that don't expose
an httpx transport (e.g. non-OpenAI/Anthropic providers).
"""
for attr in ("client", "async_client"):
client = getattr(model, attr, None)
if client is None:
continue
cid = id(client)
if cid not in _hooked_clients:
install_httpx_hook(client)
_hooked_clients.add(cid)
class StateSchema(AgentState):
copilotkit: CopilotKitProperties
StateSchema.__annotations__["ag-ui"] = CopilotKitProperties
# Internal/framework keys that should never be surfaced to the LLM as
# user-facing state. These are either reducer-managed message buckets,
# CopilotKit/AG-UI plumbing, or graph-internal scaffolding.
_RESERVED_STATE_KEYS = frozenset(
{
"messages",
"copilotkit",
# Transport-layer plumbing: forwarded request headers conveyed via a
# separate ContextVar to the httpx hook. MUST never be rendered into
# the LLM prompt — neither via App Context nor via expose_state.
"copilotkit_forwarded_headers",
"ag-ui",
"tools",
"structured_response",
"thread_id",
"remaining_steps",
}
)
class CopilotKitMiddleware(AgentMiddleware[StateSchema, Any]):
"""CopilotKit Middleware for LangGraph agents.
Handles frontend tool injection, interception for CopilotKit, and
automatic exposure of agent state to the LLM so values written via
``agent.setState`` on the frontend (or via ``Command(update=...)`` in a
tool) are visible in the next model call without needing a custom
``get_state`` tool.
Args:
expose_state: Controls how user-defined state keys are surfaced into
``request.system_message`` on every model call. Off by default
to avoid leaking arbitrary state into prompts; opt in explicitly.
- ``False`` (default) — never surface state.
- ``True`` — every state key that is not in the reserved
internal set and does not start with an underscore is
JSON-serialized into a "Current agent state:" note appended
to the system message.
- ``list``/``tuple``/``set[str]`` — only surface the named keys.
Use this when you want explicit control over what the LLM
sees (e.g. ``["liked", "todos"]``).
a2ui_params: Optional host overrides for the auto-injected
``generate_a2ui`` tool, forwarded to ``get_a2ui_tools`` when A2UI
injection fires. An ``A2UIToolParams``-shaped dict: ``guidelines``
(``generation_guidelines`` / ``design_guidelines`` /
``composition_guide``), ``default_catalog_id``,
``default_surface_id``, ``tool_name``, ``recovery``, etc. Lets a
host steer the subagent (e.g. override the default design
guidelines to favor a repeating-card layout) on the auto-inject
path, which otherwise only ever uses the toolkit defaults.
The middleware always injects ``model`` from the bound request
model (the host cannot supply the live, header-hooked model), and
folds the registered catalog id + component schema into the params
unless the host already set them — so host values win.
"""
state_schema = StateSchema
tools: ClassVar[list] = []
def __init__(
self,
*,
expose_state: Union[bool, Iterable[str]] = False,
a2ui_params: "Optional[A2UIToolParams]" = None,
):
super().__init__()
if isinstance(expose_state, bool):
self._expose_state: Union[bool, frozenset[str]] = expose_state
else:
self._expose_state = frozenset(expose_state)
# Host-supplied A2UI tool overrides (guidelines, catalog id, tool name,
# recovery, ...). Copied so later mutation of the caller's dict can't
# bleed into the middleware. ``model`` + the registered catalog are
# layered in at build time; everything here is host-owned and wins.
self._a2ui_params: dict = dict(a2ui_params or {})
@property
def name(self) -> str:
return "CopilotKitMiddleware"
# ------------------------------------------------------------------
# State-to-prompt surfacing
# ------------------------------------------------------------------
def _build_state_note(self, state: dict) -> str | None:
"""Serialize a snapshot of user state into a system-prompt note.
Returns ``None`` when nothing should be appended (feature disabled
or no non-empty user keys present).
"""
if self._expose_state is False:
return None
if isinstance(self._expose_state, frozenset):
# Allowlist branch: honor user intent for other reserved keys
# (e.g. ``thread_id``) so the override test in this suite still
# passes, but hard-exclude ``copilotkit_forwarded_headers`` —
# rendering it would leak the raw forwarded request headers into
# the LLM prompt, which is what the reserved-keys comment above
# promises will never happen "via App Context nor via expose_state".
keys: list[str] = [
k
for k in self._expose_state
if k in state and k != "copilotkit_forwarded_headers"
]
else:
keys = [
k
for k in state
if k not in _RESERVED_STATE_KEYS and not str(k).startswith("_")
]
snapshot: dict[str, Any] = {}
for k in keys:
v = state.get(k)
# Skip empty / no-op values to keep the note tight.
if v in (None, "", [], {}):
continue
snapshot[k] = v
if not snapshot:
return None
try:
body = json.dumps(snapshot, default=str, ensure_ascii=False, indent=2)
except (TypeError, ValueError):
body = str(snapshot)
return f"Current agent state:\n{body}"
def _apply_state_note(self, request: ModelRequest) -> ModelRequest:
note = self._build_state_note(request.state or {})
if not note:
return request
existing = request.system_message
if existing is None:
return request.override(system_message=SystemMessage(content=note))
base = (
existing.content
if isinstance(existing.content, str)
else str(existing.content)
)
return request.override(
system_message=SystemMessage(content=f"{base}\n\n{note}")
)
# ------------------------------------------------------------------
# Auto-A2UI tool injection
# ------------------------------------------------------------------
@staticmethod
def _resolve_a2ui_catalog(state: dict) -> "tuple[str | None, str | None] | None":
"""Find the frontend-registered A2UI catalog wherever it was passed.
Returns ``(component_schema, catalog_id)`` when a catalog is present,
else ``None`` (so the tool is never advertised when the client can't
render A2UI). Two delivery paths are supported, because the catalog
lands in different places depending on how the agent is served:
- **AG-UI native endpoint** → ``state["ag-ui"]["a2ui_schema"]``, a JSON
string ``{"catalogId": ..., "components": [...]}``.
- **CopilotKit runtime proxy** → a ``state["copilotkit"]["context"]``
entry describing the A2UI catalog (catalog id + component schemas as
text).
``component_schema`` is the text/JSON the subagent should compose from;
``catalog_id`` binds generated surfaces to the frontend's catalog (so
BYOC custom catalogs render their own components, not the basic one).
"""
# AG-UI native path.
ag_ui = state.get("ag-ui") or {}
a2ui_schema = ag_ui.get("a2ui_schema")
if a2ui_schema:
catalog_id = None
try:
parsed = (
json.loads(a2ui_schema)
if isinstance(a2ui_schema, str)
else a2ui_schema
)
if isinstance(parsed, dict):
catalog_id = parsed.get("catalogId")
except (TypeError, ValueError):
pass
# Native path: the toolkit reads ``a2ui_schema`` from state itself,
# so no composition_guide is needed — just surface the catalog id.
return None, catalog_id
# CopilotKit runtime-proxy path: the catalog arrives as a context entry.
context = (state.get("copilotkit") or {}).get("context") or []
for entry in context:
if not isinstance(entry, dict):
continue
description = entry.get("description") or ""
value = entry.get("value") or ""
if "A2UI catalog" not in description or not value:
continue
# The value lists catalogs as "- <catalogId>" lines; the first is
# the custom catalog the client registered.
match = re.search(r"(?m)^\s*-\s+(\S+)", value)
catalog_id = match.group(1) if match else None
return value, catalog_id
return None
@staticmethod
def _a2ui_inject_decision(state: dict) -> "bool | str | None":
"""Return the A2UI ``injectA2UITool`` decision, or ``None``.
The ``@ag-ui/a2ui-middleware`` forwards its ``injectA2UITool`` setting on
``forwardedProps``, which ``ag-ui-langgraph`` surfaces into agent state at
``state["ag-ui"]["inject_a2ui_tool"]`` — present only when the host turned
the runtime A2UI tool on (truthy or a custom tool-name string). ``None``
means no signal at all (off, or no A2UI middleware in the pipeline), in
which case we do not auto-inject.
"""
return (state.get("ag-ui") or {}).get("inject_a2ui_tool")
def _maybe_build_a2ui_tool(self, request: ModelRequest) -> Any | None:
"""Build a ``generate_a2ui`` tool bound to the agent's own model when
A2UI tool injection is turned on for this run.
Gating, in order:
1. **Opt-in.** Only inject when the A2UI ``injectA2UITool`` flag is
truthy (forwarded by ``@ag-ui/a2ui-middleware`` and surfaced at
``state["ag-ui"]["inject_a2ui_tool"]``). No flag → no injection. This
is the whole contract: "no injectA2UITool, no A2UI tool injection."
2. **No double-inject.** If the agent already exposes a tool with the
same name (e.g. a backend-defined ``generate_a2ui``), don't inject —
the host owns it, and a duplicate would show the model two tools with
one name.
The model is inferred from ``request.model`` (the bound agent model); the
component schema and catalog id come from the registered catalog (when
present) so the subagent composes the right components and surfaces bind
to the frontend's catalog — otherwise the toolkit's basic catalog is
used. The built tool is stashed for the tool-call hook to execute.
Returns the tool or ``None`` when A2UI is not applicable.
"""
if get_a2ui_tools is None:
return None
state = request.state or {}
# (1) Opt-in: only inject when the host turned the A2UI tool on.
if not self._a2ui_inject_decision(state):
return None
# Bind to the frontend's catalog when one was registered (optional).
resolved = self._resolve_a2ui_catalog(state)
component_schema, catalog_id = resolved if resolved else (None, None)
# Shared A2UIToolParams: a single params object owned by the toolkit.
# Start from the host overrides (guidelines / catalog id / tool name /
# recovery) so a host can steer the subagent, then layer in only what
# the host cannot know — the bound model, and the registered catalog id
# + component schema — without clobbering any host-set value.
params: "A2UIToolParams" = dict(self._a2ui_params)
params["model"] = request.model
if catalog_id and "default_catalog_id" not in params:
params["default_catalog_id"] = catalog_id
# Feed the registered component schema to the subagent so it composes
# only catalog components (the toolkit appends this to its prompt).
# Merge into any host ``guidelines`` bag; a host-set composition_guide
# wins, and host generation/design overrides are preserved.
if component_schema:
guidelines = dict(params.get("guidelines") or {})
guidelines.setdefault("composition_guide", component_schema)
params["guidelines"] = guidelines
tool = get_a2ui_tools(params)
# (2) Don't double-inject if the agent already defines this tool.
existing_names = {getattr(t, "name", None) for t in (request.tools or [])}
if tool.name in existing_names:
return None
_a2ui_tools_by_thread[_current_thread_id() or _DEFAULT_THREAD_KEY] = tool
return tool
# Inject frontend + A2UI tools and surface user state before model call
def wrap_model_call(
self,
request: ModelRequest,
handler: Callable[[ModelRequest], ModelResponse],
) -> ModelResponse:
_extract_forwarded_headers_from_config()
_ensure_httpx_hook(request.model)
request = self._apply_state_note(request)
a2ui_tool = self._maybe_build_a2ui_tool(request)
frontend_tools = request.state.get("copilotkit", {}).get("actions", [])
if a2ui_tool is not None:
# Our generate_a2ui replaces the runtime's render tool — don't
# advertise both. Drop the render tool the A2UI middleware injected.
decision = self._a2ui_inject_decision(request.state or {})
drop = decision if isinstance(decision, str) else "render_a2ui"
frontend_tools = [
t
for t in frontend_tools
if ((t.get("function") or {}).get("name") or t.get("name")) != drop
]
if not frontend_tools and a2ui_tool is None:
return handler(request)
extra_tools = [a2ui_tool] if a2ui_tool is not None else []
merged_tools = [*request.tools, *extra_tools, *frontend_tools]
return handler(request.override(tools=merged_tools))
@staticmethod
def _fix_messages_for_bedrock(messages: list) -> list:
"""Fix messages loaded from checkpoint before sending to Bedrock.
Handles four issues caused by CopilotKit's after_agent restoring
frontend tool_calls to the checkpoint:
1. Strip unanswered tool_calls (no matching ToolMessage) — Bedrock
rejects toolUse without a corresponding toolResult.
2. Sync msg.content tool_use blocks with msg.tool_calls.
3. Fix tool_use content blocks with string input (must be dict).
4. Deduplicate ToolMessages by tool_call_id — patch_orphan_tool_calls
injects a placeholder with a new random ID on every checkpoint load;
when the real result is later appended alongside it, Bedrock rejects
the duplicate toolResult IDs. We keep the real result (non-interrupted)
over the placeholder, falling back to the last occurrence if both look
real.
"""
# 4. Deduplicate ToolMessages by tool_call_id before all other processing.
# patch_orphan_tool_calls adds "…was interrupted before completion."
# placeholders with fresh random IDs on every checkpoint load. The real
# result comes in as a separate message with a different ID, so both end
# up in the list. Keep the real (non-interrupted) one; if multiple real
# ones exist, keep the last.
_INTERRUPTED_PAT = re.compile(
r"^Tool call '.+' with id '.+' was interrupted before completion\.$"
)
# Group ToolMessages by tool_call_id, preserving position
tc_groups: dict[str, list] = {}
for i, msg in enumerate(messages):
if isinstance(msg, ToolMessage):
tc_id = getattr(msg, "tool_call_id", None)
if tc_id:
tc_groups.setdefault(tc_id, []).append(i)
drop_indices: set = set()
for tc_id, indices in tc_groups.items():
if len(indices) <= 1:
continue
# Separate interrupted placeholders from real results
real_indices = [
i
for i in indices
if not (
isinstance(messages[i].content, str)
and _INTERRUPTED_PAT.match(messages[i].content)
)
]
interrupted_indices = [i for i in indices if i not in real_indices]
if real_indices and interrupted_indices:
# Replace the first placeholder (correct position, adjacent to AI
# message) with the last real result (likely appended at the end).
# This keeps the tool result in the right position for Bedrock.
messages[interrupted_indices[0]] = messages[real_indices[-1]]
drop_indices.update(interrupted_indices[1:])
drop_indices.update(real_indices) # drop all originals (we moved one)
elif real_indices:
# No placeholders, multiple real — keep only the last
drop_indices.update(real_indices[:-1])
else:
# All interrupted — keep only the last
drop_indices.update(interrupted_indices[:-1])
if drop_indices:
messages[:] = [
msg for i, msg in enumerate(messages) if i not in drop_indices
]
for idx, msg in enumerate(messages):
if not isinstance(msg, AIMessage):
continue
tool_calls = getattr(msg, "tool_calls", None) or []
# 1. Sync content with tool_calls: remove tool_use content blocks
# that aren't in msg.tool_calls (e.g. stripped by after_model
# but content blocks left behind in checkpoint).
if tool_calls and isinstance(msg.content, list):
tc_ids = {tc.get("id") for tc in tool_calls}
msg.content = [
block
for block in msg.content
if not (
isinstance(block, dict)
and block.get("type") == "tool_use"
and block.get("id") not in tc_ids
)
]
elif not tool_calls and isinstance(msg.content, list):
# No tool_calls at all — strip ALL tool_use content blocks
msg.content = [
block
for block in msg.content
if not (isinstance(block, dict) and block.get("type") == "tool_use")
]
if not tool_calls:
continue
# 2. Strip unanswered tool_calls — only consider ToolMessages that
# are ADJACENT (immediately following this AIMessage, before the
# next non-ToolMessage). A ToolMessage at the wrong position
# won't satisfy Bedrock's Converse API requirement that toolResult
# blocks appear in the user turn right after the assistant turn.
adjacent_tc_ids: set = set()
j = idx + 1
while j < len(messages) and isinstance(messages[j], ToolMessage):
tc_id = getattr(messages[j], "tool_call_id", None)
if tc_id:
adjacent_tc_ids.add(tc_id)
j += 1
unanswered = [
tc for tc in tool_calls if tc.get("id") not in adjacent_tc_ids
]
if unanswered:
unanswered_ids = {tc["id"] for tc in unanswered}
msg.tool_calls = [
tc for tc in tool_calls if tc.get("id") in adjacent_tc_ids
]
# Also strip matching content blocks
if isinstance(msg.content, list):
msg.content = [
block
for block in msg.content
if not (
isinstance(block, dict)
and block.get("type") == "tool_use"
and block.get("id") in unanswered_ids
)
]
# 3. Fix string args in tool_calls
for tc in msg.tool_calls or []:
if isinstance(tc.get("args"), str):
try:
tc["args"] = json.loads(tc["args"])
except (json.JSONDecodeError, TypeError):
tc["args"] = {}
# 4. Fix string input in content blocks
if isinstance(msg.content, list):
for block in msg.content:
if isinstance(block, dict) and block.get("type") == "tool_use":
inp = block.get("input")
if isinstance(inp, str):
try:
block["input"] = json.loads(inp) if inp else {}
except (json.JSONDecodeError, TypeError):
block["input"] = {}
elif inp is None:
block["input"] = {}
# 5. Remove orphan ToolMessages whose tool_call_id no longer matches
# any remaining tool_call in any AIMessage. These can be left over
# after stripping unanswered tool_calls above.
remaining_tc_ids: set = set()
for msg in messages:
if isinstance(msg, AIMessage):
for tc in getattr(msg, "tool_calls", None) or []:
tc_id = tc.get("id")
if tc_id:
remaining_tc_ids.add(tc_id)
messages[:] = [
msg
for msg in messages
if not isinstance(msg, ToolMessage)
or getattr(msg, "tool_call_id", None) in remaining_tc_ids
]
return messages
async def awrap_model_call(
self,
request: ModelRequest,
handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
) -> ModelResponse:
_extract_forwarded_headers_from_config()
_ensure_httpx_hook(request.model)
self._fix_messages_for_bedrock(request.messages)
request = self._apply_state_note(request)
a2ui_tool = self._maybe_build_a2ui_tool(request)
frontend_tools = request.state.get("copilotkit", {}).get("actions", [])
if a2ui_tool is not None:
# Our generate_a2ui replaces the runtime's render tool — don't
# advertise both. Drop the render tool the A2UI middleware injected.
decision = self._a2ui_inject_decision(request.state or {})
drop = decision if isinstance(decision, str) else "render_a2ui"
frontend_tools = [
t
for t in frontend_tools
if ((t.get("function") or {}).get("name") or t.get("name")) != drop
]
if not frontend_tools and a2ui_tool is None:
return await handler(request)
extra_tools = [a2ui_tool] if a2ui_tool is not None else []
merged_tools = [*request.tools, *extra_tools, *frontend_tools]
return await handler(request.override(tools=merged_tools))
# ------------------------------------------------------------------
# Auto-A2UI tool execution
# ------------------------------------------------------------------
# The generate_a2ui tool is advertised dynamically in wrap_model_call and is
# NOT in create_agent's static tool registry, so the tool node cannot
# execute it on its own. These hooks supply the implementation (built with
# the inferred model) for that one tool; their presence also disables
# create_agent's "unknown tool" guard for dynamically-advertised tools.
def _resolve_a2ui_request(self, request: Any) -> Any:
"""Return a request overridden with the stashed A2UI tool when this
tool call targets it, else the original request unchanged."""
tool = _a2ui_tools_by_thread.get(_current_thread_id() or _DEFAULT_THREAD_KEY)
if (
tool is not None
and getattr(request, "tool", None) is None
and request.tool_call.get("name") == tool.name
):
return request.override(tool=tool)
return request
def wrap_tool_call(
self,
request: Any,
handler: Callable[[Any], Any],
) -> Any:
return handler(self._resolve_a2ui_request(request))
async def awrap_tool_call(
self,
request: Any,
handler: Callable[[Any], Awaitable[Any]],
) -> Any:
return await handler(self._resolve_a2ui_request(request))
# Inject app context before agent runs
def before_agent(
self,
state: StateSchema,
runtime: Runtime[Any],
) -> dict[str, Any] | None:
messages = state.get("messages", [])
if not messages:
return None
# Get app context from state or runtime
copilotkit_state = state.get("copilotkit", {})
app_context = copilotkit_state.get("context") or getattr(
runtime, "context", None
)
# Strip the reserved transport-layer key ``copilotkit_forwarded_headers``
# so it is never rendered into the LLM prompt. langgraph-api auto-copies
# ``config.configurable`` into ``runtime.context``, which means the
# forwarded-headers wrapper dict shows up here even though it is only
# meant for the httpx hook (which reads it from a separate ContextVar
# via ``_extract_forwarded_headers_from_config``).
if isinstance(app_context, dict):
app_context = {
k: v
for k, v in app_context.items()
if k != "copilotkit_forwarded_headers"
}
# Check if app_context is missing or empty
if not app_context:
return None
if isinstance(app_context, str) and app_context.strip() == "":
return None
if isinstance(app_context, dict) and len(app_context) == 0:
return None
# Create the context content
if isinstance(app_context, str):
context_content = app_context
else:
# Handle Pydantic models (e.g. ag_ui Context)
if hasattr(app_context, "model_dump"):
app_context = app_context.model_dump()
elif isinstance(app_context, list):
app_context = [
item.model_dump() if hasattr(item, "model_dump") else item
for item in app_context
]
context_content = json.dumps(app_context, indent=2)
context_message_content = f"App Context:\n{context_content}"
context_message_prefix = "App Context:\n"
# Helper to get message content as string
def get_content_string(msg: Any) -> str | None:
content = getattr(msg, "content", None)
if isinstance(content, str):
return content
if isinstance(content, list) and content and isinstance(content[0], dict):
return content[0].get("text")
return None
# Find the first system/developer message (not our context message)
# to determine where to insert our context message (right after it)
first_system_index = -1
for i, msg in enumerate(messages):
msg_type = getattr(msg, "type", None)
if msg_type in ("system", "developer"):
content = get_content_string(msg)
# Skip if this is our own context message
if content and content.startswith(context_message_prefix):
continue
first_system_index = i
break
# Check if our context message already exists
existing_context_index = -1
for i, msg in enumerate(messages):
msg_type = getattr(msg, "type", None)
if msg_type in ("system", "developer"):
content = get_content_string(msg)
if content and content.startswith(context_message_prefix):
existing_context_index = i
break
# Create the context message.
# When replacing an existing context message, reuse its ID so the
# add_messages reducer updates in-place instead of appending a
# duplicate at the end of the message list.
if existing_context_index != -1:
existing_id = getattr(messages[existing_context_index], "id", None)
context_message = SystemMessage(
content=context_message_content, id=existing_id
)
else:
context_message = SystemMessage(content=context_message_content)
if existing_context_index != -1:
# Replace existing context message
updated_messages = list(messages)
updated_messages[existing_context_index] = context_message
else:
# Insert after the first system message, or at position 0 if no system message
insert_index = first_system_index + 1 if first_system_index != -1 else 0
updated_messages = [
*messages[:insert_index],
context_message,
*messages[insert_index:],
]
return {
**state,
"messages": updated_messages,
}
async def abefore_agent(
self,
state: StateSchema,
runtime: Runtime[Any],
) -> dict[str, Any] | None:
# Delegate to sync implementation
return self.before_agent(state, runtime)
# Intercept frontend tool calls after model returns, before ToolNode executes
def after_model(
self,
state: StateSchema,
runtime: Runtime[Any],
) -> dict[str, Any] | None:
frontend_tools = state.get("copilotkit", {}).get("actions", [])
if not frontend_tools:
return None
frontend_tool_names = {
t.get("function", {}).get("name") or t.get("name") for t in frontend_tools
}
# Find last AI message with tool calls
messages = state.get("messages", [])
if not messages:
return None
last_message = messages[-1]
if not isinstance(last_message, AIMessage):
return None
tool_calls = getattr(last_message, "tool_calls", None) or []
if not tool_calls:
return None
backend_tool_calls = []
frontend_tool_calls = []
for call in tool_calls:
if call.get("name") in frontend_tool_names:
frontend_tool_calls.append(call)
else:
backend_tool_calls.append(call)
if not frontend_tool_calls:
return None
# Create updated AIMessage with only backend tool calls
updated_ai_message = AIMessage(
content=last_message.content,
tool_calls=backend_tool_calls,
id=last_message.id,
)
return {
"messages": [*messages[:-1], updated_ai_message],
"copilotkit": {
"intercepted_tool_calls": frontend_tool_calls,
"original_ai_message_id": last_message.id,
},
}
async def aafter_model(
self,
state: StateSchema,
runtime: Runtime[Any],
) -> dict[str, Any] | None:
# Delegate to sync implementation
return self.after_model(state, runtime)
# Restore frontend tool calls to AIMessage before agent exits
def after_agent(
self,
state: StateSchema,
runtime: Runtime[Any],
) -> dict[str, Any] | None:
# Drop the bridged A2UI tool for this run — all tool calls for the turn
# have executed by now; the next model call re-stashes if needed.
_a2ui_tools_by_thread.pop(_current_thread_id() or _DEFAULT_THREAD_KEY, None)
copilotkit_state = state.get("copilotkit", {})
intercepted_tool_calls = copilotkit_state.get("intercepted_tool_calls")
original_message_id = copilotkit_state.get("original_ai_message_id")
if not intercepted_tool_calls or not original_message_id:
return None
messages = state.get("messages", [])
updated_messages = []
for msg in messages:
if isinstance(msg, AIMessage) and msg.id == original_message_id:
existing_tool_calls = getattr(msg, "tool_calls", None) or []
updated_messages.append(
AIMessage(
content=msg.content,
tool_calls=[*existing_tool_calls, *intercepted_tool_calls],
id=msg.id,
)
)
else:
updated_messages.append(msg)
return {
"messages": updated_messages,
"copilotkit": {
"intercepted_tool_calls": None,
"original_ai_message_id": None,
},
}
async def aafter_agent(
self,
state: StateSchema,
runtime: Runtime[Any],
) -> dict[str, Any] | None:
# Delegate to sync implementation
return self.after_agent(state, runtime)
+45
View File
@@ -0,0 +1,45 @@
"""
CrewAI
"""
from .crewai_agent import CrewAIAgent
from .crewai_sdk import (
CopilotKitProperties,
CopilotKitState,
copilotkit_emit_state,
copilotkit_emit_message,
copilotkit_emit_tool_call,
copilotkit_stream,
copilotkit_exit,
copilotkit_predict_state,
)
from .copilotkit_integration import (
CopilotKitFlow,
CopilotKitToolCallEvent,
register_tool_call_listener,
tool_calls_log,
create_tool_proxy,
FlowInputState,
CopilotKitStateUpdateEvent,
emit_copilotkit_state_update_event,
)
__all__ = [
"CrewAIAgent",
"CopilotKitProperties",
"CopilotKitState",
"copilotkit_emit_state",
"copilotkit_emit_message",
"copilotkit_emit_tool_call",
"copilotkit_stream",
"copilotkit_exit",
"copilotkit_predict_state",
"CopilotKitFlow",
"CopilotKitToolCallEvent",
"register_tool_call_listener",
"tool_calls_log",
"create_tool_proxy",
"FlowInputState",
"CopilotKitStateUpdateEvent",
"emit_copilotkit_state_update_event",
]
@@ -0,0 +1,356 @@
#!/usr/bin/env python
from typing import Dict, Any, List, Optional, Generic
import datetime
from crewai.flow import Flow
from crewai import LLM
from crewai.utilities.events import crewai_event_bus
import logging
from crewai.utilities.events.base_events import BaseEvent
from pydantic import Field
from typing import TypeVar
from pydantic import BaseModel
# Define a generic type variable for the state
S = TypeVar("S")
logger = logging.getLogger(__name__)
# Tool calls log for tracking
tool_calls_log = []
class FlowInputState(BaseModel):
"""Defines the expected input state for the AgenticChatFlow."""
messages: List[Dict[str, str]] = [] # Current message(s) from the user
tools: List[
Dict[str, Any]
] = [] # CopilotKit tool format: name, description, parameters
conversation_history: List[
Dict[str, str]
] = [] # Full conversation history (persisted between runs)
class CopilotKitToolCallEvent(BaseEvent):
"""Event emitted when a tool call is made through CopilotKit"""
type: str = "copilotkit_frontend_tool_call"
tool_name: str
args: Dict[str, Any]
timestamp: str = Field(default_factory=lambda: datetime.datetime.now().isoformat())
def __init__(self, **data):
# If timestamp is not provided, it will use the default_factory
super().__init__(**data)
class CopilotKitStateUpdateEvent(BaseEvent):
"""Event for state updates in CopilotKit"""
type: str = "copilotkit_state_update"
tool_name: str
args: dict[str, Any]
timestamp: str = Field(default_factory=lambda: datetime.datetime.now().isoformat())
def __init__(self, **data):
# If timestamp is not provided, it will use the default_factory
super().__init__(**data)
def create_tool_proxy(tool_name):
def tool_proxy(**kwargs):
event = CopilotKitToolCallEvent(tool_name=tool_name, args=kwargs)
tool_calls_log.append(
{"tool_name": tool_name, "args": kwargs, "timestamp": event.timestamp}
)
assert hasattr(crewai_event_bus, "emit")
logger.info(
f"create_tool_proxy: Emitting tool call event for {tool_name} with parameters: {kwargs}"
)
crewai_event_bus.emit(None, event=event)
return f"\n\nTool {tool_name} called successfully with parameters: {kwargs}\n\n"
return tool_proxy
class CopilotKitFlow(Flow[S], Generic[S]): # Make it generic
_tools_from_input: List[Dict[str, Any]] = [] # Store raw tool definitions
def __class_getitem__(cls, item):
# Pass type info down to Flow's __class_getitem__
super().__class_getitem__(item)
cls._initial_state_T = item
return cls
def kickoff(
self, state: Optional[S] = None, inputs: Optional[Dict[str, Any]] = None
):
# CrewAI's Flow class initializes self.state from the 'state' parameter or
# by instantiating S using 'inputs' if 'state' is None and 'inputs' is a dict.
# We need to ensure tools from 'inputs' (if any) are captured if not part of S's direct fields
# or if S is initialized before this kickoff by CrewAI.
# If inputs dict contains 'tools', store them for get_available_tools
if isinstance(inputs, dict) and "tools" in inputs:
# Be careful with class-level _tools_from_input if multiple instances run concurrently
# It might be better to store this on self.
CopilotKitFlow._tools_from_input = inputs.get("tools", [])
print(f"Tools from inputs dict: {CopilotKitFlow._tools_from_input}")
# The actual_input for super().kickoff should be the state model instance S
# or the dict 'inputs' if state is None.
# The base Flow's kickoff will handle initializing self.state.
# If state is already an instance of S, pass it.
# If state is None and inputs is a dict, Flow.__init__ will use inputs to create S.
# Let the base Flow handle state initialization.
# Our main job here is to potentially intercept 'inputs' if it has a structure
# not directly mapping to S (e.g., tools in a separate key).
# However, with AgentInputState having 'tools', this should be cleaner.
# Call parent's kickoff - note that base Flow.kickoff() only accepts 'inputs'
# If state is not None, we should convert it to dict and use as inputs
if state is not None and inputs is None:
# If we have a state model instance but no inputs, convert state to dict for inputs
if hasattr(state, "dict") and callable(getattr(state, "dict")):
inputs_dict = state.dict()
result = super().kickoff(inputs=inputs_dict)
else:
# If state can't be converted via .dict(), use it directly as inputs
result = super().kickoff(inputs=state)
else:
# Normal case: just pass inputs (which might be None)
result = super().kickoff(inputs=inputs)
return result # Return what the base Flow.kickoff returns
def get_message_history(
self, system_prompt: Optional[str] = None, max_messages: int = 20
) -> List[Dict[str, str]]:
messages: List[Dict[str, str]] = []
# PRIORITIZE conversation_history if available (for persistence between runs)
if (
hasattr(self.state, "conversation_history")
and isinstance(self.state.conversation_history, list)
and self.state.conversation_history
):
# If we have conversation history, use it as the primary source of messages
messages.extend(self.state.conversation_history)
logger.info(
f"get_message_history: Loaded {len(self.state.conversation_history)} messages from conversation history"
)
# If there are new messages not in the history, add them temporarily (they'll be saved to history later)
if hasattr(self.state, "messages") and isinstance(
self.state.messages, list
):
for msg in self.state.messages:
if msg not in messages:
messages.append(msg)
logger.info(
f"get_message_history: Added new message (not yet in history): {msg.get('content', '')[:30]}..."
)
# If no conversation history, try current messages
elif hasattr(self.state, "messages") and isinstance(self.state.messages, list):
messages.extend(self.state.messages)
print(
f"get_message_history: Loaded {len(self.state.messages)} messages from current messages"
)
# Fallback for raw input if state isn't populated as expected (less ideal)
elif (
hasattr(self, "_raw_input")
and isinstance(self._raw_input, dict)
and "messages" in self._raw_input
):
messages.extend(self._raw_input["messages"])
logger.info(
f"get_message_history: Loaded {len(self._raw_input['messages'])} messages from _raw_input"
)
# Add system prompt if needed
if system_prompt:
# Check if we already have a system message
has_system_message = any(msg.get("role") == "system" for msg in messages)
if not has_system_message:
# Add system message at the beginning
messages.insert(0, {"role": "system", "content": system_prompt})
logger.info(f"get_message_history: Added system prompt message")
# Limit to max_messages, but keep the system message if present
if len(messages) > max_messages:
# If first message is system message, keep it and take the (max_messages-1) most recent messages
if messages and messages[0].get("role") == "system":
system_msg = messages[0]
recent_msgs = messages[-(max_messages - 1) :]
messages = [system_msg] + recent_msgs
logger.info(
f"get_message_history: Truncated to {len(messages)} messages (including system message)"
)
else:
# Otherwise just take most recent messages
messages = messages[-max_messages:]
logger.info(
f"get_message_history: Truncated to {len(messages)} most recent messages"
)
return messages
def get_available_tools(self) -> List[Dict[str, Any]]:
raw_tools: List[Dict[str, Any]] = []
# Primary source: self.state.tools (from AgentInputState)
if hasattr(self.state, "tools") and isinstance(self.state.tools, list):
raw_tools = self.state.tools
logger.info(
f"get_available_tools: Loaded {len(raw_tools)} tools from self.state.tools"
)
# Fallback to _tools_from_input (populated in kickoff from raw 'inputs' dict)
# This is useful if 'tools' was passed separately and not as part of the state model S.
elif CopilotKitFlow._tools_from_input:
raw_tools = CopilotKitFlow._tools_from_input
logger.info(
f"get_available_tools: Loaded {len(raw_tools)} tools from _tools_from_input"
)
# Fallback for raw input (less ideal)
elif (
hasattr(self, "_raw_input")
and isinstance(self._raw_input, dict)
and "tools" in self._raw_input
):
raw_tools = self._raw_input["tools"]
logger.info(
f"get_available_tools: Loaded {len(raw_tools)} tools from _raw_input"
)
return raw_tools
def format_tools_for_llm(
self, tools_definitions: List[Dict[str, Any]]
) -> tuple[List[Dict[str, Any]], Dict[str, callable]]:
formatted_tools = []
available_functions = {}
logger.info(
f"format_tools_for_llm: Processing {len(tools_definitions)} tool definitions."
)
for tool_def in tools_definitions:
if (
"name" in tool_def
and "parameters" in tool_def
and "description" in tool_def
):
# Standard OpenAI tool format
formatted_tool = {
"type": "function",
"function": {
"name": tool_def["name"],
"description": tool_def["description"],
"parameters": tool_def["parameters"],
},
}
formatted_tools.append(formatted_tool)
# Create and store the proxy function
tool_name = tool_def["name"]
available_functions[tool_name] = create_tool_proxy(tool_name)
logger.info(
f"format_tools_for_llm: Created proxy for tool: {tool_name}"
)
else:
logger.info(
f"format_tools_for_llm: Skipped invalid tool definition: {tool_def.get('name', 'N/A')}"
)
return formatted_tools, available_functions
def handle_tool_responses(
self,
llm: LLM,
response_text: str, # Changed from 'response' to 'response_text' for clarity
messages: List[Dict[str, str]],
tools_called_count_before_llm_call: int, # More descriptive name
follow_up_prompt: Optional[str] = None,
) -> str:
new_tools_called_during_interaction = (
len(tool_calls_log) > tools_called_count_before_llm_call
)
# Check if a follow-up is needed (tools were called but no substantive natural language content)
need_followup = new_tools_called_during_interaction and (
not response_text.strip()
or all(
f"Tool {call['tool_name']}" in response_text
for call in tool_calls_log[tools_called_count_before_llm_call:]
)
)
if need_followup:
logger.info("handle_tool_responses: Follow-up needed after tool call.")
follow_up_messages = messages.copy()
# Add the assistant's response that included tool calls (or was just tool call confirmations)
follow_up_messages.append({"role": "assistant", "content": response_text})
# Add tool call results as messages (CopilotKit might do this differently, adjust if needed)
# For OpenAI, tool results are typically added with role 'tool'
# This part might need alignment with how CopilotKit expects tool results to be fed back.
# The current [create_tool_proxy](cci:1://file:///Users/croonnicola/Downloads/agentic_chat/src/agentic_chat/copilotkit_integration.py:22:0-42:21) returns a string. This string becomes the 'content'
# of the assistant's message. If the LLM needs explicit tool result messages,
# this needs adjustment. For now, we assume the proxy's string output is sufficient.
prompt_for_final_answer = (
follow_up_prompt
or "Tools have been called. Continue with your response."
)
follow_up_messages.append(
{"role": "user", "content": prompt_for_final_answer}
)
logger.info(
f"handle_tool_responses: Calling LLM for follow-up with {len(follow_up_messages)} messages."
)
# Call LLM without tools for a final natural language response
final_response_text = llm.call(
messages=follow_up_messages, tools=None, available_functions=None
)
# Combine initial tool call confirmations with the final natural language response
# This behavior might need tuning based on desired output verbosity
# combined_response = response_text + "\n\n" + final_response_text
# Often, you just want the final_response_text
return final_response_text
else:
return response_text # No follow-up needed, return original LLM response
def get_tools_summary(self) -> str: # Remains the same
summary = f"\nTotal tool calls: {len(tool_calls_log)}\n"
for i, call in enumerate(tool_calls_log):
summary += f"\n[{i + 1}] Tool: {call['tool_name']}"
summary += f"\n Args: {call['args']}"
summary += f"\n Time: {call['timestamp']}\n"
return summary
# Register event listener (remains the same)
def register_tool_call_listener():
@crewai_event_bus.on(CopilotKitToolCallEvent)
def on_tool_call_event(source, event):
print(
f"Received CopilotKit tool call event: Tool: {event.tool_name}, Args: {event.args}, Time: {event.timestamp}"
)
pass
# Use this function to emit state updates to the client UI (STATE_SNAPSHOT)
# This is particularly useful when you need to update the UI state from within a tool call
# or when you want to reflect state changes in the AG-UI interface
# Example: emit_copilotkit_state_update_event("write_document", {"document": state.data["document"]})
def emit_copilotkit_state_update_event(tool_name: str, args: dict[str, Any]):
event = CopilotKitStateUpdateEvent(tool_name=tool_name, args=args)
crewai_event_bus.emit(None, event=event)
@@ -0,0 +1,517 @@
"""
CrewAI Agent
"""
import uuid
import json
from copy import deepcopy
from typing import Optional, List, Callable
from typing_extensions import TypedDict, NotRequired, Any, Dict, cast
from pydantic import BaseModel
from crewai import Crew, Flow
from crewai.flow import start
from crewai.cli.crew_chat import (
initialize_chat_llm as crew_chat_initialize_chat_llm,
generate_crew_chat_inputs as crew_chat_generate_crew_chat_inputs,
generate_crew_tool_schema as crew_chat_generate_crew_tool_schema,
build_system_message as crew_chat_build_system_message,
create_tool_function as crew_chat_create_tool_function,
)
from litellm import completion
from copilotkit.agent import Agent
from copilotkit.types import Message
from copilotkit.action import ActionDict
from copilotkit.protocol import (
emit_runtime_events,
agent_state_message,
)
from copilotkit.crewai.crewai_sdk import (
copilotkit_messages_to_crewai_flow,
crewai_flow_messages_to_copilotkit,
crewai_flow_async_runner,
copilotkit_stream,
copilotkit_exit,
logger,
)
from copilotkit.runloop import copilotkit_run, CopilotKitRunExecution
class CopilotKitConfig(TypedDict):
"""
CopilotKit config for CrewAIAgent
This is used for advanced cases where you want to customize how CopilotKit interacts with
CrewAI.
```python
# Function signatures:
def merge_state(
*,
state: dict,
messages: List[BaseMessage],
actions: List[Any],
agent_name: str
):
# ...implementation...
```
Parameters
----------
merge_state : Callable
This function lets you customize how CopilotKit merges the agent state.
"""
merge_state: NotRequired[Callable]
class CrewAIFlowExecutionState(TypedDict):
"""
State for an execution of a CrewAI Flow agent
"""
should_exit: bool
node_name: str
is_finished: bool
predict_state_configuration: Dict[str, Any]
predicted_state: Dict[str, Any]
argument_buffer: str
current_tool_call: Optional[str]
class CrewAIAgent(Agent):
"""
CrewAIAgent lets you define your agent for use with CopilotKit.
To install, run:
```bash
pip install copilotkit[crewai]
```
Every agent must have the `name` and either `crew` or `flow` properties defined. An optional
`description` can also be provided. This is used when CopilotKit is dynamically routing requests
to the agent.
## Serving a Crew based agent
To serve a Crew based agent, pass in a `Crew` object to the `crew` parameter.
Note:
You need to make sure to have a `chat_llm` set on the `Crew` object.
See [the CrewAI docs](https://docs.crewai.com/concepts/cli#9-chat) for more information.
```python
from copilotkit import CrewAIAgent
CrewAIAgent(
name="email_agent_crew",
description="This crew based agent sends emails",
crew=SendEmailCrew(),
)
```
## Serving a Flow based agent
To serve a Flow based agent, pass in a `Flow` object to the `flow` parameter.
```python
CrewAIAgent(
name="email_agent_flow",
description="This flow based agent sends emails",
flow=SendEmailFlow(),
)
```
Note:
Either a `crew` or `flow` must be provided to CrewAIAgent.
Parameters
----------
name : str
The name of the agent.
crew : Crew
When using a Crew based agent, pass in a `Crew` object to the `crew` parameter.
flow : Flow
When using a Flow based agent, pass in a `Flow` object to the `flow` parameter.
description : Optional[str]
The description of the agent.
copilotkit_config : Optional[CopilotKitConfig]
The CopilotKit config to use with the agent.
"""
def __init__(
self,
*,
name: str,
description: Optional[str] = None,
crew: Optional[Crew] = None,
flow: Optional[Flow] = None,
copilotkit_config: Optional[CopilotKitConfig] = None,
):
super().__init__(
name=name,
description=description,
)
if (crew is None) == (flow is None):
raise ValueError("Either crew or flow must be provided to CrewAIAgent")
self.crew = crew
self.flow = flow
self.copilotkit_config = copilotkit_config or {}
def execute( # pylint: disable=too-many-arguments
self,
*,
state: dict,
thread_id: str,
messages: List[Message],
actions: Optional[List[ActionDict]] = None,
**kwargs,
):
"""Execute the agent"""
if self.crew:
crew = deepcopy(self.crew)
return self.execute_crew(
state=state,
messages=messages,
thread_id=thread_id,
actions=actions,
crew=crew,
**kwargs,
)
if self.flow:
flow = deepcopy(self.flow)
return self.execute_flow(
state=state,
messages=messages,
thread_id=thread_id,
actions=actions,
flow=flow,
**kwargs,
)
raise ValueError("Either crew or flow must be provided to CrewAIAgent")
def execute_crew( # pylint: disable=too-many-arguments,unused-argument
self,
*,
state: dict,
crew: Crew,
thread_id: str,
messages: List[Message],
actions: Optional[List[ActionDict]] = None,
**kwargs,
):
"""Execute a `Crew` based agent"""
flow = ChatWithCrewFlow(
crew=crew,
crew_name=self.name,
thread_id=thread_id,
cache_key=f"crew_{id(self.crew)}",
)
return self.execute_flow(
state=state,
messages=messages,
thread_id=thread_id,
actions=actions,
flow=flow,
**kwargs,
)
async def execute_flow( # pylint: disable=too-many-arguments,unused-argument,too-many-locals
self,
*,
state: dict,
messages: List[Message],
thread_id: Optional[str] = None,
actions: Optional[List[ActionDict]] = None,
flow: Flow,
**kwargs,
):
"""Execute a `Flow` based agent"""
if thread_id is None:
raise ValueError("Thread ID is required")
run_id = str(uuid.uuid4())
merge_state = self.copilotkit_config.get(
"merge_state", crewai_flow_default_merge_state
)
crewai_flow_messages = copilotkit_messages_to_crewai_flow(messages)
state = merge_state(
state=state,
messages=crewai_flow_messages,
actions=actions or [],
agent_name=self.name,
flow=flow,
)
execution: CopilotKitRunExecution = CopilotKitRunExecution(
thread_id=thread_id,
agent_name=self.name,
run_id=run_id,
should_exit=False,
node_name="start",
is_finished=False,
predict_state_configuration={},
predicted_state={},
argument_buffer="",
current_tool_call=None,
state=state,
)
async for event in copilotkit_run(
fn=lambda: crewai_flow_async_runner(flow, deepcopy(state)),
execution=execution,
):
yield event
state = {
**(
flow.state.model_dump()
if isinstance(flow.state, BaseModel)
else flow.state
)
}
if "messages" in state:
state["messages"] = crewai_flow_messages_to_copilotkit(state["messages"])
# emit the final state
yield emit_runtime_events(
agent_state_message(
thread_id=thread_id,
agent_name=self.name,
node_name=execution["node_name"],
run_id=run_id,
active=False,
role="assistant",
state=json.dumps(filter_state(state, exclude_keys=["id"])),
running=not execution["should_exit"],
)
)
async def get_state(
self,
*,
thread_id: str,
):
if self.flow and self.flow._persistence: # pylint: disable=protected-access
try:
stored_state = self.flow._persistence.load_state(thread_id) # pylint: disable=protected-access
messages = []
if "messages" in stored_state and stored_state["messages"]:
try:
messages = crewai_flow_messages_to_copilotkit(
stored_state["messages"]
)
except Exception as e: # pylint: disable=broad-except
# If conversion fails, we'll return empty messages
logger.warning(
f"Failed to convert messages from stored state: {str(e)}"
)
return {
"threadId": thread_id,
"threadExists": True,
"state": stored_state,
"messages": messages,
}
except Exception as e: # pylint: disable=broad-except
logger.warning(f"Failed to load state for thread {thread_id}: {str(e)}")
return {
"threadId": thread_id,
"threadExists": False,
"state": {},
"messages": [],
}
def dict_repr(self):
super_repr = super().dict_repr()
return {**super_repr, "type": "crewai"}
def crewai_flow_default_merge_state( # pylint: disable=unused-argument, too-many-arguments
*,
state: dict,
flow: Flow,
messages: List[Any],
actions: List[Any],
agent_name: str,
):
"""Default merge state for CrewAI"""
if len(messages) > 0:
if "role" in messages[0] and messages[0]["role"] == "system":
messages = messages[1:]
actions = [
{
"type": "function",
"function": {
**action,
},
}
for action in actions
]
new_state = {**state, "messages": messages, "copilotkit": {"actions": actions}}
return new_state
def filter_state(
state: Dict[str, Any], exclude_keys: Optional[List[str]] = None
) -> Dict[str, Any]:
"""Filter out messages and id from the state"""
exclude_keys = exclude_keys or ["messages", "id"]
return {k: v for k, v in state.items() if k not in exclude_keys}
CREW_EXIT_TOOL = {
"type": "function",
"function": {
"name": "crew_exit",
"description": "Call this when the user has indicated that they are done with the crew",
"parameters": {
"type": "object",
"properties": {},
"required": [],
},
},
}
_CREW_INPUTS_CACHE = {}
class ChatWithCrewFlow(Flow):
"""Chat with crew"""
def __init__(self, *, crew: Crew, crew_name: str, thread_id: str, cache_key: str):
super().__init__()
self.crew = cast(Any, crew).crew()
if self.crew.chat_llm is None:
raise ValueError("Crew chat LLM is not set")
self.crew_name = crew_name
self.thread_id = thread_id
self.chat_llm = crew_chat_initialize_chat_llm(self.crew)
if cache_key not in _CREW_INPUTS_CACHE:
self.crew_chat_inputs = crew_chat_generate_crew_chat_inputs(
self.crew, self.crew_name, self.chat_llm
)
_CREW_INPUTS_CACHE[cache_key] = self.crew_chat_inputs
else:
self.crew_chat_inputs = _CREW_INPUTS_CACHE[cache_key]
self.crew_tool_schema = crew_chat_generate_crew_tool_schema(
self.crew_chat_inputs
)
self.system_message = crew_chat_build_system_message(self.crew_chat_inputs)
super().__init__()
@start()
async def chat(self):
"""Chat with the crew"""
system_message = self.system_message
if self.state.get("inputs"):
system_message += "\n\nCurrent inputs: " + json.dumps(self.state["inputs"])
messages = [
{
"role": "system",
"content": system_message,
"id": self.thread_id + "-system",
},
*self.state["messages"],
]
tools = [
action
for action in self.state["copilotkit"]["actions"]
if action["function"]["name"] != self.crew_name
]
tools += [self.crew_tool_schema, CREW_EXIT_TOOL]
response = await copilotkit_stream(
completion(
model=self.crew.chat_llm,
messages=messages,
tools=tools,
parallel_tool_calls=False,
stream=True,
)
)
message = cast(Any, response).choices[0]["message"]
self.state["messages"].append(message)
if message.get("tool_calls"):
if message["tool_calls"][0]["function"]["name"] == self.crew_name:
# run the crew
crew_function = crew_chat_create_tool_function(self.crew, messages)
args = json.loads(message["tool_calls"][0]["function"]["arguments"])
result = crew_function(**args)
if isinstance(result, str):
self.state["outputs"] = result
elif hasattr(result, "json_dict"):
self.state["outputs"] = result.json_dict
elif hasattr(result, "raw"):
self.state["outputs"] = result.raw
else:
raise ValueError("Unexpected result type", type(result))
self.state["messages"].append(
{
"role": "tool",
"content": result,
"tool_call_id": message["tool_calls"][0]["id"],
}
)
elif (
message["tool_calls"][0]["function"]["name"]
== CREW_EXIT_TOOL["function"]["name"]
):
await copilotkit_exit()
self.state["messages"].append(
{
"role": "tool",
"content": "Crew exited",
"tool_call_id": message["tool_calls"][0]["id"],
}
)
response = await copilotkit_stream(
completion( # pylint: disable=too-many-arguments
model=self.crew.chat_llm,
messages=[
{
"role": "system",
"content": "Indicate to the user that the crew has exited",
"id": self.thread_id + "-system",
},
*self.state["messages"],
],
tools=tools,
parallel_tool_calls=False,
stream=True,
tool_choice="none",
)
)
message = cast(Any, response).choices[0]["message"]
self.state["messages"].append(message)
+678
View File
@@ -0,0 +1,678 @@
"""
CrewAI integration for CopilotKit
"""
import uuid
import json
import asyncio
from typing_extensions import Any, Dict, List, Literal, Optional
from copilotkit.exc import CopilotKitMisuseError
from pydantic import BaseModel, Field
from litellm.types.utils import (
ModelResponse,
Choices,
Message as LiteLLMMessage,
ChatCompletionMessageToolCall,
Function as LiteLLMFunction,
)
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from crewai.flow.flow import FlowState, Flow
try:
from crewai.utilities.events.flow_events import (
FlowEvent as CrewAIFlowEvent,
FlowStartedEvent,
MethodExecutionStartedEvent,
MethodExecutionFinishedEvent,
FlowFinishedEvent,
)
except ImportError:
from crewai.events.types.flow_events import ( # type: ignore[no-redef]
FlowEvent as CrewAIFlowEvent,
FlowStartedEvent,
MethodExecutionStartedEvent,
MethodExecutionFinishedEvent,
FlowFinishedEvent,
)
from crewai.utilities.events import crewai_event_bus as _crewai_event_bus
from copilotkit.types import Message
from copilotkit.logging import get_logger
from copilotkit.runloop import queue_put, get_context_execution
from copilotkit.protocol import (
RuntimeEventTypes,
RunStarted,
RunFinished,
RunError,
NodeStarted,
NodeFinished,
agent_state_message,
text_message_start,
text_message_content,
text_message_end,
action_execution_start,
action_execution_args,
action_execution_end,
meta_event,
RuntimeMetaEventName,
PredictStateConfig,
)
logger = get_logger(__name__)
class CopilotKitProperties(BaseModel):
"""CopilotKit properties"""
actions: List[Any] = Field(default_factory=list)
class CopilotKitState(FlowState):
"""CopilotKit state"""
messages: List[Any] = Field(default_factory=list)
copilotkit: CopilotKitProperties = Field(default_factory=CopilotKitProperties)
async def crewai_flow_async_runner(flow: Flow, inputs: Dict[str, Any]):
"""
Runs a flow in a separate thread. Workaround since the flow will use
asyncio.run().
"""
async def crewai_flow_event_subscriber(flow: Any, event: CrewAIFlowEvent):
if isinstance(event, FlowStartedEvent):
await queue_put(
RunStarted(type=RuntimeEventTypes.RUN_STARTED, state=flow.state),
priority=True,
)
elif isinstance(event, MethodExecutionStartedEvent):
await queue_put(
NodeStarted(
type=RuntimeEventTypes.NODE_STARTED,
node_name=event.method_name,
state=flow.state,
),
priority=True,
)
elif isinstance(event, MethodExecutionFinishedEvent):
await queue_put(
NodeFinished(
type=RuntimeEventTypes.NODE_FINISHED,
node_name=event.method_name,
state=flow.state,
),
priority=True,
)
elif isinstance(event, FlowFinishedEvent):
await queue_put(
RunFinished(type=RuntimeEventTypes.RUN_FINISHED, state=flow.state),
priority=True,
)
def _global_event_listener(_sender: Any, _event: CrewAIFlowEvent, **_kw): # noqa: D401
# Forward to the async handler inside the flow's loop
loop = asyncio.get_running_loop()
loop.call_soon(
lambda: asyncio.create_task(crewai_flow_event_subscriber(flow, _event))
)
# Register for the specific event classes we care about to avoid noise
for _ev_cls in (
FlowStartedEvent,
MethodExecutionStartedEvent,
MethodExecutionFinishedEvent,
FlowFinishedEvent,
):
_crewai_event_bus.on(_ev_cls)(_global_event_listener) # type: ignore
try:
await flow.kickoff_async(inputs=inputs)
except Exception as e: # pylint: disable=broad-except
await queue_put(RunError(type=RuntimeEventTypes.RUN_ERROR, error=e))
async def copilotkit_emit_state(state: Any) -> Literal[True]:
"""
Emits intermediate state to CopilotKit.
Useful if you have a longer running node and you want to update the user with the current state of the node.
To install the CopilotKit SDK, run:
```bash
pip install copilotkit[crewai]
```
### Examples
```python
from copilotkit.crewai import copilotkit_emit_state
for i in range(10):
await some_long_running_operation(i)
await copilotkit_emit_state({"progress": i})
```
Parameters
----------
state : Any
The state to emit (Must be JSON serializable).
Returns
-------
Awaitable[bool]
Always return True.
"""
execution = get_context_execution()
state_as_dict = state.model_dump() if isinstance(state, BaseModel) else state
state = {
k: v for k, v in state_as_dict.items() if k not in ["messages", "copilotkit"]
}
await queue_put(
agent_state_message(
thread_id=execution["thread_id"],
agent_name=execution["agent_name"],
node_name=execution["node_name"],
run_id=execution["run_id"],
active=True,
role="assistant",
state=json.dumps(state_as_dict),
running=True,
)
)
return True
async def copilotkit_emit_message(message: str) -> str:
"""
Manually emits a message to CopilotKit. Useful in longer running nodes to update the user.
Important: You still need to return the messages from the node.
### Examples
```python
from copilotkit.crewai import copilotkit_emit_message
message = "Step 1 of 10 complete"
await copilotkit_emit_message(message)
# Return the message from the node
return {
"messages": [AIMessage(content=message)]
}
```
Parameters
----------
message : str
The message to emit.
Returns
-------
Awaitable[bool]
Always return True.
"""
message_id = str(uuid.uuid4())
await queue_put(
text_message_start(message_id=message_id, parent_message_id=None),
text_message_content(message_id=message_id, content=message),
text_message_end(message_id=message_id),
)
return message_id
async def copilotkit_emit_tool_call(
*, name: str, args: Dict[str, Any], tool_call_id: Optional[str] = None
) -> str:
"""
Manually emits a tool call to CopilotKit.
```python
from copilotkit.crewai import copilotkit_emit_tool_call
auto_id = await copilotkit_emit_tool_call(name="SearchTool", args={"steps": 10})
# With a custom ID for correlation/idempotency:
custom_id = await copilotkit_emit_tool_call(name="SearchTool", args={"steps": 10}, tool_call_id="my-custom-id")
```
Parameters
----------
name : str
The name of the tool to emit.
args : Dict[str, Any]
The arguments to emit.
tool_call_id : Optional[str]
Optional tool call ID. If not provided, a random UUID is generated.
When provided, this ID is used as both the toolCallId and
parentMessageId in AG-UI protocol events.
The caller is responsible for ensuring uniqueness.
Returns
-------
str
The tool call ID used for the emitted tool call.
"""
if not isinstance(name, str) or not name.strip():
raise CopilotKitMisuseError(
"Tool name must be a non-empty string for copilotkit_emit_tool_call"
)
if tool_call_id is not None:
if not isinstance(tool_call_id, str) or not tool_call_id.strip():
raise CopilotKitMisuseError(
"Tool call id must be a non-empty string when provided for copilotkit_emit_tool_call"
)
try:
args_json = json.dumps(args)
except (TypeError, ValueError) as e:
raise CopilotKitMisuseError(
f"Tool arguments for '{name}' are not JSON-serializable: {e}"
) from e
message_id = tool_call_id if tool_call_id is not None else str(uuid.uuid4())
try:
await queue_put(
action_execution_start(
action_execution_id=message_id,
action_name=name,
parent_message_id=message_id,
),
action_execution_args(action_execution_id=message_id, args=args_json),
action_execution_end(action_execution_id=message_id),
)
except Exception:
try:
await queue_put(
action_execution_end(action_execution_id=message_id),
)
except Exception:
logger.error(
"Failed to emit compensating action_execution_end for %s",
message_id,
exc_info=True,
)
raise
return message_id
async def copilotkit_stream(response):
"""
Stream litellm responses token by token to CopilotKit.
```python
response = await copilotkit_stream(
completion(
model="openai/gpt-4o",
messages=messages,
tools=tools,
stream=True # this must be set to True for streaming
)
)
```
"""
if isinstance(response, ModelResponse):
return _copilotkit_stream_response(response)
if isinstance(response, CustomStreamWrapper):
return await _copilotkit_stream_custom_stream_wrapper(response)
raise ValueError("Invalid response type")
async def _copilotkit_stream_custom_stream_wrapper(response: CustomStreamWrapper):
message_id: str = ""
tool_call_id: str = ""
content = ""
created = 0
model = ""
system_fingerprint = ""
finish_reason = None
mode = None
all_tool_calls = []
for chunk in response:
if message_id is None:
message_id = chunk["id"]
tool_calls = chunk["choices"][0]["delta"]["tool_calls"]
finish_reason = chunk["choices"][0]["finish_reason"]
created = chunk["created"]
model = chunk["model"]
system_fingerprint = chunk["system_fingerprint"]
if mode == "text" and (tool_calls is not None or finish_reason is not None):
# end the current text message
await queue_put(text_message_end(message_id=message_id))
elif mode == "tool" and (tool_calls is None or finish_reason is not None):
# end the current tool call
await queue_put(action_execution_end(action_execution_id=tool_call_id))
if finish_reason is not None:
break
if mode != "text" and tool_calls is None:
# start a new text message
await queue_put(
text_message_start(message_id=message_id, parent_message_id=None)
)
elif mode != "tool" and tool_calls is not None and tool_calls[0].id is not None:
# start a new tool call
tool_call_id = tool_calls[0].id
await queue_put(
action_execution_start(
action_execution_id=tool_call_id,
action_name=tool_calls[0].function["name"],
parent_message_id=message_id,
)
)
all_tool_calls.append(
{
"id": tool_call_id,
"name": tool_calls[0].function["name"],
"arguments": "",
}
)
mode = "tool" if tool_calls is not None else "text"
if mode == "text":
text_content = chunk["choices"][0]["delta"]["content"]
if text_content is not None:
content += text_content
await queue_put(
text_message_content(message_id=message_id, content=text_content)
)
elif mode == "tool":
tool_arguments = tool_calls[0].function["arguments"]
if tool_arguments is not None:
await queue_put(
action_execution_args(
action_execution_id=tool_call_id, args=tool_arguments
)
)
all_tool_calls[-1]["arguments"] += tool_arguments
tool_calls = [
ChatCompletionMessageToolCall(
function=LiteLLMFunction(
arguments=tool_call["arguments"], name=tool_call["name"]
),
id=tool_call["id"],
type="function",
)
for tool_call in all_tool_calls
]
return ModelResponse(
id=message_id,
created=created,
model=model,
object="chat.completion",
system_fingerprint=system_fingerprint,
choices=[
Choices(
finish_reason=finish_reason,
index=0,
message=LiteLLMMessage(
content=content,
role="assistant",
tool_calls=tool_calls if len(tool_calls) > 0 else None,
function_call=None,
),
)
],
)
def _copilotkit_stream_response(response: ModelResponse):
return response
async def copilotkit_exit() -> Literal[True]:
"""
Exits the current agent after the run completes. Calling copilotkit_exit() will
not immediately stop the agent. Instead, it signals to CopilotKit to stop the agent after
the run completes.
### Examples
```python
from copilotkit.crewai import copilotkit_exit
def my_function():
await copilotkit_exit()
return state
```
Returns
-------
Awaitable[bool]
Always return True.
"""
await queue_put(meta_event(name=RuntimeMetaEventName.EXIT, value=True))
return True
async def copilotkit_predict_state(
config: Dict[str, PredictStateConfig],
) -> Literal[True]:
"""
Stream tool calls as state to CopilotKit.
To emit a tool call as streaming CrewAI state, pass the destination key in state,
the tool name and optionally the tool argument. (If you don't pass the argument name,
all arguments are emitted under the state key.)
```python
from copilotkit.crewai import copilotkit_predict_state
await copilotkit_predict_state(
{
"steps": {
"tool_name": "SearchTool",
"tool_argument": "steps",
},
}
)
```
Parameters
----------
config : Dict[str, CopilotKitPredictStateConfig]
The configuration to predict the state.
Returns
-------
Awaitable[bool]
Always return True.
"""
await queue_put(meta_event(name=RuntimeMetaEventName.PREDICT_STATE, value=config))
return True
def copilotkit_messages_to_crewai_flow(messages: List[Message]) -> List[Any]:
"""
Convert CopilotKit messages to CrewAI Flow messages
"""
result = []
processed_action_executions = set()
for message in messages:
message_id = message["id"]
message_type = message.get("type")
if message_type == "TextMessage":
result.append(
{
"id": message_id,
"role": message.get("role"),
"content": message.get("content"),
}
)
elif message_type == "ActionExecutionMessage":
# convert multiple tool calls to a single message
original_message_id = message.get("parentMessageId", message_id)
if original_message_id in processed_action_executions:
continue
processed_action_executions.add(original_message_id)
all_tool_calls = []
# Find all tool calls for this message
for msg in messages:
msg_id = msg["id"]
if (
msg.get("parentMessageId", None) == original_message_id
or msg_id == original_message_id
):
all_tool_calls.append(msg)
tool_calls = [
{
"type": "function",
"function": {
"name": t["name"],
"arguments": json.dumps(t["arguments"]),
},
"id": t["id"],
}
for t in all_tool_calls
]
result.append(
{
"id": original_message_id,
"role": "assistant",
"content": "",
"tool_calls": tool_calls,
}
)
elif message_type == "ResultMessage":
result.append(
{
"id": message_id,
"role": "tool",
"tool_call_id": message.get("actionExecutionId"),
"content": message.get("result"),
}
)
return result
def crewai_flow_messages_to_copilotkit(messages: List[Dict]) -> List[Message]: # pylint: disable=too-many-branches
"""
Convert CrewAI Flow messages to CopilotKit messages
"""
result = []
tool_call_names = {}
message_ids = {id(m): m.get("id", str(uuid.uuid4())) for m in messages}
for message in messages:
if "content" in message and message.get("role") == "assistant":
if message.get("tool_calls"):
for tool_call in message["tool_calls"]:
tc_id = tool_call.get("id")
if tc_id is None:
continue
if tool_call.get("function"):
tool_call_names[tc_id] = tool_call["function"].get("name", "")
else:
tool_call_names[tc_id] = tool_call.get("name", "")
for message in messages:
message_id = message_ids[id(message)]
if message.get("role") == "tool":
result.append(
{
"actionExecutionId": message["tool_call_id"],
"actionName": tool_call_names.get(
message["tool_call_id"], message.get("name", "")
),
"result": message["content"],
"id": message_id,
}
)
elif message.get("tool_calls"):
# Always emit the assistant message, even with empty content.
# Tool call entries reference it via parentMessageId; omitting it
# orphans tool calls and breaks frontend thread reconstruction.
result.append(
{
"role": message["role"],
"content": message.get("content")
if message.get("content") is not None
else "",
"id": message_id,
}
)
for tool_call in message["tool_calls"]:
tc_id = tool_call.get("id")
if tc_id is None:
continue
if tool_call.get("function"):
result.append(
{
"id": tc_id,
"name": tool_call["function"].get("name", ""),
"arguments": json.loads(tool_call["function"]["arguments"]),
"parentMessageId": message_id,
}
)
else:
result.append(
{
"id": tc_id,
"name": tool_call.get("name", ""),
"arguments": tool_call.get("arguments", {}),
"parentMessageId": message_id,
}
)
elif message.get("content"):
result.append(
{
"role": message["role"],
"content": message["content"],
"id": message_id,
}
)
# Create a dictionary to map message ids to their corresponding messages
results_dict = {
msg["actionExecutionId"]: msg for msg in result if "actionExecutionId" in msg
}
# since we are splitting multiple tool calls into multiple messages,
# we need to reorder the corresponding result messages to be after the tool call
reordered_result = []
for msg in result:
# add all messages that are not tool call results
if not "actionExecutionId" in msg:
reordered_result.append(msg)
# if the message is a tool call, also add the corresponding result message
# immediately after the tool call
if msg.get("name"):
msg_id = msg["id"]
if msg_id in results_dict:
reordered_result.append(results_dict[msg_id])
else:
logger.warning("Tool call result message not found for id: %s", msg_id)
return reordered_result
+54
View File
@@ -0,0 +1,54 @@
"""Exceptions for CopilotKit."""
class CopilotKitError(Exception):
"""Base exception for all CopilotKit errors.
Catch this to handle any CopilotKit-specific exception.
"""
pass
class ActionNotFoundException(CopilotKitError):
"""Exception raised when an action or agent is not found."""
def __init__(self, name: str):
self.name = name
super().__init__(f"Action '{name}' not found.")
class AgentNotFoundException(CopilotKitError):
"""Exception raised when an agent is not found."""
def __init__(self, name: str):
self.name = name
super().__init__(f"Agent '{name}' not found.")
class ActionExecutionException(CopilotKitError):
"""Exception raised when an action fails to execute."""
def __init__(self, name: str, error: Exception):
self.name = name
self.error = error
super().__init__(f"Action '{name}' failed to execute: {error}")
class AgentExecutionException(CopilotKitError):
"""Exception raised when an agent fails to execute."""
def __init__(self, name: str, error: Exception):
self.name = name
self.error = error
super().__init__(f"Agent '{name}' failed to execute: {error}")
class CopilotKitMisuseError(CopilotKitError, ValueError):
"""Exception raised when CopilotKit detects incorrect usage of its APIs.
Inherits from both CopilotKitError (for ``except CopilotKitError``) and
ValueError (for backward compatibility with ``except ValueError`` handlers).
"""
pass
+201
View File
@@ -0,0 +1,201 @@
"""Forward CopilotKit request-context headers onto outbound LLM/provider HTTP calls
so downstream services (e.g. the aimock test server, proxies, request routing /
fixture-matching infrastructure) can correlate the outbound provider call with the
original inbound request.
What this module does
---------------------
On each inbound request the application stores a small set of ``x-*`` prefixed
headers (for example ``x-aimock-context``, ``x-aimock-session``, ``x-request-id``,
``x-trace-id``) on a per-request ``contextvars.ContextVar``. When the application
later makes an outbound HTTP call to an LLM provider (OpenAI, Anthropic, or any
client that wraps ``httpx``), an httpx request event hook reads that ContextVar
and copies those same headers onto the outbound request so downstream services
can correlate the two.
This is plain header propagation, not data collection. Scope and limits:
* Only headers the application itself set on the request context via
``set_forwarded_headers`` are forwarded. The module never reads request
bodies, cookies, user data, credentials, or anything off the inbound
request beyond the headers explicitly handed to it.
* Only ``x-*`` prefixed headers pass the filter; ``authorization``,
``content-type``, and any other non ``x-*`` headers are dropped.
* Nothing is collected, persisted, logged, or sent anywhere by this module
itself — it only attaches headers to an HTTP request that the caller was
already going to make. There is no telemetry, no out-of-band channel, and
no end-user data flow.
Mechanics
---------
``install_httpx_hook`` does two small things:
1. It walks the ``._client`` chain on the given object (modern provider SDKs
wrap their httpx client behind several layers of ``._client``) to find the
first object that exposes an httpx-style ``event_hooks`` mapping.
2. It attaches a request event hook to that mapping. The hook flavor matches
the client: an async coroutine hook for ``httpx.AsyncClient`` (httpx awaits
request hooks on async clients), and a plain sync hook for ``httpx.Client``.
Installation is idempotent via a marker attribute on the installed callable.
This mirrors the CopilotKit runtime's ``extractForwardableHeaders()`` behavior
on the Node side so the Python SDK forwards the same set of context headers.
"""
import contextvars
import warnings
from typing import Any, Dict, Optional
# Per-request storage for the set of headers the application has asked to forward
# onto outbound LLM/provider calls (populated by ``set_forwarded_headers``).
_forwarded_headers: contextvars.ContextVar[Dict[str, str]] = contextvars.ContextVar(
"copilotkit_forwarded_headers"
)
# Marker used to identify hooks we have already installed, so install_httpx_hook
# is idempotent across repeated calls on the same client.
_HOOK_MARKER = "_copilotkit_forwarded_header_hook"
# Bound on how deep we'll walk a ``._client`` chain looking for event_hooks.
# The modern OpenAI SDK shape is:
# ChatOpenAI.client -> Completions/AsyncCompletions resource
# -> ._client = openai.OpenAI / AsyncOpenAI (no event_hooks)
# -> ._client._client = httpx wrapper (HAS event_hooks)
# 5 hops is plenty of headroom for similar SDKs without risking pathological loops.
_MAX_CHAIN_DEPTH = 5
def set_forwarded_headers(headers: Dict[str, str]) -> None:
"""Record the set of headers to forward onto outbound LLM/provider calls
made later in this request context.
Only ``x-*`` prefixed headers are kept; everything else is dropped.
"""
filtered = {k.lower(): v for k, v in headers.items() if k.lower().startswith("x-")}
_forwarded_headers.set(filtered)
def get_forwarded_headers() -> Dict[str, str]:
"""Return the headers the application has asked to forward onto outbound
LLM/provider calls in the current request context."""
return _forwarded_headers.get({})
def _find_event_hooks_target(client: Any) -> Optional[Any]:
"""Walk the ``._client`` chain looking for the first object that exposes
an httpx-style ``event_hooks`` mapping.
Returns the target object, or ``None`` if no such object is found within
``_MAX_CHAIN_DEPTH`` hops.
"""
current = client
for _ in range(_MAX_CHAIN_DEPTH + 1):
if current is None:
return None
if hasattr(current, "event_hooks"):
return current
nxt = getattr(current, "_client", None)
if nxt is current or nxt is None:
return None
current = nxt
return None
def install_httpx_hook(client: Any) -> None:
"""Attach a request event hook to ``client``'s underlying httpx client so
that headers recorded via ``set_forwarded_headers`` are copied onto
outbound requests.
Works with OpenAI and Anthropic Python SDKs (both wrap httpx internally,
sometimes via several layers of ``._client`` indirection), as well as raw
``httpx.Client`` / ``httpx.AsyncClient`` instances.
For ``httpx.AsyncClient`` an async hook is attached (httpx awaits request
hooks on async clients); for sync clients a sync hook is attached.
Idempotent: a marker attribute on the installed callable prevents double
installation on the same target.
Parameters
----------
client : object
An OpenAI/Anthropic client instance, or a raw httpx.Client/AsyncClient.
"""
target = _find_event_hooks_target(client)
if target is None:
warnings.warn(
f"install_httpx_hook: client of type {type(client).__name__} has no "
"recognized event_hooks attribute; x-* headers will not be forwarded",
stacklevel=2,
)
return
request_hooks = target.event_hooks.get("request", [])
# Idempotency: don't double-install on the same target.
for existing in request_hooks:
if getattr(existing, _HOOK_MARKER, False):
return
# Choose sync vs async hook flavor based on the target class.
# httpx.AsyncClient awaits request hooks; a sync hook returning None would
# raise "TypeError: object NoneType can't be used in 'await' expression",
# which surfaces as APIConnectionError to the caller.
is_async = _is_async_httpx_target(target)
if is_async:
async def _inject_headers_async(request):
headers = get_forwarded_headers()
for key, value in headers.items():
request.headers[key] = value
setattr(_inject_headers_async, _HOOK_MARKER, True)
request_hooks.append(_inject_headers_async)
else:
def _inject_headers(request):
headers = get_forwarded_headers()
for key, value in headers.items():
request.headers[key] = value
setattr(_inject_headers, _HOOK_MARKER, True)
request_hooks.append(_inject_headers)
# In case ``event_hooks`` returned a fresh list (defensive), make sure the
# mutation is reflected on the target.
target.event_hooks["request"] = request_hooks
def _is_async_httpx_target(target: Any) -> bool:
"""Best-effort detection: is this object an httpx async client?
Tries ``isinstance`` against the real ``httpx.AsyncClient`` / ``httpx.Client``
first (the authoritative answer for real clients). If httpx is not
importable, or the target is neither of those (e.g. a wrapped or
duck-typed client used in tests), falls back to an EXACT MRO class-name
match against ``"AsyncClient"``. Avoids a broad ``startswith("Async")``
check, which would misclassify a sync client whose MRO happens to
include an ``Async*``-named base (e.g. ``AsyncContextManager``) as
async — attaching an async hook that httpx calls synchronously would
leave the coroutine unawaited and the forwarded headers would not be
attached to the outbound request.
"""
try:
import httpx # local import keeps httpx an optional concern at import time
if isinstance(target, httpx.AsyncClient):
return True
if isinstance(target, httpx.Client):
return False
except (
ImportError
): # pragma: no cover - httpx should always be importable in practice
pass
# Fall back to exact class-name match for wrapped/duck-typed clients.
for cls in type(target).__mro__:
if cls.__name__ == "AsyncClient":
return True
return False
+178
View File
@@ -0,0 +1,178 @@
"""
HTML templates, used when the info endpoint is accessed from the browser.
"""
import json
from copilotkit.sdk import InfoDict
HEAD_HTML = """
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>CopilotKit Remote Endpoint v0.1.12</title>
<style>
body {
font-family: Arial, sans-serif;
line-height: 1.6;
color: #333;
background-color: #f4f4f4;
margin: 0;
padding: 0;
}
.container {
max-width: 800px;
margin: 0 auto;
padding: 30px;
}
header {
display: flex;
align-items: center;
justify-content: center;
margin-bottom: 40px;
}
h1 {
font-size: 2rem;
margin: 0;
}
h2 {
font-size: 1.8rem;
margin-bottom: 20px;
}
h3 {
font-size: 1.4rem;
margin-bottom: 10px;
}
.version {
font-family: 'Courier New', Courier, monospace;
font-size: 1.2rem;
}
.kite-icon {
font-size: 38px;
margin-right: 16px;
}
.grid {
display: grid;
grid-template-columns: 1fr;
gap: 20px;
margin-bottom: 40px;
}
.card {
background-color: #fff;
border-radius: 8px;
box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
padding: 20px;
}
.badge {
display: inline-block;
padding: 4px 8px;
font-size: 0.75rem;
font-weight: bold;
border-radius: 4px;
margin-left: 10px;
background-color: #dbeafe;
color: #1e40af;
}
pre {
background-color: #f1f1f1;
padding: 10px;
border-radius: 4px;
overflow-x: auto;
white-space: pre-wrap;
word-wrap: break-word;
}
code {
font-family: 'Courier New', Courier, monospace;
}
</style>
</head>
"""
INFO_TEMPLATE = """
<!DOCTYPE html>
<html lang="en">
{head_html}
<body>
<div class="container">
<header>
<h1><span class="kite-icon">🪁</span>CopilotKit Remote Endpoint <span class="version">(v{version})</span></h1>
</header>
<main>
<section>
<h2>Actions</h2>
<div class="grid">
{action_html}
</div>
</section>
<section>
<h2>Agents</h2>
<div class="grid">
{agent_html}
</div>
</section>
</main>
</div>
</body>
</html>
"""
ACTION_TEMPLATE = """
<div class="card">
<h3>{name}</h3>
<p>{description}</p>
<h4>Arguments:</h4>
<pre><code>{arguments}</code></pre>
</div>
"""
AGENT_TEMPLATE = """
<div class="card">
<h3>{name} <span class="badge">{type}</span></h3>
<p>{description}</p>
</div>
"""
NO_ACTIONS_FOUND_HTML = """
<div class="card">
<p>No actions found</p>
</div>
"""
NO_AGENTS_FOUND_HTML = """
<div class="card">
<p>No agents found</p>
</div>
"""
def generate_info_html(info: InfoDict) -> str:
"""
Generate HTML for the info endpoint
"""
print(info, flush=True)
action_html = ""
for action in info["actions"]:
action_html += ACTION_TEMPLATE.format(
name=action["name"],
description=action["description"],
arguments=json.dumps(action.get("parameters", []), indent=2),
)
agent_html = ""
for agent in info["agents"]:
agent_type = agent.get("type", "Unknown")
if agent_type == "langgraph":
agent_type = "LangGraph"
elif agent_type == "crewai":
agent_type = "CrewAI"
agent_html += AGENT_TEMPLATE.format(
name=agent["name"],
type=agent_type,
description=agent["description"],
)
return INFO_TEMPLATE.format(
head_html=HEAD_HTML,
version=info["sdkVersion"],
action_html=action_html or NO_ACTIONS_FOUND_HTML,
agent_html=agent_html or NO_AGENTS_FOUND_HTML,
)
@@ -0,0 +1,332 @@
"""FastAPI integration"""
import logging
import asyncio
import re
import uuid
import warnings
from concurrent.futures import ThreadPoolExecutor
from typing import List, Any, cast, Optional
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse, StreamingResponse, HTMLResponse
from fastapi.encoders import jsonable_encoder
from ..sdk import CopilotKitRemoteEndpoint, CopilotKitContext
from ..types import Message, MetaEvent
from ..exc import (
ActionNotFoundException,
ActionExecutionException,
AgentNotFoundException,
AgentExecutionException,
)
from ..action import ActionDict
from ..html import generate_info_html
from ..header_propagation import set_forwarded_headers
logging.basicConfig(level=logging.ERROR)
logger = logging.getLogger(__name__)
def add_fastapi_endpoint(
fastapi_app: FastAPI,
sdk: CopilotKitRemoteEndpoint,
prefix: str,
*,
use_thread_pool: bool = False,
max_workers: int = 10,
):
"""Add FastAPI endpoint with configurable ThreadPoolExecutor size"""
if use_thread_pool:
warnings.warn(
"The 'use_thread_pool' parameter is deprecated "
+ "and will be removed in a future version.",
DeprecationWarning,
)
def run_handler_in_thread(request: Request, sdk: CopilotKitRemoteEndpoint):
# Run the handler coroutine in the event loop
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
return loop.run_until_complete(handler(request, sdk))
async def make_handler(request: Request):
if use_thread_pool:
executor = ThreadPoolExecutor(max_workers=max_workers)
loop = asyncio.get_event_loop()
future = loop.run_in_executor(executor, run_handler_in_thread, request, sdk)
return await future
return await handler(request, sdk)
# Ensure the prefix starts with a slash and remove trailing slashes
normalized_prefix = "/" + prefix.strip("/")
fastapi_app.add_api_route(
f"{normalized_prefix}/{{path:path}}",
make_handler,
methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
)
def body_get_or_raise(body: Any, key: str):
"""Get value from body or raise an error"""
value = body.get(key)
if value is None:
raise HTTPException(status_code=400, detail=f"{key} is required")
return value
async def handler(request: Request, sdk: CopilotKitRemoteEndpoint):
"""Handle FastAPI request"""
try:
body = await request.json()
except: # pylint: disable=bare-except
body = None
# Propagate x-aimock-* headers to outgoing LLM calls via ContextVar
set_forwarded_headers(dict(request.headers))
path = request.path_params.get("path")
method = request.method
context = cast(
CopilotKitContext,
{
"properties": (body or {}).get("properties", {}),
"frontend_url": (body or {}).get("frontendUrl", None),
"headers": request.headers,
},
)
# handle / request for info endpoint
if method in ["GET", "POST"] and path == "":
accept_header = request.headers.get("accept", "")
return await handle_info(
sdk=sdk,
context=context,
as_html="text/html" in accept_header,
)
# handle /agent/name request for executing an agent
if method == "POST" and (match := re.match(r"agent/([a-zA-Z0-9_-]+)", path)):
name = match.group(1)
body = body or {}
thread_id = body.get("threadId", str(uuid.uuid4()))
state = body.get("state", {})
messages = body.get("messages", [])
actions = body.get("actions", [])
# used for LangGraph only
node_name = body.get("nodeName")
return handle_execute_agent(
sdk=sdk,
context=context,
thread_id=thread_id,
node_name=node_name,
name=name,
state=state,
messages=messages,
actions=actions,
)
# handle /agent/name/state request for getting agent state
if method == "POST" and (match := re.match(r"agent/([a-zA-Z0-9_-]+)/state", path)):
name = match.group(1)
thread_id = body_get_or_raise(body, "threadId")
return await handle_get_agent_state(
sdk=sdk,
context=context,
thread_id=thread_id,
name=name,
)
# handle /action/name request for executing an action
if method == "POST" and (match := re.match(r"action/([a-zA-Z0-9_-]+)", path)):
name = match.group(1)
arguments = body.get("arguments", {})
return await handle_execute_action(
sdk=sdk,
context=context,
name=name,
arguments=arguments,
)
# v2: POST /agents/name/state
# Deal with backwards compatibility
result_v1 = await handler_v1(
sdk=sdk,
method=method,
path=path,
body=body,
context=context,
)
if result_v1 is not None:
return result_v1
raise HTTPException(status_code=404, detail="Not found")
async def handler_v1(
sdk: CopilotKitRemoteEndpoint,
method: str,
path: str,
body: Any,
context: CopilotKitContext,
):
"""Handle FastAPI request for v1"""
if body is None:
raise HTTPException(status_code=400, detail="Request body is required")
if method == "POST" and path == "info":
return await handle_info(sdk=sdk, context=context)
if method == "POST" and path == "actions/execute":
name = body_get_or_raise(body, "name")
arguments = body.get("arguments", {})
return await handle_execute_action(
sdk=sdk,
context=context,
name=name,
arguments=arguments,
)
if method == "POST" and path == "agents/execute":
thread_id = body.get("threadId")
node_name = body.get("nodeName")
config = body.get("config")
name = body_get_or_raise(body, "name")
state = body_get_or_raise(body, "state")
messages = body_get_or_raise(body, "messages")
actions = cast(List[ActionDict], body.get("actions", []))
meta_events = cast(List[MetaEvent], body.get("metaEvents", []))
return handle_execute_agent(
sdk=sdk,
context=context,
thread_id=thread_id,
node_name=node_name,
name=name,
state=state,
config=config,
messages=messages,
actions=actions,
meta_events=meta_events,
)
if method == "POST" and path == "agents/state":
thread_id = body_get_or_raise(body, "threadId")
name = body_get_or_raise(body, "name")
return await handle_get_agent_state(
sdk=sdk,
context=context,
thread_id=thread_id,
name=name,
)
return None
async def handle_info(
*,
sdk: CopilotKitRemoteEndpoint,
context: CopilotKitContext,
as_html: bool = False,
):
"""Handle info request with FastAPI"""
result = sdk.info(context=context)
if as_html:
return HTMLResponse(content=generate_info_html(result))
return JSONResponse(content=jsonable_encoder(result))
async def handle_execute_action(
*,
sdk: CopilotKitRemoteEndpoint,
context: CopilotKitContext,
name: str,
arguments: dict,
):
"""Handle execute action request with FastAPI"""
try:
result = await sdk.execute_action(
context=context, name=name, arguments=arguments
)
return JSONResponse(content=jsonable_encoder(result))
except ActionNotFoundException as exc:
logger.error("Action not found: %s", exc)
return JSONResponse(content={"error": str(exc)}, status_code=404)
except ActionExecutionException as exc:
logger.error("Action execution error: %s", exc)
return JSONResponse(content={"error": str(exc)}, status_code=500)
except Exception as exc: # pylint: disable=broad-except
logger.error("Action execution error: %s", exc)
return JSONResponse(content={"error": str(exc)}, status_code=500)
def handle_execute_agent( # pylint: disable=too-many-arguments
*,
sdk: CopilotKitRemoteEndpoint,
context: CopilotKitContext,
thread_id: str,
name: str,
state: dict,
config: Optional[dict] = None,
messages: List[Message],
actions: List[ActionDict],
node_name: str,
meta_events: Optional[List[MetaEvent]] = None,
):
"""Handle continue agent execution request with FastAPI"""
try:
events = sdk.execute_agent(
context=context,
thread_id=thread_id,
name=name,
node_name=node_name,
state=state,
config=config,
messages=messages,
actions=actions,
meta_events=meta_events,
)
return StreamingResponse(events, media_type="application/json")
except AgentNotFoundException as exc:
logger.error("Agent not found: %s", exc, exc_info=True)
return JSONResponse(content={"error": str(exc)}, status_code=404)
except AgentExecutionException as exc:
logger.error("Agent execution error: %s", exc, exc_info=True)
return JSONResponse(content={"error": str(exc)}, status_code=500)
except Exception as exc: # pylint: disable=broad-except
logger.error("Agent execution error: %s", exc, exc_info=True)
return JSONResponse(content={"error": str(exc)}, status_code=500)
async def handle_get_agent_state(
*,
sdk: CopilotKitRemoteEndpoint,
context: CopilotKitContext,
thread_id: str,
name: str,
):
"""Handle get agent state request with FastAPI"""
try:
result = await sdk.get_agent_state(
context=context,
thread_id=thread_id,
name=name,
)
return JSONResponse(content=jsonable_encoder(result))
except AgentNotFoundException as exc:
logger.error("Agent not found: %s", exc, exc_info=True)
return JSONResponse(content={"error": str(exc)}, status_code=404)
except Exception as exc: # pylint: disable=broad-except
logger.error("Agent get state error: %s", exc, exc_info=True)
return JSONResponse(content={"error": str(exc)}, status_code=500)
+30
View File
@@ -0,0 +1,30 @@
"""
copilotkit.langchain is deprecated. Use copilotkit.langgraph instead.
"""
import warnings
from copilotkit.langgraph import (
langchain_messages_to_copilotkit,
copilotkit_customize_config,
copilotkit_exit,
copilotkit_emit_state,
copilotkit_emit_message,
copilotkit_emit_tool_call,
copilotkit_interrupt,
)
warnings.warn(
"copilotkit.langchain is deprecated. Use copilotkit.langgraph instead.",
DeprecationWarning,
stacklevel=2,
)
__all__ = [
"langchain_messages_to_copilotkit",
"copilotkit_customize_config",
"copilotkit_exit",
"copilotkit_emit_state",
"copilotkit_emit_message",
"copilotkit_emit_tool_call",
"copilotkit_interrupt",
]
+495
View File
@@ -0,0 +1,495 @@
"""
LangChain specific utilities for CopilotKit
"""
import uuid
import json
import warnings
import asyncio
from typing import List, Optional, Any, Union, Dict
from typing_extensions import TypedDict
from langgraph.graph import MessagesState
from langchain_core.messages import (
HumanMessage,
SystemMessage,
BaseMessage,
AIMessage,
ToolMessage,
)
from langchain_core.runnables import RunnableConfig
from langchain_core.callbacks.manager import adispatch_custom_event
from langgraph.types import interrupt
from .types import Message, IntermediateStateConfig
from .exc import CopilotKitMisuseError
from .logging import get_logger
logger = get_logger(__name__)
class CopilotContextItem(TypedDict):
"""Copilot context item"""
description: str
value: Any
class CopilotKitProperties(TypedDict):
"""CopilotKit state"""
actions: List[Any]
context: List[CopilotContextItem]
# Private state for CopilotKit middleware
intercepted_tool_calls: Any
original_ai_message_id: Any
class CopilotKitState(MessagesState):
"""CopilotKit state"""
copilotkit: CopilotKitProperties
def langchain_messages_to_copilotkit(messages: List[BaseMessage]) -> List[Message]:
"""
Convert LangChain messages to CopilotKit messages
"""
result = []
tool_call_names = {}
for message in messages:
if isinstance(message, AIMessage):
for tool_call in message.tool_calls or []:
tool_call_names[tool_call["id"]] = tool_call["name"]
for message in messages:
content = None
if hasattr(message, "content"):
content = message.content
# Content can be a list of content blocks (e.g. Anthropic models).
# Extract and concatenate all text parts instead of only taking
# the first element.
if isinstance(content, list):
text_parts = []
for part in content:
if isinstance(part, str):
text_parts.append(part)
elif isinstance(part, dict) and part.get("type") == "text":
text_parts.append(part.get("text", ""))
elif isinstance(part, dict) and "text" in part:
text_parts.append(part.get("text", ""))
content = "".join(text_parts)
# Anthropic models return a dict with a "text" key
if isinstance(content, dict):
content = content.get("text", "")
if isinstance(message, HumanMessage):
result.append(
{
"role": "user",
"content": content,
"id": message.id,
}
)
elif isinstance(message, SystemMessage):
result.append(
{
"role": "system",
"content": content,
"id": message.id,
}
)
elif isinstance(message, AIMessage):
# Always emit the assistant message, even with empty content.
# Tool call entries reference it via parentMessageId; omitting it
# orphans tool calls and breaks frontend thread reconstruction.
result.append(
{
"role": "assistant",
"content": content if content is not None else "",
"id": message.id,
}
)
if message.tool_calls:
for tool_call in message.tool_calls:
result.append(
{
"id": tool_call["id"],
"name": tool_call["name"],
"arguments": tool_call["args"],
"parentMessageId": message.id,
}
)
elif isinstance(message, ToolMessage):
result.append(
{
"actionExecutionId": message.tool_call_id,
"actionName": tool_call_names.get(
message.tool_call_id, message.name or ""
),
"result": content,
"id": message.id,
}
)
# Create a dictionary to map message ids to their corresponding messages
results_dict = {
msg["actionExecutionId"]: msg for msg in result if "actionExecutionId" in msg
}
# since we are splitting multiple tool calls into multiple messages,
# we need to reorder the corresponding result messages to be after the tool call
reordered_result = []
for msg in result:
# add all messages that are not tool call results
if not "actionExecutionId" in msg:
reordered_result.append(msg)
# if the message is a tool call, also add the corresponding result message
# immediately after the tool call
if "arguments" in msg:
msg_id = msg["id"]
if msg_id in results_dict:
reordered_result.append(results_dict[msg_id])
else:
logger.warning("Tool call result message not found for id: %s", msg_id)
return reordered_result
def copilotkit_customize_config(
base_config: Optional[RunnableConfig] = None,
*,
emit_messages: Optional[bool] = None,
emit_tool_calls: Optional[Union[bool, str, List[str]]] = None,
emit_intermediate_state: Optional[List[IntermediateStateConfig]] = None,
emit_all: Optional[bool] = None, # deprecated
) -> RunnableConfig:
"""
Customize the LangGraph configuration for use in CopilotKit.
To install the CopilotKit SDK, run:
```bash
pip install copilotkit
```
### Examples
Disable emitting messages and tool calls:
```python
from copilotkit.langgraph import copilotkit_customize_config
config = copilotkit_customize_config(
config,
emit_messages=False,
emit_tool_calls=False
)
```
To emit a tool call as streaming LangGraph state, pass the destination key in state,
the tool name and optionally the tool argument. (If you don't pass the argument name,
all arguments are emitted under the state key.)
```python
from copilotkit.langgraph import copilotkit_customize_config
config = copilotkit_customize_config(
config,
emit_intermediate_state=[
{
"state_key": "steps",
"tool": "SearchTool",
"tool_argument": "steps"
},
]
)
```
Parameters
----------
base_config : Optional[RunnableConfig]
The LangChain/LangGraph configuration to customize. Pass None to make a new configuration.
emit_messages : Optional[bool]
Configure how messages are emitted. By default, all messages are emitted. Pass False to
disable emitting messages.
emit_tool_calls : Optional[Union[bool, str, List[str]]]
Configure how tool calls are emitted. By default, all tool calls are emitted. Pass False to
disable emitting tool calls. Pass a string or list of strings to emit only specific tool calls.
emit_intermediate_state : Optional[List[IntermediateStateConfig]]
Lets you emit tool calls as streaming LangGraph state.
Returns
-------
RunnableConfig
The customized LangGraph configuration.
"""
if emit_all is not None:
warnings.warn(
"The `emit_all` parameter is deprecated and will be removed in a future version. "
"CopilotKit will now emit all messages and tool calls by default.",
DeprecationWarning,
stacklevel=2,
)
metadata = base_config.get("metadata", {}) if base_config else {}
if emit_all is True:
metadata["copilotkit:emit-tool-calls"] = True
metadata["copilotkit:emit-messages"] = True
else:
if emit_tool_calls is not None:
metadata["copilotkit:emit-tool-calls"] = emit_tool_calls
if emit_messages is not None:
metadata["copilotkit:emit-messages"] = emit_messages
if emit_intermediate_state:
metadata["copilotkit:emit-intermediate-state"] = emit_intermediate_state
base_config = base_config or {}
return {**base_config, "metadata": metadata}
async def copilotkit_exit(config: RunnableConfig):
"""
Exits the current agent after the run completes. Calling copilotkit_exit() will
not immediately stop the agent. Instead, it signals to CopilotKit to stop the agent after
the run completes.
### Examples
```python
from copilotkit.langgraph import copilotkit_exit
def my_node(state: Any):
await copilotkit_exit(config)
return state
```
Parameters
----------
config : RunnableConfig
The LangGraph configuration.
Returns
-------
Awaitable[bool]
Always return True.
"""
await adispatch_custom_event(
"copilotkit_exit",
{},
config=config,
)
await asyncio.sleep(0.02)
return True
async def copilotkit_emit_state(config: RunnableConfig, state: Any):
"""
Emits intermediate state to CopilotKit. Useful if you have a longer running node and you want to
update the user with the current state of the node.
### Examples
```python
from copilotkit.langgraph import copilotkit_emit_state
for i in range(10):
await some_long_running_operation(i)
await copilotkit_emit_state(config, {"progress": i})
```
Parameters
----------
config : RunnableConfig
The LangGraph configuration.
state : Any
The state to emit (Must be JSON serializable).
Returns
-------
Awaitable[bool]
Always return True.
"""
await adispatch_custom_event(
"copilotkit_manually_emit_intermediate_state",
state,
config=config,
)
await asyncio.sleep(0.02)
return True
async def copilotkit_emit_message(config: RunnableConfig, message: str):
"""
Manually emits a message to CopilotKit. Useful in longer running nodes to update the user.
Important: You still need to return the messages from the node.
### Examples
```python
from copilotkit.langgraph import copilotkit_emit_message
message = "Step 1 of 10 complete"
await copilotkit_emit_message(config, message)
# Return the message from the node
return {
"messages": [AIMessage(content=message)]
}
```
Parameters
----------
config : RunnableConfig
The LangGraph configuration.
message : str
The message to emit.
Returns
-------
Awaitable[bool]
Always return True.
"""
await adispatch_custom_event(
"copilotkit_manually_emit_message",
{"message": message, "message_id": str(uuid.uuid4()), "role": "assistant"},
config=config,
)
await asyncio.shield(asyncio.sleep(0.02))
return True
async def copilotkit_emit_tool_call(
config: RunnableConfig,
*,
name: str,
args: Dict[str, Any],
tool_call_id: Optional[str] = None,
) -> str:
"""
Manually emits a tool call to CopilotKit.
```python
from copilotkit.langgraph import copilotkit_emit_tool_call
auto_id = await copilotkit_emit_tool_call(config, name="SearchTool", args={"steps": 10})
# With a custom ID for correlation/idempotency:
custom_id = await copilotkit_emit_tool_call(config, name="SearchTool", args={"steps": 10}, tool_call_id="my-custom-id")
```
Parameters
----------
config : RunnableConfig
The LangGraph configuration.
name : str
The name of the tool to emit.
args : Dict[str, Any]
The arguments to emit.
tool_call_id : Optional[str]
Optional tool call ID. If not provided, a random UUID is generated.
When provided, this ID is used as the toolCallId and parentMessageId
in AG-UI protocol events. The caller is responsible for ensuring uniqueness.
Returns
-------
str
The tool call ID used for the emitted tool call.
"""
if not isinstance(name, str) or not name.strip():
raise CopilotKitMisuseError(
"Tool name must be a non-empty string for copilotkit_emit_tool_call"
)
if tool_call_id is not None:
if not isinstance(tool_call_id, str) or not tool_call_id.strip():
raise CopilotKitMisuseError(
"Tool call id must be a non-empty string when provided for copilotkit_emit_tool_call"
)
else:
tool_call_id = str(uuid.uuid4())
try:
json.dumps(args)
except (TypeError, ValueError) as e:
raise CopilotKitMisuseError(
f"Tool arguments for '{name}' are not JSON-serializable: {e}"
) from e
await adispatch_custom_event(
"copilotkit_manually_emit_tool_call",
{"name": name, "args": args, "id": tool_call_id},
config=config,
)
# LangGraph's adispatch_custom_event is async but does not guarantee the event
# has been flushed to the SSE stream before it returns. Without this sleep,
# a subsequent emit can interleave and corrupt event ordering on the client.
# Shielded so that task cancellation doesn't prevent us from returning the ID.
try:
await asyncio.shield(asyncio.sleep(0.02))
except asyncio.CancelledError:
logger.warning(
"copilotkit_emit_tool_call cancelled during post-dispatch flush for "
"tool_call_id=%s; event was already dispatched",
tool_call_id,
)
raise
return tool_call_id
def copilotkit_interrupt(
message: Optional[str] = None,
action: Optional[str] = None,
args: Optional[Dict[str, Any]] = None,
):
if message is None and action is None:
raise ValueError(
"Either message or action (and optional arguments) must be provided"
)
interrupt_message = None
interrupt_values = None
answer = None
if message is not None:
interrupt_values = message
interrupt_message = AIMessage(content=message, id=str(uuid.uuid4()))
else:
tool_id = str(uuid.uuid4())
interrupt_message = AIMessage(
content="", tool_calls=[{"id": tool_id, "name": action, "args": args or {}}]
)
interrupt_values = {"action": action, "args": args or {}}
response = interrupt(
{
"__copilotkit_interrupt_value__": interrupt_values,
"__copilotkit_messages__": [interrupt_message],
}
)
if isinstance(response, str):
answer = response
elif isinstance(response, dict):
answer = json.dumps(response)
elif isinstance(response, list):
answer = response[-1].content
else:
answer = str(response)
return answer, response
@@ -0,0 +1,298 @@
import json
import logging
from typing import Dict, Any, List, Optional, Union, AsyncGenerator
from enum import Enum
from .exc import CopilotKitMisuseError
logger = logging.getLogger(__name__)
from ag_ui_langgraph import LangGraphAgent
from ag_ui.core import (
EventType,
CustomEvent,
TextMessageStartEvent,
TextMessageContentEvent,
TextMessageEndEvent,
ToolCallStartEvent,
ToolCallArgsEvent,
ToolCallEndEvent,
StateSnapshotEvent,
)
from langgraph.graph.state import CompiledStateGraph
from langchain_core.runnables import RunnableConfig
try:
from langchain.schema import BaseMessage
except ImportError:
# Langchain >= 1.0.0
from langchain_core.messages import BaseMessage
class CustomEventNames(Enum):
"""Custom event names for CopilotKit"""
ManuallyEmitMessage = "copilotkit_manually_emit_message"
ManuallyEmitToolCall = "copilotkit_manually_emit_tool_call"
ManuallyEmitState = "copilotkit_manually_emit_intermediate_state"
class LangGraphEventTypes(Enum):
"""LangGraph event types"""
OnChatModelStream = "on_chat_model_stream"
OnCustomEvent = "on_custom_event"
class PredictStateTool:
def __init__(self, tool: str, state_key: str, tool_argument: str):
self.tool = tool
self.state_key = state_key
self.tool_argument = tool_argument
State = Dict[str, Any]
SchemaKeys = Dict[str, List[str]]
TextMessageEvents = Union[
TextMessageStartEvent, TextMessageContentEvent, TextMessageEndEvent
]
ToolCallEvents = Union[ToolCallStartEvent, ToolCallArgsEvent, ToolCallEndEvent]
class LangGraphAGUIAgent(LangGraphAgent):
def __init__(
self,
*,
name: str,
graph: CompiledStateGraph,
description: Optional[str] = None,
config: Union[Optional[RunnableConfig], dict] = None,
):
super().__init__(name=name, graph=graph, description=description, config=config)
self.constant_schema_keys = self.constant_schema_keys + ["copilotkit"]
def _dispatch_event(self, event) -> str:
"""Override the dispatch event method to handle custom CopilotKit events and filtering.
Note: Returns None for filtered events (which violates the str return type annotation,
but the base class also violates it by returning event objects). The None values are
filtered out in run() before reaching the encoder.
"""
if event.type == EventType.CUSTOM:
custom_event = event
if custom_event.name == CustomEventNames.ManuallyEmitMessage.value:
# Emit the message events
super()._dispatch_event(
TextMessageStartEvent(
type=EventType.TEXT_MESSAGE_START,
role="assistant",
message_id=custom_event.value["message_id"],
raw_event=event,
)
)
super()._dispatch_event(
TextMessageContentEvent(
type=EventType.TEXT_MESSAGE_CONTENT,
message_id=custom_event.value["message_id"],
delta=custom_event.value["message"],
raw_event=event,
)
)
super()._dispatch_event(
TextMessageEndEvent(
type=EventType.TEXT_MESSAGE_END,
message_id=custom_event.value["message_id"],
raw_event=event,
)
)
return super()._dispatch_event(event)
if custom_event.name == CustomEventNames.ManuallyEmitToolCall.value:
value = custom_event.value
if not isinstance(value, dict):
raise CopilotKitMisuseError(
f"ManuallyEmitToolCall event 'value' must be a dict, got {type(value).__name__}"
)
tool_call_id = value.get("id")
tool_call_name = value.get("name")
tool_call_args = value.get("args")
if not isinstance(tool_call_id, str) or not tool_call_id.strip():
raise CopilotKitMisuseError(
f"ManuallyEmitToolCall event missing valid 'id': got {type(tool_call_id).__name__}"
)
if not isinstance(tool_call_name, str) or not tool_call_name.strip():
raise CopilotKitMisuseError(
f"ManuallyEmitToolCall event missing valid 'name': got {type(tool_call_name).__name__}"
)
if tool_call_args is None:
raise CopilotKitMisuseError(
f"ManuallyEmitToolCall event missing 'args' for tool_call_id={tool_call_id}"
)
try:
delta = (
tool_call_args
if isinstance(tool_call_args, str)
else json.dumps(tool_call_args)
)
except (TypeError, ValueError) as e:
raise CopilotKitMisuseError(
f"ManuallyEmitToolCall 'args' is not JSON-serializable for tool_call_id={tool_call_id}: {e}"
) from e
dispatched_start = False
end_dispatched = False
try:
super()._dispatch_event(
ToolCallStartEvent(
type=EventType.TOOL_CALL_START,
tool_call_id=tool_call_id,
tool_call_name=tool_call_name,
parent_message_id=tool_call_id,
raw_event=event,
)
)
dispatched_start = True
super()._dispatch_event(
ToolCallArgsEvent(
type=EventType.TOOL_CALL_ARGS,
tool_call_id=tool_call_id,
delta=delta,
raw_event=event,
)
)
super()._dispatch_event(
ToolCallEndEvent(
type=EventType.TOOL_CALL_END,
tool_call_id=tool_call_id,
raw_event=event,
)
)
end_dispatched = True
except Exception:
if dispatched_start and not end_dispatched:
try:
super()._dispatch_event(
ToolCallEndEvent(
type=EventType.TOOL_CALL_END,
tool_call_id=tool_call_id,
raw_event=event,
)
)
except Exception:
logger.error(
"Failed to emit compensating TOOL_CALL_END for %s",
tool_call_id,
exc_info=True,
)
raise
return super()._dispatch_event(event)
if custom_event.name == CustomEventNames.ManuallyEmitState.value:
self.active_run["manually_emitted_state"] = custom_event.value
return super()._dispatch_event(
StateSnapshotEvent(
type=EventType.STATE_SNAPSHOT,
snapshot=self.get_state_snapshot(
self.active_run["manually_emitted_state"]
),
raw_event=event,
)
)
if custom_event.name == "copilotkit_exit":
return super()._dispatch_event(
CustomEvent(
type=EventType.CUSTOM,
name="Exit",
value=True,
raw_event=event,
)
)
# Handle filtering based on metadata for text messages and tool calls
raw_event = getattr(event, "raw_event", None)
if raw_event:
is_message_event = event.type in [
EventType.TEXT_MESSAGE_START,
EventType.TEXT_MESSAGE_CONTENT,
EventType.TEXT_MESSAGE_END,
]
is_tool_event = event.type in [
EventType.TOOL_CALL_START,
EventType.TOOL_CALL_ARGS,
EventType.TOOL_CALL_END,
]
# Handle both dict and object cases for raw_event
# See: https://github.com/CopilotKit/CopilotKit/issues/2066
metadata = (
raw_event.get("metadata", {})
if isinstance(raw_event, dict)
else getattr(raw_event, "metadata", {})
) or {}
if "copilotkit:emit-tool-calls" in metadata:
if metadata["copilotkit:emit-tool-calls"] is False and is_tool_event:
return None # Don't dispatch this event
if "copilotkit:emit-messages" in metadata:
if metadata["copilotkit:emit-messages"] is False and is_message_event:
return None # Don't dispatch this event
return super()._dispatch_event(event)
async def run(self, input):
"""Override run to filter out None events from _dispatch_event filtering."""
async for event in super().run(input):
if event is not None:
yield event
async def _handle_single_event(
self, event: Any, state: State
) -> AsyncGenerator[str, None]:
"""Override to add custom event processing for PredictState events"""
# First, check if this is a raw event that should generate a PredictState event
if event.get("event") == LangGraphEventTypes.OnChatModelStream.value:
predict_state_metadata = event.get("metadata", {}).get(
"copilotkit:emit-intermediate-state", None
)
if predict_state_metadata is not None:
event["metadata"]["predict_state"] = predict_state_metadata
# Call the parent method to handle all other events
async for event_str in super()._handle_single_event(event, state):
yield event_str
def langgraph_default_merge_state(
self, state: State, messages: List[BaseMessage], input: Any
) -> State:
"""Override to add CopilotKit actions to the state"""
merged_state = super().langgraph_default_merge_state(state, messages, input)
# Extract tools from the merged state and add them as CopilotKit actions
agui_properties = merged_state.get("ag-ui", {}) or merged_state
return {
**merged_state,
"copilotkit": {
"actions": [
a.model_dump() if hasattr(a, "model_dump") else a
for a in agui_properties.get("tools", [])
],
"context": [
c.model_dump() if hasattr(c, "model_dump") else c
for c in agui_properties.get("context", [])
],
},
}
def dict_repr(self):
"""Return dictionary representation of the agent"""
return {
"name": self.name,
"description": self.description or "",
"type": "langgraph_agui",
}
+27
View File
@@ -0,0 +1,27 @@
"""
Logging setup for CopilotKit.
"""
import logging
import os
import sys
def get_logger(name: str):
"""
Get a logger with the given name.
"""
logger = logging.getLogger(name)
log_level = os.getenv("LOG_LEVEL")
if log_level:
logger.setLevel(log_level.upper())
return logger
def bold(text: str) -> str:
"""
Bold the given text.
"""
if hasattr(sys.stdout, "isatty") and sys.stdout.isatty():
return f"\033[1m{text}\033[0m"
return text
+61
View File
@@ -0,0 +1,61 @@
"""Parameter classes for CopilotKit"""
from typing import TypedDict, Optional, Literal, List, Union, cast, Any
from typing_extensions import NotRequired
class SimpleParameter(TypedDict):
"""Simple parameter class"""
name: str
description: NotRequired[str]
required: NotRequired[bool]
type: NotRequired[Literal["number", "boolean", "number[]", "boolean[]"]]
class ObjectParameter(TypedDict):
"""Object parameter class"""
name: str
description: NotRequired[str]
required: NotRequired[bool]
type: Literal["object", "object[]"]
attributes: List["Parameter"]
class StringParameter(TypedDict):
"""String parameter class"""
name: str
description: NotRequired[str]
required: NotRequired[bool]
type: Literal["string", "string[]"]
enum: NotRequired[List[str]]
Parameter = Union[SimpleParameter, ObjectParameter, StringParameter]
def normalize_parameters(parameters: Optional[List[Parameter]]) -> List[Parameter]:
"""Normalize the parameters to ensure they have the correct type and format."""
if parameters is None:
return []
return [_normalize_parameter(parameter) for parameter in parameters]
def _normalize_parameter(parameter: Parameter) -> Parameter:
"""Normalize a parameter to ensure it has the correct type and format."""
if not "type" in parameter:
cast(Any, parameter)["type"] = "string"
if not "required" in parameter:
parameter["required"] = True
if not "description" in parameter:
parameter["description"] = ""
if "type" in parameter and (
parameter["type"] == "object" or parameter["type"] == "object[]"
):
cast(Any, parameter)["attributes"] = normalize_parameters(
parameter.get("attributes")
)
return parameter
+305
View File
@@ -0,0 +1,305 @@
"""
CopilotKit Protocol
"""
import json
from enum import Enum
from typing import Union, Optional
from typing_extensions import TypedDict, Literal, Any, Dict
class RuntimeEventTypes(Enum):
"""CopilotKit Runtime Event Types"""
TEXT_MESSAGE_START = "TextMessageStart"
TEXT_MESSAGE_CONTENT = "TextMessageContent"
TEXT_MESSAGE_END = "TextMessageEnd"
ACTION_EXECUTION_START = "ActionExecutionStart"
ACTION_EXECUTION_ARGS = "ActionExecutionArgs"
ACTION_EXECUTION_END = "ActionExecutionEnd"
ACTION_EXECUTION_RESULT = "ActionExecutionResult"
AGENT_STATE_MESSAGE = "AgentStateMessage"
META_EVENT = "MetaEvent"
RUN_STARTED = "RunStarted"
RUN_FINISHED = "RunFinished"
RUN_ERROR = "RunError"
NODE_STARTED = "NodeStarted"
NODE_FINISHED = "NodeFinished"
class RuntimeMetaEventName(Enum):
"""Runtime Meta Event Name"""
LANG_GRAPH_INTERRUPT_EVENT = "LangGraphInterruptEvent"
PREDICT_STATE = "PredictState"
EXIT = "Exit"
class TextMessageStart(TypedDict):
"""Text Message Start Event"""
type: Literal[RuntimeEventTypes.TEXT_MESSAGE_START]
messageId: str
parentMessageId: Optional[str]
class TextMessageContent(TypedDict):
"""Text Message Content Event"""
type: Literal[RuntimeEventTypes.TEXT_MESSAGE_CONTENT]
messageId: str
content: str
class TextMessageEnd(TypedDict):
"""Text Message End Event"""
type: Literal[RuntimeEventTypes.TEXT_MESSAGE_END]
messageId: str
class ActionExecutionStart(TypedDict):
"""Action Execution Start Event"""
type: Literal[RuntimeEventTypes.ACTION_EXECUTION_START]
actionExecutionId: str
actionName: str
parentMessageId: Optional[str]
class ActionExecutionArgs(TypedDict):
"""Action Execution Args Event"""
type: Literal[RuntimeEventTypes.ACTION_EXECUTION_ARGS]
actionExecutionId: str
args: str
class ActionExecutionEnd(TypedDict):
"""Action Execution End Event"""
type: Literal[RuntimeEventTypes.ACTION_EXECUTION_END]
actionExecutionId: str
class ActionExecutionResult(TypedDict):
"""Action Execution Result Event"""
type: Literal[RuntimeEventTypes.ACTION_EXECUTION_RESULT]
actionName: str
actionExecutionId: str
result: str
class AgentStateMessage(TypedDict):
"""Agent State Message Event"""
type: Literal[RuntimeEventTypes.AGENT_STATE_MESSAGE]
threadId: str
agentName: str
nodeName: str
runId: str
active: bool
role: str
state: str
running: bool
class MetaEvent(TypedDict):
"""Meta Event"""
type: Literal[RuntimeEventTypes.META_EVENT]
name: RuntimeMetaEventName
value: Any
class RunStarted(TypedDict):
"""Run Started Event"""
type: Literal[RuntimeEventTypes.RUN_STARTED]
state: Dict[str, Any]
class RunFinished(TypedDict):
"""Run Finished Event"""
type: Literal[RuntimeEventTypes.RUN_FINISHED]
state: Dict[str, Any]
class RunError(TypedDict):
"""Run Error Event"""
type: Literal[RuntimeEventTypes.RUN_ERROR]
error: Any
class NodeStarted(TypedDict):
"""Node Started Event"""
type: Literal[RuntimeEventTypes.NODE_STARTED]
node_name: str
state: Dict[str, Any]
class NodeFinished(TypedDict):
"""Node Finished Event"""
type: Literal[RuntimeEventTypes.NODE_FINISHED]
node_name: str
state: Dict[str, Any]
RuntimeProtocolEvent = Union[
TextMessageStart,
TextMessageContent,
TextMessageEnd,
ActionExecutionStart,
ActionExecutionArgs,
ActionExecutionEnd,
ActionExecutionResult,
AgentStateMessage,
MetaEvent,
]
RuntimeLifecycleEvent = Union[
RunStarted,
RunFinished,
RunError,
NodeStarted,
NodeFinished,
]
RuntimeEvent = Union[
RuntimeProtocolEvent,
RuntimeLifecycleEvent,
]
class PredictStateConfig(TypedDict):
"""
Predict State Config
"""
tool_name: str
tool_argument: Optional[str]
def text_message_start(
*, message_id: str, parent_message_id: Optional[str] = None
) -> TextMessageStart:
"""Utility function to create a text message start event"""
return {
"type": RuntimeEventTypes.TEXT_MESSAGE_START,
"messageId": message_id,
"parentMessageId": parent_message_id,
}
def text_message_content(*, message_id: str, content: str) -> TextMessageContent:
"""Utility function to create a text message content event"""
return {
"type": RuntimeEventTypes.TEXT_MESSAGE_CONTENT,
"messageId": message_id,
"content": content,
}
def text_message_end(*, message_id: str) -> TextMessageEnd:
"""Utility function to create a text message end event"""
return {"type": RuntimeEventTypes.TEXT_MESSAGE_END, "messageId": message_id}
def action_execution_start(
*,
action_execution_id: str,
action_name: str,
parent_message_id: Optional[str] = None,
) -> ActionExecutionStart:
"""Utility function to create an action execution start event"""
return {
"type": RuntimeEventTypes.ACTION_EXECUTION_START,
"actionExecutionId": action_execution_id,
"actionName": action_name,
"parentMessageId": parent_message_id,
}
def action_execution_args(
*, action_execution_id: str, args: str
) -> ActionExecutionArgs:
"""Utility function to create an action execution args event"""
return {
"type": RuntimeEventTypes.ACTION_EXECUTION_ARGS,
"actionExecutionId": action_execution_id,
"args": args,
}
def action_execution_end(*, action_execution_id: str) -> ActionExecutionEnd:
"""Utility function to create an action execution end event"""
return {
"type": RuntimeEventTypes.ACTION_EXECUTION_END,
"actionExecutionId": action_execution_id,
}
def action_execution_result(
*, action_name: str, action_execution_id: str, result: str
) -> ActionExecutionResult:
"""Utility function to create an action execution result event"""
return {
"type": RuntimeEventTypes.ACTION_EXECUTION_RESULT,
"actionName": action_name,
"actionExecutionId": action_execution_id,
"result": result,
}
def agent_state_message( # pylint: disable=too-many-arguments
*,
thread_id: str,
agent_name: str,
node_name: str,
run_id: str,
active: bool,
role: str,
state: str,
running: bool,
) -> AgentStateMessage:
"""Utility function to create an agent state message event"""
return {
"type": RuntimeEventTypes.AGENT_STATE_MESSAGE,
"threadId": thread_id,
"agentName": agent_name,
"nodeName": node_name,
"runId": run_id,
"active": active,
"role": role,
"state": state,
"running": running,
}
def meta_event(*, name: RuntimeMetaEventName, value: Any) -> MetaEvent:
"""Utility function to create a meta event"""
return {"type": RuntimeEventTypes.META_EVENT, "name": name, "value": value}
def emit_runtime_events(*events: RuntimeProtocolEvent) -> str:
"""Emit a list of runtime events"""
def serialize_event(event):
# Convert enum values to their string representation
if isinstance(event, dict):
return {
k: (v.value if isinstance(v, Enum) else v) for k, v in event.items()
}
return event
return "\n".join(json.dumps(serialize_event(event)) for event in events) + "\n"
def emit_runtime_event(event: RuntimeProtocolEvent) -> str:
"""Emit a single runtime event"""
return emit_runtime_events(event)
View File
+346
View File
@@ -0,0 +1,346 @@
"""
CopilotKit Run Loop
"""
import asyncio
import contextvars
import json
import traceback
from typing import Callable
from pydantic import BaseModel
from typing_extensions import Any, Dict, Optional, List, TypedDict, cast
from partialjson.json_parser import JSONParser as PartialJSONParser
from .protocol import (
RuntimeEvent,
RuntimeEventTypes,
RuntimeMetaEventName,
emit_runtime_event,
emit_runtime_events,
agent_state_message,
AgentStateMessage,
PredictStateConfig,
RuntimeProtocolEvent,
)
async def yield_control():
"""
Yield control to the event loop.
"""
loop = asyncio.get_running_loop()
future = loop.create_future()
loop.call_soon(future.set_result, None)
await future
class CopilotKitRunExecution(TypedDict):
"""
CopilotKit Run Execution
"""
thread_id: str
agent_name: str
run_id: str
should_exit: bool
node_name: str
is_finished: bool
predict_state_configuration: Dict[str, PredictStateConfig]
predicted_state: Dict[str, Any]
argument_buffer: str
current_tool_call: Optional[str]
state: Dict[str, Any]
_CONTEXT_QUEUE = contextvars.ContextVar("queue", default=None)
_CONTEXT_EXECUTION = contextvars.ContextVar("execution", default=None)
def get_context_queue() -> asyncio.Queue:
"""
Retrieve the queue from this task's context.
"""
q = _CONTEXT_QUEUE.get()
if q is None:
raise RuntimeError("No context queue is set!")
return q
def set_context_queue(q: asyncio.Queue) -> contextvars.Token:
"""
Set the queue in this task's context.
"""
token = _CONTEXT_QUEUE.set(cast(Any, q))
return token
def reset_context_queue(token: contextvars.Token):
"""
Reset the queue in this task's context.
"""
_CONTEXT_QUEUE.reset(token)
def get_context_execution() -> CopilotKitRunExecution:
"""
Get the execution from this task's context.
"""
return cast(CopilotKitRunExecution, _CONTEXT_EXECUTION.get())
def set_context_execution(execution: CopilotKitRunExecution) -> contextvars.Token:
"""
Set the execution in this task's context.
"""
token = _CONTEXT_EXECUTION.set(cast(Any, execution))
return token
def reset_context_execution(token: contextvars.Token):
"""
Reset the execution in this task's context.
"""
_CONTEXT_EXECUTION.reset(token)
async def queue_put(*events: RuntimeEvent, priority: bool = False):
"""
Put an event in the queue.
"""
if not priority:
# yield control so that priority events can be processed first
await yield_control()
q = get_context_queue()
for event in events:
await q.put(event)
# yield control so that the reader can process the event
await yield_control()
def _to_dict_if_pydantic(obj):
if isinstance(obj, BaseModel):
return obj.model_dump()
return obj
def _filter_state(
*, state: Dict[str, Any], exclude_keys: Optional[List[str]] = None
) -> Dict[str, Any]:
"""Filter out messages and id from the state"""
state = _to_dict_if_pydantic(state)
exclude_keys = exclude_keys or ["messages", "id"]
return {k: v for k, v in state.items() if k not in exclude_keys}
async def copilotkit_run(fn: Callable, *, execution: CopilotKitRunExecution):
"""
Run a task with a local queue.
"""
local_queue = asyncio.Queue()
token_queue = set_context_queue(local_queue)
token_execution = set_context_execution(execution)
task = asyncio.create_task(fn())
try:
while True:
event = await local_queue.get()
local_queue.task_done()
json_lines = handle_runtime_event(event=event, execution=execution)
if json_lines is not None:
yield json_lines
if execution["is_finished"]:
break
# return control to the containing run loop to send events
await yield_control()
await task
finally:
reset_context_queue(token_queue)
reset_context_execution(token_execution)
def handle_runtime_event(
*, event: RuntimeEvent, execution: CopilotKitRunExecution
) -> Optional[str]:
"""
Handle a runtime event.
"""
if event["type"] in [
RuntimeEventTypes.TEXT_MESSAGE_START,
RuntimeEventTypes.TEXT_MESSAGE_CONTENT,
RuntimeEventTypes.TEXT_MESSAGE_END,
RuntimeEventTypes.ACTION_EXECUTION_START,
RuntimeEventTypes.ACTION_EXECUTION_ARGS,
RuntimeEventTypes.ACTION_EXECUTION_END,
RuntimeEventTypes.ACTION_EXECUTION_RESULT,
RuntimeEventTypes.AGENT_STATE_MESSAGE,
]:
events: List[RuntimeProtocolEvent] = [cast(RuntimeProtocolEvent, event)]
if event["type"] in [
RuntimeEventTypes.ACTION_EXECUTION_START,
RuntimeEventTypes.ACTION_EXECUTION_ARGS,
]:
message = predict_state(
thread_id=execution["thread_id"],
agent_name=execution["agent_name"],
run_id=execution["run_id"],
event=event,
execution=execution,
)
if message is not None:
events.append(message)
return emit_runtime_events(*events)
if event["type"] == RuntimeEventTypes.META_EVENT:
if event["name"] == RuntimeMetaEventName.PREDICT_STATE:
execution["predict_state_configuration"] = event["value"]
return None
if event["name"] == RuntimeMetaEventName.EXIT:
execution["should_exit"] = event["value"]
return None
return None
if event["type"] == RuntimeEventTypes.RUN_STARTED:
execution["state"] = event["state"]
return None
if event["type"] == RuntimeEventTypes.NODE_STARTED:
execution["node_name"] = event["node_name"]
execution["state"] = event["state"]
return emit_runtime_event(
agent_state_message(
thread_id=execution["thread_id"],
agent_name=execution["agent_name"],
node_name=execution["node_name"],
run_id=execution["run_id"],
active=True,
role="assistant",
state=json.dumps(_filter_state(state=execution["state"])),
running=True,
)
)
if event["type"] == RuntimeEventTypes.NODE_FINISHED:
# reset the predict state configuration at the end of the method execution
execution["predict_state_configuration"] = {}
execution["current_tool_call"] = None
execution["argument_buffer"] = ""
execution["predicted_state"] = {}
execution["state"] = event["state"]
return emit_runtime_event(
agent_state_message(
thread_id=execution["thread_id"],
agent_name=execution["agent_name"],
node_name=execution["node_name"],
run_id=execution["run_id"],
active=False,
role="assistant",
state=json.dumps(_filter_state(state=execution["state"])),
running=True,
)
)
if event["type"] == RuntimeEventTypes.RUN_FINISHED:
execution["is_finished"] = True
return None
if event["type"] == RuntimeEventTypes.RUN_ERROR:
print("Flow execution error", flush=True)
error_info = event["error"]
if isinstance(error_info, Exception):
# If it's an exception, print the traceback
print("Exception occurred:", flush=True)
print(
"".join(
traceback.format_exception(
None, error_info, error_info.__traceback__
)
),
flush=True,
)
else:
# Otherwise, assume it's a string and print it
print(error_info, flush=True)
execution["is_finished"] = True
return None
def predict_state(
*,
thread_id: str,
agent_name: str,
run_id: str,
event: Any,
execution: CopilotKitRunExecution,
) -> Optional[AgentStateMessage]:
"""Predict the state"""
if event["type"] == RuntimeEventTypes.ACTION_EXECUTION_START:
execution["current_tool_call"] = event["actionName"]
execution["argument_buffer"] = ""
elif event["type"] == RuntimeEventTypes.ACTION_EXECUTION_ARGS:
execution["argument_buffer"] += event["args"]
tool_names = [
config.get("tool_name")
for config in execution["predict_state_configuration"].values()
]
if execution["current_tool_call"] not in tool_names:
return None
current_arguments = {}
try:
current_arguments = PartialJSONParser().parse(execution["argument_buffer"])
except: # pylint: disable=bare-except
return None
emit_update = False
for k, v in execution["predict_state_configuration"].items():
if v["tool_name"] == execution["current_tool_call"]:
tool_argument = v.get("tool_argument")
if tool_argument is not None:
argument_value = current_arguments.get(tool_argument)
if argument_value is not None:
execution["predicted_state"][k] = argument_value
emit_update = True
else:
execution["predicted_state"][k] = current_arguments
emit_update = True
if emit_update:
return agent_state_message(
thread_id=thread_id,
agent_name=agent_name,
node_name=execution["node_name"],
run_id=run_id,
active=True,
role="assistant",
state=json.dumps(
_filter_state(
state={
**(
execution["state"].model_dump()
if isinstance(execution["state"], BaseModel)
else execution["state"]
),
**execution["predicted_state"],
}
)
),
running=True,
)
return None
+397
View File
@@ -0,0 +1,397 @@
"""CopilotKit SDK"""
import warnings
from importlib import metadata
from pprint import pformat
from typing import List, Callable, Union, Optional, Any, Coroutine
from typing_extensions import TypedDict, Tuple, cast, Mapping
from .agent import Agent, AgentDict
from .action import Action, ActionDict, ActionResultDict
from .types import Message, MetaEvent
from .exc import (
ActionNotFoundException,
AgentNotFoundException,
ActionExecutionException,
AgentExecutionException,
)
from .logging import get_logger, bold
try:
__version__ = metadata.version(cast(str, __package__))
except metadata.PackageNotFoundError:
# Case where package metadata is not available.
__version__ = ""
del metadata # optional, avoids polluting the results of dir(__package__)
COPILOTKIT_SDK_VERSION = __version__
logger = get_logger(__name__)
class InfoDict(TypedDict):
"""
Info dictionary
"""
sdkVersion: str
actions: List[ActionDict]
agents: List[AgentDict]
class CopilotKitContext(TypedDict):
"""
CopilotKit Context
Parameters
----------
properties : Any
The properties provided to the frontend via `<CopilotKit properties={...} />`
frontend_url : Optional[str]
The current URL of the frontend
headers : Mapping[str, str]
The headers of the request
"""
properties: Any
frontend_url: Optional[str]
headers: Mapping[str, str]
# Alias for backwards compatibility
CopilotKitSDKContext = CopilotKitContext
class CopilotKitRemoteEndpoint:
"""
CopilotKitRemoteEndpoint lets you connect actions and agents written in Python to your
CopilotKit application.
To install CopilotKit for Python, run:
```bash
pip install copilotkit
# or to include crewai
pip install copilotkit[crewai]
```
## Adding actions
In this example, we provide a simple action to the Copilot:
```python
from copilotkit import CopilotKitRemoteEndpoint, Action
sdk = CopilotKitRemoteEndpoint(
actions=[
Action(
name="greet_user",
handler=greet_user_handler,
description="Greet the user",
parameters=[
{
"name": "name",
"type": "string",
"description": "The name of the user"
}
]
)
]
)
```
You can also dynamically build actions by providing a callable that returns a list of actions.
In this example, we use "name" from the `properties` object to parameterize the action handler.
```python
from copilotkit import CopilotKitRemoteEndpoint, Action
sdk = CopilotKitRemoteEndpoint(
actions=lambda context: [
Action(
name="greet_user",
handler=make_greet_user_handler(context["properties"]["name"]),
description="Greet the user"
)
]
)
```
Using the same approach, you can restrict the actions available to the Copilot:
```python
from copilotkit import CopilotKitRemoteEndpoint, Action
sdk = CopilotKitRemoteEndpoint(
actions=lambda context: (
[action_a, action_b] if is_admin(context["properties"]["token"]) else [action_a]
)
)
```
## Adding agents
Serving agents works in a similar way to serving actions:
```python
from copilotkit import CopilotKitRemoteEndpoint, LangGraphAGUIAgent
from my_agent.agent import graph
sdk = CopilotKitRemoteEndpoint(
agents=[
LangGraphAGUIAgent(
name="email_agent",
description="This agent sends emails",
graph=graph,
)
]
)
```
To dynamically build agents, provide a callable that returns a list of agents:
```python
from copilotkit import CopilotKitRemoteEndpoint, LangGraphAGUIAgent
from my_agent.agent import graph
sdk = CopilotKitRemoteEndpoint(
agents=lambda context: [
LangGraphAGUIAgent(
name="email_agent",
description="This agent sends emails",
graph=graph,
langgraph_config={
"token": context["properties"]["token"]
}
)
]
)
```
To restrict the agents available to the Copilot, simply return a different list of agents based on the `context`:
```python
from copilotkit import CopilotKitRemoteEndpoint
from my_agents import agent_a, agent_b, is_admin
sdk = CopilotKitRemoteEndpoint(
agents=lambda context: (
[agent_a, agent_b] if is_admin(context["properties"]["token"]) else [agent_a]
)
)
```
## Serving the CopilotKit SDK
To serve the CopilotKit SDK, you can use the `add_fastapi_endpoint` function from the `copilotkit.integrations.fastapi` module:
```python
from copilotkit.integrations.fastapi import add_fastapi_endpoint
from fastapi import FastAPI
app = FastAPI()
sdk = CopilotKitRemoteEndpoint(...)
add_fastapi_endpoint(app, sdk, "/copilotkit")
def main():
uvicorn.run(
"your_package:app",
host="0.0.0.0",
port=8000,
reload=True,
)
```
Parameters
----------
actions : Optional[Union[List[Action], Callable[[CopilotKitContext], List[Action]]]]
The actions to make available to the Copilot.
agents : Optional[Union[List[Agent], Callable[[CopilotKitContext], List[Agent]]]]
The agents to make available to the Copilot.
"""
def __init__(
self,
*,
actions: Optional[
Union[List[Action], Callable[[CopilotKitContext], List[Action]]]
] = None,
agents: Optional[
Union[List[Agent], Callable[[CopilotKitContext], List[Agent]]]
] = None,
):
self.agents = agents or []
self.actions = actions or []
def info(self, *, context: CopilotKitContext) -> InfoDict:
"""
Returns information about available actions and agents
"""
actions = self.actions(context) if callable(self.actions) else self.actions
agents = self.agents(context) if callable(self.agents) else self.agents
actions_list = [action.dict_repr() for action in actions]
agents_list = [agent.dict_repr() for agent in agents]
self._log_request_info(
title="Handling info request:",
data=[
("Context", context),
("Actions", actions_list),
("Agents", agents_list),
],
)
return {
"actions": actions_list,
"agents": agents_list,
"sdkVersion": COPILOTKIT_SDK_VERSION,
}
def _get_action(
self,
*,
context: CopilotKitContext,
name: str,
) -> Action:
"""
Get an action by name
"""
actions = self.actions(context) if callable(self.actions) else self.actions
action = next((action for action in actions if action.name == name), None)
if action is None:
raise ActionNotFoundException(name)
return action
def execute_action(
self,
*,
context: CopilotKitContext,
name: str,
arguments: dict,
) -> Coroutine[Any, Any, ActionResultDict]:
"""
Execute an action
"""
action = self._get_action(context=context, name=name)
self._log_request_info(
title="Handling execute action request:",
data=[
("Context", context),
("Action", action.dict_repr()),
("Arguments", arguments),
],
)
try:
result = action.execute(arguments=arguments)
return result
except Exception as error:
raise ActionExecutionException(name, error) from error
def execute_agent( # pylint: disable=too-many-arguments
self,
*,
context: CopilotKitContext,
name: str,
thread_id: str,
state: dict,
config: Optional[dict] = None,
messages: List[Message],
actions: List[ActionDict],
node_name: str,
meta_events: Optional[List[MetaEvent]] = None,
) -> Any:
"""
Execute an agent
"""
agents = self.agents(context) if callable(self.agents) else self.agents
agent = next((agent for agent in agents if agent.name == name), None)
if agent is None:
raise AgentNotFoundException(name)
self._log_request_info(
title="Handling execute agent request:",
data=[
("Context", context),
("Agent", agent.dict_repr()),
("Thread ID", thread_id),
("Node Name", node_name),
("State", state),
("Config", config),
("Messages", messages),
("Actions", actions),
("MetaEvents", meta_events),
],
)
try:
return agent.execute(
thread_id=thread_id,
node_name=node_name,
state=state,
config=config,
messages=messages,
actions=actions,
meta_events=meta_events,
)
except Exception as error:
raise AgentExecutionException(name, error) from error
async def get_agent_state(
self,
*,
context: CopilotKitContext,
thread_id: str,
name: str,
):
"""
Get agent state
"""
agents = self.agents(context) if callable(self.agents) else self.agents
agent = next((agent for agent in agents if agent.name == name), None)
if agent is None:
raise AgentNotFoundException(name)
self._log_request_info(
title="Handling get agent state request:",
data=[
("Context", context),
("Agent", agent.dict_repr()),
("Thread ID", thread_id),
],
)
try:
return await agent.get_state(thread_id=thread_id)
except Exception as error:
raise AgentExecutionException(name, error) from error
def _log_request_info(self, title: str, data: List[Tuple[str, Any]]):
"""
Log request info
"""
logger.info(bold(title))
logger.info("--------------------------")
for key, value in data:
logger.info(bold(key + ":"))
logger.info(pformat(value))
logger.info("--------------------------")
# Alias for backwards compatibility
class CopilotKitSDK(CopilotKitRemoteEndpoint):
"""Deprecated: Use CopilotKitRemoteEndpoint instead. This class will be removed in a future version."""
def __init__(self, *args, **kwargs):
warnings.warn(
"CopilotKitSDK is deprecated since version 0.1.31. "
"Use CopilotKitRemoteEndpoint instead.",
DeprecationWarning,
stacklevel=2,
)
super().__init__(*args, **kwargs)
+59
View File
@@ -0,0 +1,59 @@
"""State for CopilotKit"""
from typing import TypedDict
from enum import Enum
from typing_extensions import NotRequired
class MessageRole(Enum):
"""Message role"""
ASSISTANT = "assistant"
SYSTEM = "system"
USER = "user"
class Message(TypedDict):
"""Message"""
id: str
createdAt: str
class TextMessage(Message):
"""Text message"""
parentMessageId: NotRequired[str]
role: MessageRole
content: str
class ActionExecutionMessage(Message):
"""Action execution message"""
parentMessageId: NotRequired[str]
name: str
arguments: dict
class ResultMessage(Message):
"""Result message"""
actionExecutionId: str
actionName: str
result: str
class IntermediateStateConfig(TypedDict):
"""Intermediate state config"""
state_key: str
tool: str
tool_argument: NotRequired[str]
class MetaEvent(TypedDict):
"""Type definition for meta events"""
name: str
response: NotRequired[str]
+5
View File
@@ -0,0 +1,5 @@
def filter_by_schema_keys(obj, schema):
try:
return {k: v for k, v in obj.items() if k in schema or k == "messages"}
except Exception:
return obj
+9
View File
@@ -0,0 +1,9 @@
{
"python_version": "3.11",
"dockerfile_lines": [],
"dependencies": ["."],
"graphs": {
"autotale_ai": "./copilotkit/demos/autotale_ai/agent.py:graph"
},
"env": ".env"
}
+5645
View File
File diff suppressed because it is too large Load Diff
+51
View File
@@ -0,0 +1,51 @@
[tool.poetry]
name = "copilotkit"
version = "0.1.94"
description = "CopilotKit python SDK"
authors = ["Markus Ecker <markus.ecker@gmail.com>"]
license = "MIT"
readme = "README.md"
homepage = "https://copilotkit.ai"
repository = "https://github.com/CopilotKit/CopilotKit/tree/main/sdk-python"
keywords = [
"copilot",
"copilotkit",
"langgraph",
"langchain",
"ai",
"langsmith",
"langserve",
]
classifiers = [
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
"Programming Language :: Python :: 3.14",
]
packages = [{ include = "copilotkit" }]
include = ["copilotkit/*.json", "copilotkit/py.typed"]
[tool.poetry.dependencies]
python = ">=3.10,<3.15"
langgraph = { version = ">=0.3.25,<2" }
langchain = { version = ">=0.3.0" }
crewai = { version = ">=0.118.0", optional = true, python = ">=3.10,<3.14" }
ag-ui-langgraph = { version = ">=0.0.42", extras = ["fastapi"] }
ag-ui-protocol = ">=0.1.15"
fastapi = ">=0.115.0,<1.0.0"
partialjson = "^0.0.8"
toml = "^0.10.2"
# Pin transitive: cp314 wheels first appear in 2.35.0; older versions
# would force a Rust source build that fails on PyO3 0.24 (capped at 3.13).
pydantic-core = ">=2.35.0"
[tool.poetry.extras]
crewai = ["crewai"]
[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[dependency-groups]
dev = ["pytest (>=9.0.2,<10.0.0)", "pytest-asyncio (>=1.0.0,<2.0.0)"]
View File
+661
View File
@@ -0,0 +1,661 @@
"""Tests for LangGraphAGUIAgent — CopilotKit's extension layer on top of AG-UI's base agent.
Covers:
1. Custom event handling (_dispatch_event with CopilotKit-specific custom events)
2. copilotkit state namespace (langgraph_default_merge_state)
3. Unknown custom events pass through without crashing
"""
import json
import pytest
from contextlib import contextmanager
from unittest.mock import MagicMock, patch
from ag_ui.core import (
EventType,
CustomEvent,
TextMessageStartEvent,
TextMessageContentEvent,
TextMessageEndEvent,
ToolCallStartEvent,
ToolCallArgsEvent,
ToolCallEndEvent,
StateSnapshotEvent,
)
from ag_ui_langgraph import LangGraphAgent as AGUIBase
from copilotkit import CopilotKitRemoteEndpoint
from copilotkit.langgraph_agui_agent import (
LangGraphAGUIAgent,
CustomEventNames,
)
# The source code at langgraph_agui_agent.py:134 checks for the literal string
# "copilotkit_exit" (not via CustomEventNames enum — exit was never added to it).
# We intentionally match the source's literal here.
COPILOTKIT_EXIT_EVENT_NAME = "copilotkit_exit"
# The output event name "Exit" is a hardcoded downstream value in the source code
# (langgraph_agui_agent.py:138), distinct from any CustomEventNames constant.
EXIT_OUTPUT_NAME = "Exit"
@pytest.fixture
def agent():
"""Create a LangGraphAGUIAgent with a mocked graph."""
mock_graph = MagicMock()
mock_graph.get_state = MagicMock()
a = LangGraphAGUIAgent(name="test", graph=mock_graph)
a.active_run = {"id": "run-1", "thread_id": "t-1"}
return a
@contextmanager
def track_parent_dispatches(agent):
"""Patch AGUIBase._dispatch_event to record all dispatched events.
Yields the list of captured events. The original dispatch is still called
so that event serialisation behaves normally.
Usage::
with track_parent_dispatches(agent) as dispatched:
agent._dispatch_event(some_event)
assert EventType.TEXT_MESSAGE_START in [e.type for e in dispatched]
"""
dispatched = []
original = AGUIBase._dispatch_event
def _tracking(self_inner, event):
dispatched.append(event)
return original(self_inner, event)
with patch.object(AGUIBase, "_dispatch_event", new=_tracking):
yield dispatched
# ---------- Custom event: ManuallyEmitMessage ----------
class TestManuallyEmitMessage:
"""copilotkit_manually_emit_message → TEXT_MESSAGE_START/CONTENT/END sequence."""
def test_emits_text_message_sequence(self, agent):
"""Should call super()._dispatch_event for start, content, end, and the custom event."""
with track_parent_dispatches(agent) as dispatched:
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitMessage.value,
value={
"message_id": "msg-123",
"message": "Hello world",
"role": "assistant",
},
)
agent._dispatch_event(event)
types = [getattr(e, "type", None) for e in dispatched]
assert EventType.TEXT_MESSAGE_START in types
assert EventType.TEXT_MESSAGE_CONTENT in types
assert EventType.TEXT_MESSAGE_END in types
def test_message_ids_match(self, agent):
"""All emitted events should carry the same message_id from the custom event."""
with track_parent_dispatches(agent) as dispatched:
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitMessage.value,
value={
"message_id": "msg-456",
"message": "test",
"role": "assistant",
},
)
agent._dispatch_event(event)
message_events = [e for e in dispatched if hasattr(e, "message_id")]
for e in message_events:
assert e.message_id == "msg-456"
def test_content_delta_matches(self, agent):
"""TextMessageContentEvent should carry the message text as delta."""
with track_parent_dispatches(agent) as dispatched:
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitMessage.value,
value={
"message_id": "msg-789",
"message": "specific content",
"role": "assistant",
},
)
agent._dispatch_event(event)
content_events = [
e for e in dispatched if e.type == EventType.TEXT_MESSAGE_CONTENT
]
assert len(content_events) == 1
assert content_events[0].delta == "specific content"
# ---------- Custom event: ManuallyEmitToolCall ----------
class TestManuallyEmitToolCall:
"""copilotkit_manually_emit_tool_call → TOOL_CALL_START/ARGS/END sequence."""
def test_emits_tool_call_sequence(self, agent):
"""Should dispatch start, args, end for a tool call."""
with track_parent_dispatches(agent) as dispatched:
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={
"id": "tc-123",
"name": "SearchTool",
"args": {"query": "test"},
},
)
agent._dispatch_event(event)
types = [getattr(e, "type", None) for e in dispatched]
assert EventType.TOOL_CALL_START in types
assert EventType.TOOL_CALL_ARGS in types
assert EventType.TOOL_CALL_END in types
def test_tool_call_ids_match(self, agent):
"""All tool call events should carry the same tool_call_id."""
with track_parent_dispatches(agent) as dispatched:
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={
"id": "tc-456",
"name": "MyTool",
"args": {"key": "val"},
},
)
agent._dispatch_event(event)
tool_events = [e for e in dispatched if hasattr(e, "tool_call_id")]
for e in tool_events:
assert e.tool_call_id == "tc-456"
def test_args_serialized_as_json_when_dict(self, agent):
"""When args is a dict, it should be JSON-serialized in the TOOL_CALL_ARGS event."""
with track_parent_dispatches(agent) as dispatched:
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={
"id": "tc-789",
"name": "MyTool",
"args": {"key": "value"},
},
)
agent._dispatch_event(event)
args_events = [e for e in dispatched if e.type == EventType.TOOL_CALL_ARGS]
assert len(args_events) == 1
assert args_events[0].delta == json.dumps({"key": "value"})
def test_args_passed_as_string_when_string(self, agent):
"""When args is already a string, it should be passed through as-is."""
with track_parent_dispatches(agent) as dispatched:
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={
"id": "tc-str",
"name": "MyTool",
"args": '{"already": "serialized"}',
},
)
agent._dispatch_event(event)
args_events = [e for e in dispatched if e.type == EventType.TOOL_CALL_ARGS]
assert len(args_events) == 1
assert args_events[0].delta == '{"already": "serialized"}'
# ---------- Custom event: ManuallyEmitState ----------
class TestManuallyEmitState:
"""copilotkit_manually_emit_intermediate_state → StateSnapshotEvent."""
def test_emits_state_snapshot(self, agent):
"""Should set active_run['manually_emitted_state'] and dispatch a STATE_SNAPSHOT."""
agent.get_state_snapshot = MagicMock(return_value={"progress": 50})
with track_parent_dispatches(agent) as dispatched:
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitState.value,
value={"progress": 50},
)
agent._dispatch_event(event)
assert agent.active_run["manually_emitted_state"] == {"progress": 50}
types = [getattr(e, "type", None) for e in dispatched]
assert EventType.STATE_SNAPSHOT in types
# ---------- Custom event: copilotkit_exit ----------
class TestCopilotKitExit:
"""copilotkit_exit → CustomEvent with name='Exit'."""
def test_emits_exit_event(self, agent):
"""Should dispatch a CustomEvent with name 'Exit'."""
with track_parent_dispatches(agent) as dispatched:
event = CustomEvent(
type=EventType.CUSTOM,
name=COPILOTKIT_EXIT_EVENT_NAME,
value={},
)
agent._dispatch_event(event)
exit_events = [
e
for e in dispatched
if e.type == EventType.CUSTOM
and getattr(e, "name", None) == EXIT_OUTPUT_NAME
]
assert len(exit_events) == 1
# ---------- Unknown custom events ----------
class TestUnknownCustomEvent:
"""Unknown custom event names should pass through to super() without crashing."""
def test_unknown_custom_event_passes_through(self, agent):
"""An unrecognized custom event name should not raise and should call super()."""
event = CustomEvent(
type=EventType.CUSTOM,
name="some_unknown_event",
value={"data": "test"},
)
result = agent._dispatch_event(event)
assert result is not None
# ---------- Emit filtering (independent of tool calls) ----------
class TestEmitFilteringIndependent:
"""Test that both emit-messages and emit-tool-calls can be set independently."""
def test_emit_messages_false_tool_calls_true(self, agent):
"""emit-messages=False should filter text events but not tool events."""
metadata = {
"copilotkit:emit-messages": False,
"copilotkit:emit-tool-calls": True,
}
text_event = TextMessageContentEvent(
messageId="msg-1",
delta="hello",
rawEvent={"metadata": metadata},
)
tool_event = ToolCallStartEvent(
toolCallId="tc-1",
toolCallName="tool",
rawEvent={"metadata": metadata},
)
assert agent._dispatch_event(text_event) is None
assert agent._dispatch_event(tool_event) is not None
def test_emit_tool_calls_false_messages_true(self, agent):
"""emit-tool-calls=False should filter tool events but not text events."""
metadata = {
"copilotkit:emit-messages": True,
"copilotkit:emit-tool-calls": False,
}
text_event = TextMessageContentEvent(
messageId="msg-1",
delta="hello",
rawEvent={"metadata": metadata},
)
tool_event = ToolCallStartEvent(
toolCallId="tc-1",
toolCallName="tool",
rawEvent={"metadata": metadata},
)
assert agent._dispatch_event(text_event) is not None
assert agent._dispatch_event(tool_event) is None
def test_both_false_filters_both(self, agent):
"""Both flags false should filter both types."""
metadata = {
"copilotkit:emit-messages": False,
"copilotkit:emit-tool-calls": False,
}
text_event = TextMessageContentEvent(
messageId="msg-1",
delta="hello",
rawEvent={"metadata": metadata},
)
tool_event = ToolCallStartEvent(
toolCallId="tc-1",
toolCallName="tool",
rawEvent={"metadata": metadata},
)
assert agent._dispatch_event(text_event) is None
assert agent._dispatch_event(tool_event) is None
# ---------- copilotkit state namespace ----------
class TestLanggraphDefaultMergeState:
"""langgraph_default_merge_state adds copilotkit namespace with actions and context."""
def test_copilotkit_actions_from_agui_tools(self, agent):
"""Tools from ag-ui should appear under copilotkit.actions."""
tools = [{"name": "tool1"}, {"name": "tool2"}]
with patch.object(
AGUIBase,
"langgraph_default_merge_state",
return_value={
"ag-ui": {"tools": tools, "context": []},
"messages": [],
},
):
result = agent.langgraph_default_merge_state({}, [], MagicMock())
assert "copilotkit" in result
assert result["copilotkit"]["actions"] == tools
def test_copilotkit_context_from_agui(self, agent):
"""Context from ag-ui should appear under copilotkit.context."""
context = [{"description": "user info", "value": "test"}]
with patch.object(
AGUIBase,
"langgraph_default_merge_state",
return_value={
"ag-ui": {"tools": [], "context": context},
"messages": [],
},
):
result = agent.langgraph_default_merge_state({}, [], MagicMock())
assert result["copilotkit"]["context"] == context
def test_no_agui_key_no_crash(self, agent):
"""If no ag-ui key in state, should use merged_state as fallback without crashing."""
with patch.object(
AGUIBase,
"langgraph_default_merge_state",
return_value={"messages": [], "tools": [{"name": "fallback"}]},
):
result = agent.langgraph_default_merge_state({}, [], MagicMock())
assert "copilotkit" in result
assert result["copilotkit"]["actions"] == [{"name": "fallback"}]
def test_empty_state_no_crash(self, agent):
"""Completely empty state should not crash."""
with patch.object(
AGUIBase,
"langgraph_default_merge_state",
return_value={},
):
result = agent.langgraph_default_merge_state({}, [], MagicMock())
assert "copilotkit" in result
assert result["copilotkit"]["actions"] == []
assert result["copilotkit"]["context"] == []
def test_preserves_original_state_keys(self, agent):
"""Original keys from super() should be preserved alongside copilotkit."""
with patch.object(
AGUIBase,
"langgraph_default_merge_state",
return_value={
"ag-ui": {"tools": [], "context": []},
"messages": [{"role": "user", "content": "hi"}],
"custom_key": "custom_value",
},
):
result = agent.langgraph_default_merge_state({}, [], MagicMock())
assert result["custom_key"] == "custom_value"
assert result["messages"] == [{"role": "user", "content": "hi"}]
def test_duplicate_tools_not_in_copilotkit_actions(self, agent):
"""Tools should not be duplicated in copilotkit.actions when same name appears in input and state."""
tools = [{"name": "tool_a"}, {"name": "tool_b"}]
with patch.object(
AGUIBase,
"langgraph_default_merge_state",
return_value={
"ag-ui": {"tools": tools, "context": []},
"messages": [],
},
):
result = agent.langgraph_default_merge_state({}, [], MagicMock())
action_names = [a.get("name") for a in result["copilotkit"]["actions"]]
assert action_names.count("tool_a") == 1, (
"Duplicate tool names should not appear in copilotkit.actions"
)
assert action_names.count("tool_b") == 1
def test_copilotkit_actions_ordering_matches_tools(self, agent):
"""copilotkit.actions should have the same ordering as the merged tools list."""
tools = [{"name": "first"}, {"name": "second"}, {"name": "third"}]
with patch.object(
AGUIBase,
"langgraph_default_merge_state",
return_value={
"ag-ui": {"tools": tools, "context": []},
"messages": [],
},
):
result = agent.langgraph_default_merge_state({}, [], MagicMock())
action_names = [a.get("name") for a in result["copilotkit"]["actions"]]
assert action_names == ["first", "second", "third"]
# ---------- Reasoning content preservation ----------
class TestReasoningContentPreservation:
"""Verify that LangGraphAGUIAgent does not drop or mutate reasoning events."""
def test_unknown_custom_event_does_not_suppress_reasoning(self, agent):
"""An unrecognized custom event should pass through (not suppress subsequent reasoning events)."""
# Dispatch a non-CopilotKit custom event — should pass through without crash
unknown_event = CustomEvent(
type=EventType.CUSTOM,
name="some_unknown_custom_event",
value={"reasoning": "chain-of-thought text"},
)
result = agent._dispatch_event(unknown_event)
# Should pass through to super() and return something (not None)
assert result is not None
def test_manually_emit_tool_call_with_empty_args(self, agent):
"""ManuallyEmitToolCall with empty args dict should still emit the tool call sequence."""
with track_parent_dispatches(agent) as dispatched:
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={
"id": "tc-empty",
"name": "MyTool",
"args": {},
},
)
agent._dispatch_event(event)
types = [getattr(e, "type", None) for e in dispatched]
assert EventType.TOOL_CALL_START in types
assert EventType.TOOL_CALL_ARGS in types
assert EventType.TOOL_CALL_END in types
# ---------- AG-UI-style (unprefixed) event integration ----------
class TestAGUIStyleEventIntegration:
"""Verify that AG-UI-native (unprefixed) events flow correctly through CopilotKit.
CopilotKit's _dispatch_event only intercepts copilotkit_-prefixed CUSTOM events;
unprefixed AG-UI events (manually_emit_message, manually_emit_tool_call, exit) are
delegated to the AG-UI base class — never suppressed or double-converted by CopilotKit.
"""
def _make_agent(self):
mock_graph = MagicMock()
mock_graph.get_state = MagicMock()
agent = LangGraphAGUIAgent(name="test", graph=mock_graph)
agent.active_run = {
"id": "run-1",
"thread_id": "t1",
"reasoning_process": None,
"node_name": "agent",
"has_function_streaming": False,
"model_made_tool_call": False,
"state_reliable": True,
"streamed_messages": [],
"manually_emitted_state": None,
"schema_keys": {
"input": ["messages", "tools"],
"output": ["messages", "tools"],
"config": [],
"context": [],
},
}
return agent
def test_agui_manually_emit_message_produces_text_events(self):
"""on_custom_event/"manually_emit_message" through CopilotKit should produce
TEXT_MESSAGE_START/CONTENT/END via the AG-UI base handler — not suppressed."""
import asyncio
agent = self._make_agent()
lg_event = {
"event": "on_custom_event",
"name": "manually_emit_message",
"data": {"message_id": "msg-1", "message": "Hello", "role": "assistant"},
"metadata": {},
}
async def _run():
async for _ in agent._handle_single_event(lg_event, {}):
pass
with track_parent_dispatches(agent) as dispatched:
asyncio.run(_run())
types = [e.type for e in dispatched]
assert EventType.TEXT_MESSAGE_START in types
assert EventType.TEXT_MESSAGE_CONTENT in types
assert EventType.TEXT_MESSAGE_END in types
def test_agui_manually_emit_tool_call_produces_tool_events(self):
"""on_custom_event/"manually_emit_tool_call" through CopilotKit should produce
TOOL_CALL_START/ARGS/END via the AG-UI base handler."""
import asyncio
agent = self._make_agent()
lg_event = {
"event": "on_custom_event",
"name": "manually_emit_tool_call",
"data": {"id": "tc-1", "name": "search", "args": {"q": "weather"}},
"metadata": {},
}
async def _run():
async for _ in agent._handle_single_event(lg_event, {}):
pass
with track_parent_dispatches(agent) as dispatched:
asyncio.run(_run())
types = [e.type for e in dispatched]
assert EventType.TOOL_CALL_START in types
assert EventType.TOOL_CALL_ARGS in types
assert EventType.TOOL_CALL_END in types
def test_agui_exit_produces_custom_event(self):
"""on_custom_event/"exit" through CopilotKit should emit a CUSTOM event —
the exit signal is not suppressed by the CopilotKit layer."""
import asyncio
agent = self._make_agent()
lg_event = {
"event": "on_custom_event",
"name": "exit",
"data": {},
"metadata": {},
}
async def _run():
async for _ in agent._handle_single_event(lg_event, {}):
pass
with track_parent_dispatches(agent) as dispatched:
asyncio.run(_run())
types = [e.type for e in dispatched]
assert EventType.CUSTOM in types
def test_agui_style_event_not_suppressed_by_dispatch(self):
"""A CUSTOM event with an AG-UI-style unprefixed name passed to CopilotKit's
_dispatch_event should be forwarded to the base class unchanged.
CopilotKit only converts copilotkit_-prefixed events (e.g. copilotkit_manually_emit_message
→ TEXT_MESSAGE_*). An unprefixed "manually_emit_message" CustomEvent must not trigger
that conversion path — no TEXT_MESSAGE_START should be injected, and the original
event object must be forwarded as-is."""
agent = self._make_agent()
original = CustomEvent(
type=EventType.CUSTOM,
name="manually_emit_message",
value={"message_id": "msg-2", "message": "test", "role": "assistant"},
)
with track_parent_dispatches(agent) as dispatched:
agent._dispatch_event(original)
types = [e.type for e in dispatched]
assert EventType.CUSTOM in types
assert EventType.TEXT_MESSAGE_START not in types
assert dispatched[0] is original
class TestAgentMetadata:
"""Regression coverage for info() serializing LangGraphAGUIAgent metadata."""
def test_dict_repr_includes_type_without_parent_support(self, agent):
"""dict_repr should not depend on AG-UI's base class implementing dict_repr."""
assert agent.dict_repr() == {
"name": "test",
"description": "",
"type": "langgraph_agui",
}
def test_remote_endpoint_info_serializes_langgraph_agui_agent(self):
"""sdk.info should include LangGraphAGUIAgent metadata without raising."""
mock_graph = MagicMock()
mock_graph.get_state = MagicMock()
agent = LangGraphAGUIAgent(
name="demo",
graph=mock_graph,
description="demo agent",
)
sdk = CopilotKitRemoteEndpoint(agents=[agent], actions=[])
info = sdk.info(context={"properties": {}, "frontend_url": None, "headers": {}})
assert info["agents"] == [
{
"name": "demo",
"description": "demo agent",
"type": "langgraph_agui",
}
]
@@ -0,0 +1,71 @@
"""Tests for #3690: LangGraphAGUIAgent stores Context as Pydantic objects instead of dicts."""
import json
from unittest.mock import MagicMock, patch
from ag_ui.core import Context
class TestAGUIContextSerialization:
"""Verify that context items stored in state are JSON-serializable dicts, not Pydantic objects."""
def test_context_items_are_dicts_not_pydantic(self):
"""Context from ag-ui properties should be model_dump'd before storage."""
from copilotkit.langgraph_agui_agent import LangGraphAGUIAgent
mock_graph = MagicMock()
mock_graph.get_state = MagicMock()
agent = LangGraphAGUIAgent(name="test", graph=mock_graph)
# Pydantic Context objects as they arrive from AG-UI
ctx1 = Context(name="user_info", description="User details", value="John")
ctx2 = Context(name="session", description="Session data", value="abc123")
merged_state = {
"ag-ui": {
"tools": [],
"context": [ctx1, ctx2],
},
"messages": [],
}
with patch.object(
type(agent).__mro__[1], # LangGraphAgent (parent class)
"langgraph_default_merge_state",
return_value=merged_state,
):
result = agent.langgraph_default_merge_state({}, [], None)
# Context items must be plain dicts, not Pydantic objects
for item in result["copilotkit"]["context"]:
assert isinstance(item, dict), f"Expected dict, got {type(item)}"
json.dumps(item) # Must be JSON-serializable
def test_context_with_mixed_types(self):
"""If context contains both Pydantic objects and plain dicts, handle both."""
from copilotkit.langgraph_agui_agent import LangGraphAGUIAgent
mock_graph = MagicMock()
mock_graph.get_state = MagicMock()
agent = LangGraphAGUIAgent(name="test", graph=mock_graph)
ctx_pydantic = Context(name="key1", description="desc", value="val1")
ctx_dict = {"name": "key2", "description": "desc2", "value": "val2"}
merged_state = {
"ag-ui": {
"tools": [],
"context": [ctx_pydantic, ctx_dict],
},
"messages": [],
}
with patch.object(
type(agent).__mro__[1],
"langgraph_default_merge_state",
return_value=merged_state,
):
result = agent.langgraph_default_merge_state({}, [], None)
for item in result["copilotkit"]["context"]:
assert isinstance(item, dict), f"Expected dict, got {type(item)}"
json.dumps(item)
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,246 @@
"""Tests for crewai_flow_messages_to_copilotkit assistant message emission.
Covers the parentMessageId orphan bug where the elif chain in the message
conversion skipped emitting the assistant message for tool-call messages.
Tool call entries reference their parent assistant message via parentMessageId,
so the assistant message must always be emitted — even when content is empty.
"""
import importlib
import importlib.util
import json
import sys
from unittest.mock import MagicMock
# crewai_sdk.py imports litellm/crewai at module level. Stub them out
# so the function under test (which needs none of these) can be imported.
# We load crewai_sdk.py directly to bypass copilotkit/crewai/__init__.py
# which pulls in crewai_agent.py and its heavy transitive dependencies.
_STUBS = [
"litellm",
"litellm.types",
"litellm.types.utils",
"litellm.litellm_core_utils",
"litellm.litellm_core_utils.streaming_handler",
"crewai",
"crewai.flow",
"crewai.flow.flow",
"crewai.utilities",
"crewai.utilities.events",
"crewai.utilities.events.flow_events",
"copilotkit.runloop",
"copilotkit.protocol",
]
_originals = {}
for _name in _STUBS:
if _name in sys.modules:
_originals[_name] = sys.modules[_name]
else:
sys.modules[_name] = MagicMock()
_pkg_path = importlib.util.find_spec("copilotkit").submodule_search_locations[0] # type: ignore[union-attr,index]
_spec = importlib.util.spec_from_file_location(
"copilotkit.crewai.crewai_sdk",
f"{_pkg_path}/crewai/crewai_sdk.py",
)
_mod = importlib.util.module_from_spec(_spec) # type: ignore[arg-type]
_spec.loader.exec_module(_mod) # type: ignore[union-attr]
crewai_flow_messages_to_copilotkit = _mod.crewai_flow_messages_to_copilotkit
# Restore original modules
for _name in _STUBS:
if _name in _originals:
sys.modules[_name] = _originals[_name]
else:
sys.modules.pop(_name, None)
def _convert_and_split(messages):
"""Convert messages and split result into assistant vs tool-call entries."""
result = crewai_flow_messages_to_copilotkit(messages)
assistant_msgs = [m for m in result if m.get("role") == "assistant"]
tool_call_msgs = [m for m in result if "parentMessageId" in m]
return result, assistant_msgs, tool_call_msgs
class TestCrewAIAssistantMessageAlwaysEmitted:
"""The assistant message must always be present so tool call entries can
reference it via parentMessageId. Without it, tool calls are orphaned
and the frontend cannot reconstruct tool call rendering on reconnect."""
def test_function_style_tool_calls_with_content(self):
"""Message with content and function-style tool_calls emits assistant + tool calls."""
messages = [
{
"id": "ai-1",
"role": "assistant",
"content": "Let me help.",
"tool_calls": [
{
"id": "tc-1",
"function": {
"name": "get_help",
"arguments": json.dumps({"topic": "billing"}),
},
},
],
},
]
_, assistant_msgs, tool_call_msgs = _convert_and_split(messages)
assert len(assistant_msgs) == 1
assert assistant_msgs[0]["id"] == "ai-1"
assert assistant_msgs[0]["content"] == "Let me help."
assert len(tool_call_msgs) == 1
assert tool_call_msgs[0]["parentMessageId"] == "ai-1"
def test_function_style_tool_calls_with_empty_content(self):
"""Message with empty content (OpenAI-style) still emits the assistant message."""
messages = [
{
"id": "ai-1",
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": "tc-1",
"function": {
"name": "get_help",
"arguments": json.dumps({"topic": "billing"}),
},
},
],
},
]
_, assistant_msgs, tool_call_msgs = _convert_and_split(messages)
assert len(assistant_msgs) == 1, (
"Assistant message must be emitted even with empty content"
)
assert assistant_msgs[0]["id"] == "ai-1"
assert assistant_msgs[0]["content"] == ""
assert len(tool_call_msgs) == 1
assert tool_call_msgs[0]["parentMessageId"] == "ai-1"
def test_function_style_tool_calls_without_content_key(self):
"""Message with no content key still emits the assistant message."""
messages = [
{
"id": "ai-1",
"role": "assistant",
"tool_calls": [
{
"id": "tc-1",
"function": {"name": "get_help", "arguments": json.dumps({})},
},
],
},
]
_, assistant_msgs, tool_call_msgs = _convert_and_split(messages)
assert len(assistant_msgs) == 1, (
"Assistant message must be emitted even without content key"
)
assert assistant_msgs[0]["content"] == ""
assert len(tool_call_msgs) == 1
assert tool_call_msgs[0]["parentMessageId"] == "ai-1"
def test_direct_style_tool_calls_with_empty_content(self):
"""Message with direct-style tool_calls (no function wrapper) still emits assistant."""
messages = [
{
"id": "ai-1",
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": "tc-1",
"name": "get_help",
"arguments": {"topic": "billing"},
},
],
},
]
_, assistant_msgs, tool_call_msgs = _convert_and_split(messages)
assert len(assistant_msgs) == 1, (
"Assistant message must be emitted for direct-style tool calls"
)
assert assistant_msgs[0]["content"] == ""
assert len(tool_call_msgs) == 1
assert tool_call_msgs[0]["parentMessageId"] == "ai-1"
def test_tool_call_without_id_is_skipped(self):
"""Tool calls missing an id should be silently skipped."""
messages = [
{
"id": "ai-1",
"role": "assistant",
"content": "",
"tool_calls": [
{"function": {"name": "no_id_tool", "arguments": json.dumps({})}},
{
"id": "tc-1",
"function": {
"name": "search",
"arguments": json.dumps({"q": "x"}),
},
},
],
},
]
_, assistant_msgs, tool_call_msgs = _convert_and_split(messages)
assert len(assistant_msgs) == 1
assert len(tool_call_msgs) == 1, "Only tool calls with an id should be emitted"
assert tool_call_msgs[0]["id"] == "tc-1"
def test_no_orphaned_parent_message_ids(self):
"""Every parentMessageId must reference an existing assistant message."""
messages = [
{"id": "h-1", "role": "user", "content": "help me"},
{
"id": "ai-1",
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": "tc-1",
"function": {
"name": "get_help",
"arguments": json.dumps({"topic": "billing"}),
},
},
{
"id": "tc-2",
"function": {
"name": "search",
"arguments": json.dumps({"query": "docs"}),
},
},
],
},
{"id": "tm-1", "role": "tool", "tool_call_id": "tc-1", "content": "done"},
{"id": "tm-2", "role": "tool", "tool_call_id": "tc-2", "content": "found"},
]
result, _, tool_call_msgs = _convert_and_split(messages)
message_ids = {m["id"] for m in result if "role" in m}
for tc in tool_call_msgs:
assert tc["parentMessageId"] in message_ids, (
f"Tool call {tc['id']} has orphaned parentMessageId {tc['parentMessageId']}"
)
def test_plain_assistant_message_without_tool_calls(self):
"""Plain assistant message (no tool calls) emits just the assistant message."""
messages = [{"id": "ai-1", "role": "assistant", "content": "Hello!"}]
result, _, _ = _convert_and_split(messages)
assert len(result) == 1
assert result[0]["role"] == "assistant"
assert result[0]["content"] == "Hello!"
@@ -0,0 +1,139 @@
"""Tests for CrewAI import compatibility (#3268).
When crewai>=0.177.0 is installed, the import path for flow events changed
from crewai.utilities.events.flow_events to crewai.events.types.flow_events.
The fix uses try/except to support both import paths.
"""
import sys
import types
import pytest
from unittest.mock import patch
def _make_fake_module(attrs: dict) -> types.ModuleType:
"""Create a fake module with given attributes."""
mod = types.ModuleType("fake")
for k, v in attrs.items():
setattr(mod, k, v)
return mod
class TestCrewAIImportCompat:
"""Test that crewai_sdk.py can import flow events from either path."""
def test_old_import_path_works(self):
"""When old crewai (<0.177.0) is installed, old import path should work."""
# The old path: crewai.utilities.events.flow_events
# If it exists, the import should succeed
fake_events = _make_fake_module(
{
"FlowEvent": type("FlowEvent", (), {}),
"FlowStartedEvent": type("FlowStartedEvent", (), {}),
"MethodExecutionStartedEvent": type(
"MethodExecutionStartedEvent", (), {}
),
"MethodExecutionFinishedEvent": type(
"MethodExecutionFinishedEvent", (), {}
),
"FlowFinishedEvent": type("FlowFinishedEvent", (), {}),
}
)
with patch.dict(
sys.modules,
{
"crewai.utilities.events.flow_events": fake_events,
},
):
# Simulate what our try/except does
try:
from crewai.utilities.events.flow_events import FlowEvent
old_path_works = True
except ImportError:
old_path_works = False
assert old_path_works, "Old import path should work with old crewai"
def test_new_import_path_fallback(self):
"""When new crewai (>=0.177.0) is installed, new import path should work."""
# The new path: crewai.events.types.flow_events
fake_events = _make_fake_module(
{
"FlowEvent": type("FlowEvent", (), {}),
"FlowStartedEvent": type("FlowStartedEvent", (), {}),
"MethodExecutionStartedEvent": type(
"MethodExecutionStartedEvent", (), {}
),
"MethodExecutionFinishedEvent": type(
"MethodExecutionFinishedEvent", (), {}
),
"FlowFinishedEvent": type("FlowFinishedEvent", (), {}),
}
)
# Remove old path, add new path
old_mod_key = "crewai.utilities.events.flow_events"
saved = sys.modules.get(old_mod_key)
try:
# Make old path fail
sys.modules[old_mod_key] = None # type: ignore
with patch.dict(
sys.modules,
{
"crewai.events.types.flow_events": fake_events,
},
clear=False,
):
# Simulate fallback logic
try:
from crewai.utilities.events.flow_events import FlowEvent # type: ignore
new_path_needed = False
except ImportError:
new_path_needed = True
assert new_path_needed, "Old path should fail, triggering fallback"
from crewai.events.types.flow_events import FlowEvent # type: ignore
assert FlowEvent is not None, "New path should work"
finally:
if saved is not None:
sys.modules[old_mod_key] = saved
elif old_mod_key in sys.modules:
del sys.modules[old_mod_key]
def test_both_paths_fail_raises_import_error(self):
"""When neither path works, ImportError should propagate."""
# This simulates crewai not having flow events at all
old_key = "crewai.utilities.events.flow_events"
new_key = "crewai.events.types.flow_events"
saved_old = sys.modules.get(old_key)
saved_new = sys.modules.get(new_key)
try:
sys.modules[old_key] = None # type: ignore
sys.modules[new_key] = None # type: ignore
with pytest.raises(ImportError):
# Old path
try:
from crewai.utilities.events.flow_events import FlowEvent # type: ignore
except ImportError:
pass
# New path
from crewai.events.types.flow_events import FlowEvent # type: ignore
finally:
if saved_old is not None:
sys.modules[old_key] = saved_old
elif old_key in sys.modules:
del sys.modules[old_key]
if saved_new is not None:
sys.modules[new_key] = saved_new
elif new_key in sys.modules:
del sys.modules[new_key]
+194
View File
@@ -0,0 +1,194 @@
"""Tests for emit_messages/emit_tool_calls filtering in LangGraphAGUIAgent.
Covers the two bugs fixed in https://github.com/CopilotKit/CopilotKit/issues/2066:
1. raw_event is a dict, so metadata must be read with .get() not getattr()
2. Filtered events must return None (not ""), and run() must strip them
"""
import asyncio
import pytest
from unittest.mock import patch, AsyncMock, MagicMock
from ag_ui.core import (
EventType,
TextMessageStartEvent,
TextMessageContentEvent,
ToolCallStartEvent,
)
from copilotkit.langgraph_agui_agent import LangGraphAGUIAgent
@pytest.fixture
def agent():
"""Create a LangGraphAGUIAgent with a mocked graph."""
mock_graph = MagicMock()
mock_graph.get_state = MagicMock()
a = LangGraphAGUIAgent(name="test", graph=mock_graph)
a.active_run = {"id": "run-1", "thread_id": "t-1"}
return a
def _make_text_event(metadata: dict) -> TextMessageContentEvent:
"""Create a TEXT_MESSAGE_CONTENT event with a dict raw_event carrying metadata."""
return TextMessageContentEvent(
messageId="msg-1",
delta="hello",
rawEvent={"metadata": metadata},
)
def _make_tool_event(metadata: dict) -> ToolCallStartEvent:
"""Create a TOOL_CALL_START event with a dict raw_event carrying metadata."""
return ToolCallStartEvent(
toolCallId="tc-1",
toolCallName="some_tool",
rawEvent={"metadata": metadata},
)
# ---------- Bug 1: dict metadata reading via .get() ----------
class TestDictMetadataReading:
"""raw_event is a dict — metadata must be read with .get(), not getattr()."""
def test_emit_messages_false_filters_text_event(self, agent):
"""emit-messages=False should suppress text message events."""
event = _make_text_event({"copilotkit:emit-messages": False})
result = agent._dispatch_event(event)
assert result is None
def test_emit_tool_calls_false_filters_tool_event(self, agent):
"""emit-tool-calls=False should suppress tool call events."""
event = _make_tool_event({"copilotkit:emit-tool-calls": False})
result = agent._dispatch_event(event)
assert result is None
def test_emit_messages_true_passes_through(self, agent):
"""emit-messages=True should NOT filter — event passes to super()."""
event = _make_text_event({"copilotkit:emit-messages": True})
result = agent._dispatch_event(event)
# super()._dispatch_event returns the event object (not None or "")
assert result is not None
def test_no_metadata_key_passes_through(self, agent):
"""No emit- keys in metadata — event should pass through."""
event = _make_text_event({"some-other-key": True})
result = agent._dispatch_event(event)
assert result is not None
def test_empty_metadata_passes_through(self, agent):
"""Empty metadata dict — event should pass through."""
event = _make_text_event({})
result = agent._dispatch_event(event)
assert result is not None
# ---------- Bug 2: filtered returns None, not "" ----------
class TestFilteredReturnValue:
"""Filtered events must return None (not empty string) to avoid encoder crash."""
def test_filtered_message_returns_none_not_empty_string(self, agent):
event = _make_text_event({"copilotkit:emit-messages": False})
result = agent._dispatch_event(event)
assert result is None, f"Expected None, got {result!r}"
assert result != "", "Must not return empty string — crashes the encoder"
def test_filtered_tool_call_returns_none_not_empty_string(self, agent):
event = _make_tool_event({"copilotkit:emit-tool-calls": False})
result = agent._dispatch_event(event)
assert result is None, f"Expected None, got {result!r}"
assert result != "", "Must not return empty string — crashes the encoder"
# ---------- run() filters out None values ----------
class TestRunFiltersNone:
"""The run() override must strip None values from the event stream."""
def test_run_filters_none_events(self, agent):
"""None values from _dispatch_event should not appear in run() output."""
# Mock super().run() to yield a mix of real events and None
real_event = "data: {}\n\n"
async def mock_super_run(self, input):
yield real_event
yield None
yield real_event
with patch.object(
LangGraphAGUIAgent.__bases__[0],
"run",
new=mock_super_run,
):
results = asyncio.run(_collect_async_gen(agent.run(MagicMock())))
assert results == [real_event, real_event]
assert None not in results
async def _collect_async_gen(agen):
"""Collect all items from an async generator into a list."""
items = []
async for item in agen:
items.append(item)
return items
# ---------- Edge cases: missing / None rawEvent ----------
class TestMissingOrNoneRawEvent:
"""Events where rawEvent is absent or None should pass through to super() without crashing."""
def test_none_raw_event_passes_through(self, agent):
"""Event with raw_event=None should not crash and should pass through."""
event = TextMessageContentEvent(
messageId="msg-1",
delta="hello",
rawEvent=None,
)
# Should not raise and should call super() (returning non-None)
result = agent._dispatch_event(event)
assert result is not None
def test_event_without_raw_event_attr_passes_through(self, agent):
"""If raw_event attribute is missing entirely, should pass through."""
# TextMessageStartEvent has rawEvent as optional, passing None simulates missing
event = TextMessageStartEvent(
messageId="msg-1",
role="assistant",
rawEvent=None,
)
result = agent._dispatch_event(event)
assert result is not None
# ---------- Edge cases: None values for emit metadata keys ----------
class TestNoneEmitMetadataValues:
"""None values for emit metadata keys should NOT filter events (only False filters)."""
def test_none_emit_messages_does_not_filter(self, agent):
"""copilotkit:emit-messages=None should not suppress text events (only False does)."""
event = _make_text_event({"copilotkit:emit-messages": None})
result = agent._dispatch_event(event)
assert result is not None, (
"None value should not filter — only False suppresses"
)
def test_none_emit_tool_calls_does_not_filter(self, agent):
"""copilotkit:emit-tool-calls=None should not suppress tool events."""
event = _make_tool_event({"copilotkit:emit-tool-calls": None})
result = agent._dispatch_event(event)
assert result is not None
def test_zero_emit_messages_does_not_filter(self, agent):
"""0 (falsy but not False) should not suppress text events (only False does)."""
event = _make_text_event({"copilotkit:emit-messages": 0})
result = agent._dispatch_event(event)
assert result is not None
@@ -0,0 +1,807 @@
"""Tests for the optional `id` parameter on copilotkit_emit_tool_call.
Covers:
1. LangGraph variant: default UUID generation, custom ID passthrough, return value
2. CrewAI variant: default UUID generation, custom ID passthrough, return value
3. AG-UI agent dispatch: custom ID propagates to all three TOOL_CALL events
4. AG-UI dispatch validation: defensive CopilotKitMisuseError paths for
missing/invalid id, name, args, non-serializable args, and non-dict value
"""
import asyncio
import json
import logging
import uuid
import pytest
from unittest.mock import MagicMock, AsyncMock, patch
from ag_ui.core import (
EventType,
CustomEvent,
)
from ag_ui_langgraph import LangGraphAgent as AGUIBase
from copilotkit.langgraph_agui_agent import (
LangGraphAGUIAgent,
CustomEventNames,
)
from copilotkit.exc import CopilotKitMisuseError
# ---- Fixtures ----
@pytest.fixture
def agent():
"""Create a LangGraphAGUIAgent with a mocked graph."""
mock_graph = MagicMock()
mock_graph.get_state = MagicMock()
a = LangGraphAGUIAgent(name="test", graph=mock_graph)
a.active_run = {"id": "run-1", "thread_id": "t-1"}
return a
def _track_parent_dispatches(agent):
"""Collect events dispatched to the AG-UI base class."""
from contextlib import contextmanager
@contextmanager
def _ctx():
dispatched = []
original = AGUIBase._dispatch_event
def _tracking(self_inner, event):
dispatched.append(event)
return original(self_inner, event)
with patch.object(AGUIBase, "_dispatch_event", new=_tracking):
yield dispatched
return _ctx()
# ---- LangGraph variant tests ----
class TestLangGraphEmitToolCallOptionalId:
"""copilotkit_emit_tool_call (langgraph) with optional id parameter."""
@pytest.mark.asyncio
async def test_default_generates_uuid(self):
"""When no id is provided, a UUID v4 string should be generated and returned."""
with patch(
"copilotkit.langgraph.adispatch_custom_event", new_callable=AsyncMock
) as mock_dispatch:
from copilotkit.langgraph import copilotkit_emit_tool_call
config = {"metadata": {}}
result = await copilotkit_emit_tool_call(
config, name="MyTool", args={"key": "val"}
)
assert isinstance(result, str)
uuid.UUID(result)
payload = mock_dispatch.call_args[0][1]
assert payload["id"] == result
assert payload["name"] == "MyTool"
assert payload["args"] == {"key": "val"}
@pytest.mark.asyncio
async def test_custom_id_passthrough(self):
"""When a custom id is provided, it should be used as-is."""
with patch(
"copilotkit.langgraph.adispatch_custom_event", new_callable=AsyncMock
) as mock_dispatch:
from copilotkit.langgraph import copilotkit_emit_tool_call
config = {"metadata": {}}
result = await copilotkit_emit_tool_call(
config, name="MyTool", args={"key": "val"}, tool_call_id="custom-id-123"
)
assert result == "custom-id-123"
payload = mock_dispatch.call_args[0][1]
assert payload["id"] == "custom-id-123"
@pytest.mark.asyncio
async def test_returns_generated_id(self):
"""The return value should be the tool call ID (generated or custom)."""
with patch(
"copilotkit.langgraph.adispatch_custom_event", new_callable=AsyncMock
):
from copilotkit.langgraph import copilotkit_emit_tool_call
config = {"metadata": {}}
result_auto = await copilotkit_emit_tool_call(config, name="Tool", args={})
assert isinstance(result_auto, str)
assert len(result_auto) > 0
result_custom = await copilotkit_emit_tool_call(
config, name="Tool", args={}, tool_call_id="my-id"
)
assert result_custom == "my-id"
@pytest.mark.asyncio
async def test_none_id_generates_uuid(self):
"""Explicitly passing tool_call_id=None should behave the same as omitting it."""
with patch(
"copilotkit.langgraph.adispatch_custom_event", new_callable=AsyncMock
) as mock_dispatch:
from copilotkit.langgraph import copilotkit_emit_tool_call
config = {"metadata": {}}
result = await copilotkit_emit_tool_call(
config, name="Tool", args={}, tool_call_id=None
)
assert isinstance(result, str)
uuid.UUID(result)
assert mock_dispatch.call_args[0][1]["id"] == result
@pytest.mark.asyncio
async def test_empty_string_id_raises(self):
"""Passing an empty string should raise ValueError."""
with patch(
"copilotkit.langgraph.adispatch_custom_event", new_callable=AsyncMock
):
from copilotkit.langgraph import copilotkit_emit_tool_call
config = {"metadata": {}}
with pytest.raises(ValueError, match="non-empty string"):
await copilotkit_emit_tool_call(
config, name="Tool", args={}, tool_call_id=""
)
@pytest.mark.asyncio
async def test_whitespace_only_id_raises(self):
"""Passing a whitespace-only string should raise ValueError."""
with patch(
"copilotkit.langgraph.adispatch_custom_event", new_callable=AsyncMock
):
from copilotkit.langgraph import copilotkit_emit_tool_call
config = {"metadata": {}}
with pytest.raises(ValueError, match="non-empty string"):
await copilotkit_emit_tool_call(
config, name="Tool", args={}, tool_call_id=" "
)
@pytest.mark.asyncio
async def test_whitespace_only_name_raises(self):
"""Passing a whitespace-only name should raise CopilotKitMisuseError."""
with patch(
"copilotkit.langgraph.adispatch_custom_event", new_callable=AsyncMock
):
from copilotkit.langgraph import copilotkit_emit_tool_call
config = {"metadata": {}}
with pytest.raises(CopilotKitMisuseError, match="non-empty string"):
await copilotkit_emit_tool_call(config, name=" ", args={})
@pytest.mark.asyncio
async def test_empty_name_raises(self):
"""Passing an empty name should raise CopilotKitMisuseError."""
with patch(
"copilotkit.langgraph.adispatch_custom_event", new_callable=AsyncMock
):
from copilotkit.langgraph import copilotkit_emit_tool_call
config = {"metadata": {}}
with pytest.raises(CopilotKitMisuseError, match="non-empty string"):
await copilotkit_emit_tool_call(config, name="", args={})
@pytest.mark.asyncio
async def test_non_serializable_args_raises(self):
"""Passing non-JSON-serializable args should raise CopilotKitMisuseError."""
with patch(
"copilotkit.langgraph.adispatch_custom_event", new_callable=AsyncMock
):
from copilotkit.langgraph import copilotkit_emit_tool_call
config = {"metadata": {}}
with pytest.raises(CopilotKitMisuseError, match="not JSON-serializable"):
await copilotkit_emit_tool_call(
config, name="Tool", args={"bad": {1, 2, 3}}
)
@pytest.mark.asyncio
async def test_cancelled_error_propagates_from_post_dispatch_sleep(self):
"""CancelledError during the shielded post-dispatch sleep must propagate."""
with patch(
"copilotkit.langgraph.adispatch_custom_event", new_callable=AsyncMock
) as mock_dispatch:
from copilotkit.langgraph import copilotkit_emit_tool_call
config = {"metadata": {}}
async def _run_and_cancel():
task = asyncio.current_task()
# Schedule cancellation after dispatch completes but during sleep
original_sleep = asyncio.sleep
async def _cancel_during_sleep(delay):
task.cancel()
await original_sleep(0)
with patch(
"copilotkit.langgraph.asyncio.sleep",
side_effect=_cancel_during_sleep,
):
with patch(
"copilotkit.langgraph.asyncio.shield",
side_effect=lambda coro: coro,
):
return await copilotkit_emit_tool_call(
config,
name="CancelTool",
args={},
tool_call_id="cancel-test-id",
)
with pytest.raises(asyncio.CancelledError):
await _run_and_cancel()
mock_dispatch.assert_called_once()
@pytest.mark.asyncio
async def test_cancelled_error_logs_warning(self, caplog):
"""CancelledError during post-dispatch sleep should log with the tool_call_id."""
with patch(
"copilotkit.langgraph.adispatch_custom_event", new_callable=AsyncMock
):
from copilotkit.langgraph import copilotkit_emit_tool_call
config = {"metadata": {}}
async def _cancel_sleep(delay):
raise asyncio.CancelledError()
with caplog.at_level(logging.WARNING, logger="copilotkit.langgraph"):
with patch(
"copilotkit.langgraph.asyncio.sleep", side_effect=_cancel_sleep
):
with patch(
"copilotkit.langgraph.asyncio.shield",
side_effect=lambda coro: coro,
):
with pytest.raises(asyncio.CancelledError):
await copilotkit_emit_tool_call(
config,
name="Tool",
args={},
tool_call_id="log-cancel-id",
)
assert any("log-cancel-id" in record.message for record in caplog.records)
# ---- CrewAI variant tests ----
try:
import crewai # noqa: F401
_has_crewai = True
except ImportError:
_has_crewai = False
@pytest.mark.skipif(not _has_crewai, reason="crewai not installed")
class TestCrewAIEmitToolCallOptionalId:
"""copilotkit_emit_tool_call (crewai) with optional id parameter."""
@pytest.mark.asyncio
async def test_default_generates_uuid(self):
"""When no id is provided, a UUID v4 string should be generated and returned."""
with patch(
"copilotkit.crewai.crewai_sdk.queue_put", new_callable=AsyncMock
) as mock_queue:
from copilotkit.crewai.crewai_sdk import copilotkit_emit_tool_call
result = await copilotkit_emit_tool_call(name="MyTool", args={"key": "val"})
assert isinstance(result, str)
uuid.UUID(result)
start_ev, args_ev, end_ev = mock_queue.call_args[0]
assert start_ev["actionExecutionId"] == result
assert start_ev["parentMessageId"] == result
assert start_ev["actionName"] == "MyTool"
assert args_ev["actionExecutionId"] == result
assert json.loads(args_ev["args"]) == {"key": "val"}
assert end_ev["actionExecutionId"] == result
@pytest.mark.asyncio
async def test_custom_id_passthrough(self):
"""When a custom id is provided, it should be used as the message_id."""
with patch(
"copilotkit.crewai.crewai_sdk.queue_put", new_callable=AsyncMock
) as mock_queue:
from copilotkit.crewai.crewai_sdk import copilotkit_emit_tool_call
result = await copilotkit_emit_tool_call(
name="MyTool", args={"key": "val"}, tool_call_id="crew-custom-id"
)
assert result == "crew-custom-id"
start_ev, args_ev, end_ev = mock_queue.call_args[0]
assert start_ev["actionExecutionId"] == "crew-custom-id"
assert start_ev["parentMessageId"] == "crew-custom-id"
assert args_ev["actionExecutionId"] == "crew-custom-id"
assert end_ev["actionExecutionId"] == "crew-custom-id"
@pytest.mark.asyncio
async def test_returns_id(self):
"""Should return the tool call ID regardless of whether it was auto or custom."""
with patch("copilotkit.crewai.crewai_sdk.queue_put", new_callable=AsyncMock):
from copilotkit.crewai.crewai_sdk import copilotkit_emit_tool_call
result_auto = await copilotkit_emit_tool_call(name="T", args={})
assert isinstance(result_auto, str)
assert len(result_auto) > 0
result_custom = await copilotkit_emit_tool_call(
name="T", args={}, tool_call_id="explicit"
)
assert result_custom == "explicit"
@pytest.mark.asyncio
async def test_empty_string_id_raises(self):
"""Passing an empty string should raise ValueError."""
with patch("copilotkit.crewai.crewai_sdk.queue_put", new_callable=AsyncMock):
from copilotkit.crewai.crewai_sdk import copilotkit_emit_tool_call
with pytest.raises(ValueError, match="non-empty string"):
await copilotkit_emit_tool_call(name="Tool", args={}, tool_call_id="")
@pytest.mark.asyncio
async def test_whitespace_only_id_raises(self):
"""Passing a whitespace-only string should raise ValueError."""
with patch("copilotkit.crewai.crewai_sdk.queue_put", new_callable=AsyncMock):
from copilotkit.crewai.crewai_sdk import copilotkit_emit_tool_call
with pytest.raises(ValueError, match="non-empty string"):
await copilotkit_emit_tool_call(
name="Tool", args={}, tool_call_id=" "
)
@pytest.mark.asyncio
async def test_none_id_generates_uuid(self):
"""Explicitly passing tool_call_id=None should behave the same as omitting it."""
with patch("copilotkit.crewai.crewai_sdk.queue_put", new_callable=AsyncMock):
from copilotkit.crewai.crewai_sdk import copilotkit_emit_tool_call
result = await copilotkit_emit_tool_call(
name="Tool", args={}, tool_call_id=None
)
assert isinstance(result, str)
uuid.UUID(result)
@pytest.mark.asyncio
async def test_whitespace_only_name_raises(self):
"""Passing a whitespace-only name should raise CopilotKitMisuseError."""
with patch("copilotkit.crewai.crewai_sdk.queue_put", new_callable=AsyncMock):
from copilotkit.crewai.crewai_sdk import copilotkit_emit_tool_call
with pytest.raises(CopilotKitMisuseError, match="non-empty string"):
await copilotkit_emit_tool_call(name=" ", args={})
@pytest.mark.asyncio
async def test_non_serializable_args_raises(self):
"""Passing non-JSON-serializable args should raise CopilotKitMisuseError."""
with patch("copilotkit.crewai.crewai_sdk.queue_put", new_callable=AsyncMock):
from copilotkit.crewai.crewai_sdk import copilotkit_emit_tool_call
with pytest.raises(CopilotKitMisuseError, match="not JSON-serializable"):
await copilotkit_emit_tool_call(name="Tool", args={"bad": {1, 2, 3}})
# ---- CrewAI variant: compensating action_execution_end tests ----
@pytest.mark.skipif(not _has_crewai, reason="crewai not installed")
class TestCrewAICompensatingEnd:
"""Tests for the compensating action_execution_end when dispatch fails mid-stream.
queue_put is called once with all three events (start, args, end) in a single
batch for atomicity. If the batch fails, a compensating end is always attempted
as a best-effort measure — an orphaned END is harmless, but an orphaned START
hangs the client UI.
"""
@pytest.mark.asyncio
async def test_batch_failure_emits_compensating_end(self):
"""If the batched queue_put fails, a compensating end is emitted."""
call_count = 0
async def _failing_queue_put(*events):
nonlocal call_count
call_count += 1
if call_count == 1:
raise RuntimeError("batch failed")
with patch("copilotkit.crewai.crewai_sdk.queue_put", new=_failing_queue_put):
from copilotkit.crewai.crewai_sdk import copilotkit_emit_tool_call
with pytest.raises(RuntimeError, match="batch failed"):
await copilotkit_emit_tool_call(
name="FailTool", args={"x": 1}, tool_call_id="comp-crew-1"
)
assert call_count == 2
@pytest.mark.asyncio
async def test_compensating_end_failure_reraises_original(self):
"""If the compensating end also fails, the original error still propagates."""
call_count = 0
async def _failing_queue_put(*events):
nonlocal call_count
call_count += 1
raise RuntimeError(f"queue failure #{call_count}")
with patch("copilotkit.crewai.crewai_sdk.queue_put", new=_failing_queue_put):
from copilotkit.crewai.crewai_sdk import copilotkit_emit_tool_call
with pytest.raises(RuntimeError, match="queue failure #1"):
await copilotkit_emit_tool_call(
name="FailTool", args={}, tool_call_id="comp-crew-3"
)
assert call_count == 2
@pytest.mark.asyncio
async def test_compensating_end_failure_emits_log(self, caplog):
"""The logger.error call includes the message_id when compensating end fails."""
call_count = 0
async def _failing_queue_put(*events):
nonlocal call_count
call_count += 1
raise RuntimeError(f"queue failure #{call_count}")
with caplog.at_level(logging.ERROR, logger="copilotkit.crewai.crewai_sdk"):
with patch(
"copilotkit.crewai.crewai_sdk.queue_put", new=_failing_queue_put
):
from copilotkit.crewai.crewai_sdk import copilotkit_emit_tool_call
with pytest.raises(RuntimeError):
await copilotkit_emit_tool_call(
name="FailTool", args={}, tool_call_id="log-crew-id"
)
assert any("log-crew-id" in record.message for record in caplog.records)
# ---- AG-UI dispatch: custom ID propagates through all events ----
class TestCustomIdPropagatesThroughAGUI:
"""When a custom id is used, the downstream AG-UI events carry that exact ID."""
def test_custom_id_in_all_tool_call_events(self, agent):
"""TOOL_CALL_START, TOOL_CALL_ARGS, and TOOL_CALL_END should all carry the custom id."""
with _track_parent_dispatches(agent) as dispatched:
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={
"id": "user-provided-id-42",
"name": "CustomTool",
"args": {"x": 1},
},
)
agent._dispatch_event(event)
tool_events = [e for e in dispatched if hasattr(e, "tool_call_id")]
assert len(tool_events) == 3
for e in tool_events:
assert e.tool_call_id == "user-provided-id-42"
def test_custom_id_in_parent_message_id(self, agent):
"""ToolCallStartEvent.parent_message_id should match the custom id."""
with _track_parent_dispatches(agent) as dispatched:
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={
"id": "parent-test-id",
"name": "ParentTool",
"args": {},
},
)
agent._dispatch_event(event)
start_events = [e for e in dispatched if e.type == EventType.TOOL_CALL_START]
assert len(start_events) == 1
assert start_events[0].parent_message_id == "parent-test-id"
def test_custom_id_with_dict_args_serialized(self, agent):
"""Custom id + dict args should both work: args JSON-serialized, id preserved."""
with _track_parent_dispatches(agent) as dispatched:
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={
"id": "combo-test",
"name": "ComboTool",
"args": {"nested": {"deep": True}},
},
)
agent._dispatch_event(event)
args_events = [e for e in dispatched if e.type == EventType.TOOL_CALL_ARGS]
assert len(args_events) == 1
assert args_events[0].tool_call_id == "combo-test"
assert json.loads(args_events[0].delta) == {"nested": {"deep": True}}
def test_string_args_passed_through_unchanged(self, agent):
"""When args is already a JSON string, it should be passed through without re-serializing."""
with _track_parent_dispatches(agent) as dispatched:
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={
"id": "string-args-test",
"name": "StringArgsTool",
"args": '{"x": 1}',
},
)
agent._dispatch_event(event)
args_events = [e for e in dispatched if e.type == EventType.TOOL_CALL_ARGS]
assert len(args_events) == 1
assert args_events[0].delta == '{"x": 1}'
def test_empty_dict_args_does_not_raise(self, agent):
"""An empty dict for args is valid and should not raise."""
with _track_parent_dispatches(agent) as dispatched:
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={
"id": "empty-args-test",
"name": "EmptyArgsTool",
"args": {},
},
)
agent._dispatch_event(event)
tool_events = [e for e in dispatched if hasattr(e, "tool_call_id")]
assert len(tool_events) == 3
# ---- AG-UI dispatch: validation negative tests ----
class TestAGUIDispatchValidation:
"""Negative tests for defensive validation in _dispatch_event."""
def test_missing_id_raises(self, agent):
"""Event with no 'id' field should raise CopilotKitMisuseError."""
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={"name": "Tool", "args": {}},
)
with pytest.raises(CopilotKitMisuseError, match="valid 'id'"):
agent._dispatch_event(event)
def test_non_string_id_raises(self, agent):
"""Event with non-string 'id' should raise CopilotKitMisuseError."""
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={"id": 42, "name": "Tool", "args": {}},
)
with pytest.raises(CopilotKitMisuseError, match="valid 'id'"):
agent._dispatch_event(event)
def test_empty_string_id_raises(self, agent):
"""Event with empty string 'id' should raise CopilotKitMisuseError."""
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={"id": "", "name": "Tool", "args": {}},
)
with pytest.raises(CopilotKitMisuseError, match="valid 'id'"):
agent._dispatch_event(event)
def test_whitespace_only_id_raises(self, agent):
"""Event with whitespace-only 'id' should raise CopilotKitMisuseError."""
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={"id": " ", "name": "Tool", "args": {}},
)
with pytest.raises(CopilotKitMisuseError, match="valid 'id'"):
agent._dispatch_event(event)
def test_missing_name_raises(self, agent):
"""Event with no 'name' field should raise CopilotKitMisuseError."""
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={"id": "valid-id", "args": {}},
)
with pytest.raises(CopilotKitMisuseError, match="valid 'name'"):
agent._dispatch_event(event)
def test_whitespace_only_name_raises(self, agent):
"""Event with whitespace-only 'name' should raise CopilotKitMisuseError."""
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={"id": "valid-id", "name": " ", "args": {}},
)
with pytest.raises(CopilotKitMisuseError, match="valid 'name'"):
agent._dispatch_event(event)
def test_missing_args_raises(self, agent):
"""Event with no 'args' field should raise CopilotKitMisuseError."""
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={"id": "valid-id", "name": "Tool"},
)
with pytest.raises(CopilotKitMisuseError, match="missing 'args'"):
agent._dispatch_event(event)
def test_non_serializable_args_raises(self, agent):
"""Event with non-JSON-serializable args (set) should raise CopilotKitMisuseError."""
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={"id": "valid-id", "name": "Tool", "args": {1, 2, 3}},
)
with pytest.raises(CopilotKitMisuseError, match="not JSON-serializable"):
agent._dispatch_event(event)
def test_non_dict_value_raises(self, agent):
"""Event with non-dict value should raise CopilotKitMisuseError."""
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value=None,
)
with pytest.raises(CopilotKitMisuseError, match="must be a dict"):
agent._dispatch_event(event)
def test_list_args_accepted(self, agent):
"""Event with list args should be accepted (JSON-serializable)."""
with _track_parent_dispatches(agent) as dispatched:
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={"id": "list-args-id", "name": "Tool", "args": [1, 2, 3]},
)
agent._dispatch_event(event)
args_events = [e for e in dispatched if e.type == EventType.TOOL_CALL_ARGS]
assert len(args_events) == 1
assert args_events[0].delta == "[1, 2, 3]"
def test_int_args_accepted(self, agent):
"""Event with int args should be accepted (JSON-serializable)."""
with _track_parent_dispatches(agent) as dispatched:
event = CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={"id": "int-args-id", "name": "Tool", "args": 42},
)
agent._dispatch_event(event)
args_events = [e for e in dispatched if e.type == EventType.TOOL_CALL_ARGS]
assert len(args_events) == 1
assert args_events[0].delta == "42"
# ---- AG-UI dispatch: compensating TOOL_CALL_END on mid-stream failure ----
class TestAGUICompensatingEnd:
"""Tests for the compensating TOOL_CALL_END when dispatch fails mid-stream."""
def _make_event(self, tool_call_id="comp-test-id"):
return CustomEvent(
type=EventType.CUSTOM,
name=CustomEventNames.ManuallyEmitToolCall.value,
value={
"id": tool_call_id,
"name": "FailTool",
"args": {"x": 1},
},
)
def test_failure_after_start_emits_compensating_end(self, agent):
"""If TOOL_CALL_ARGS fails after START was sent, a compensating END is dispatched."""
call_count = 0
original = AGUIBase._dispatch_event
def _fail_on_args(self_inner, evt):
nonlocal call_count
call_count += 1
if call_count == 2:
raise RuntimeError("args dispatch failed")
return original(self_inner, evt)
with patch.object(AGUIBase, "_dispatch_event", new=_fail_on_args):
with pytest.raises(RuntimeError, match="args dispatch failed"):
agent._dispatch_event(self._make_event())
assert call_count == 3
def test_failure_on_start_does_not_emit_compensating_end(self, agent):
"""If TOOL_CALL_START itself fails, no compensating END is dispatched."""
call_count = 0
original = AGUIBase._dispatch_event
def _fail_on_start(self_inner, evt):
nonlocal call_count
call_count += 1
if call_count == 1:
raise RuntimeError("start dispatch failed")
return original(self_inner, evt)
with patch.object(AGUIBase, "_dispatch_event", new=_fail_on_start):
with pytest.raises(RuntimeError, match="start dispatch failed"):
agent._dispatch_event(self._make_event())
assert call_count == 1
def test_compensating_end_failure_reraises_original(self, agent):
"""If the compensating END also fails, the original error propagates."""
call_count = 0
original = AGUIBase._dispatch_event
def _fail_on_args_and_end(self_inner, evt):
nonlocal call_count
call_count += 1
if call_count == 1:
return original(self_inner, evt)
raise RuntimeError(f"dispatch failure #{call_count}")
with patch.object(AGUIBase, "_dispatch_event", new=_fail_on_args_and_end):
with pytest.raises(RuntimeError, match="dispatch failure #2"):
agent._dispatch_event(self._make_event())
assert call_count == 3
def test_compensating_end_failure_emits_log(self, agent, caplog):
"""The logger.error call includes the tool_call_id when compensating END fails."""
call_count = 0
original = AGUIBase._dispatch_event
def _fail_on_args_and_end(self_inner, evt):
nonlocal call_count
call_count += 1
if call_count == 1:
return original(self_inner, evt)
raise RuntimeError(f"dispatch failure #{call_count}")
with caplog.at_level(logging.ERROR, logger="copilotkit.langgraph_agui_agent"):
with patch.object(AGUIBase, "_dispatch_event", new=_fail_on_args_and_end):
with pytest.raises(RuntimeError):
agent._dispatch_event(self._make_event("log-test-id"))
assert any("log-test-id" in record.message for record in caplog.records)
def test_failure_on_end_emits_compensating_end(self, agent):
"""If TOOL_CALL_END itself fails, a compensating END is still attempted."""
call_count = 0
original = AGUIBase._dispatch_event
def _fail_on_end(self_inner, evt):
nonlocal call_count
call_count += 1
if call_count == 3:
raise RuntimeError("end dispatch failed")
return original(self_inner, evt)
with patch.object(AGUIBase, "_dispatch_event", new=_fail_on_end):
with pytest.raises(RuntimeError, match="end dispatch failed"):
agent._dispatch_event(self._make_event("end-fail-id"))
assert call_count == 4
+464
View File
@@ -0,0 +1,464 @@
"""Tests for header propagation (x-* prefixed headers to outgoing LLM calls)."""
import asyncio
import contextvars
import inspect
import warnings
import httpx
import pytest
from copilotkit.header_propagation import (
get_forwarded_headers,
install_httpx_hook,
set_forwarded_headers,
)
class TestSetForwardedHeaders:
"""set_forwarded_headers filters to x-* prefixed headers only."""
def test_filters_to_x_prefixed_headers(self):
set_forwarded_headers(
{
"x-aimock-strict": "true",
"x-aimock-session": "abc123",
"x-request-id": "req-456",
"x-custom-trace": "xyz",
"authorization": "Bearer token",
"content-type": "application/json",
}
)
result = get_forwarded_headers()
assert result == {
"x-aimock-strict": "true",
"x-aimock-session": "abc123",
"x-request-id": "req-456",
"x-custom-trace": "xyz",
}
def test_case_insensitive_prefix_match(self):
set_forwarded_headers(
{
"X-AIMock-Strict": "true",
"X-AIMOCK-SESSION": "xyz",
}
)
result = get_forwarded_headers()
assert result == {
"x-aimock-strict": "true",
"x-aimock-session": "xyz",
}
def test_empty_when_no_x_headers(self):
set_forwarded_headers(
{
"authorization": "Bearer token",
"content-type": "application/json",
}
)
result = get_forwarded_headers()
assert result == {}
def test_empty_input(self):
set_forwarded_headers({})
result = get_forwarded_headers()
assert result == {}
class TestGetForwardedHeaders:
"""get_forwarded_headers returns empty dict by default."""
def test_default_is_empty_dict(self):
# Reset to default by running in a fresh context
ctx = contextvars.copy_context()
result = ctx.run(get_forwarded_headers)
assert result == {}
class TestRoundTrip:
"""set + get round-trip."""
def test_round_trip(self):
headers = {"x-aimock-strict": "true", "x-aimock-foo": "bar"}
set_forwarded_headers(headers)
assert get_forwarded_headers() == headers
def test_overwrite(self):
set_forwarded_headers({"x-aimock-a": "1"})
set_forwarded_headers({"x-aimock-b": "2"})
assert get_forwarded_headers() == {"x-aimock-b": "2"}
class TestInstallHttpxHook:
"""install_httpx_hook appends to event hooks."""
def test_appends_to_raw_httpx_client(self):
"""Mock a raw httpx client with event_hooks dict."""
class FakeClient:
def __init__(self):
self.event_hooks = {"request": []}
client = FakeClient()
install_httpx_hook(client)
assert len(client.event_hooks["request"]) == 1
def test_appends_to_sdk_wrapped_client(self):
"""Mock an OpenAI/Anthropic SDK client with _client attribute."""
class FakeTransport:
def __init__(self):
self.event_hooks = {"request": []}
class FakeSDKClient:
def __init__(self):
self._client = FakeTransport()
client = FakeSDKClient()
install_httpx_hook(client)
assert len(client._client.event_hooks["request"]) == 1
def test_hook_injects_headers(self):
"""The installed hook reads from ContextVar and injects headers."""
class FakeHeaders(dict):
"""Dict that also supports item assignment like httpx Headers."""
pass
class FakeRequest:
def __init__(self):
self.headers = FakeHeaders()
class FakeClient:
def __init__(self):
self.event_hooks = {"request": []}
client = FakeClient()
install_httpx_hook(client)
# Set headers in ContextVar
set_forwarded_headers({"x-aimock-strict": "true"})
# Simulate httpx calling the hook
request = FakeRequest()
client.event_hooks["request"][0](request)
assert request.headers["x-aimock-strict"] == "true"
def test_hook_noop_when_no_headers(self):
"""Hook is a no-op when ContextVar is empty (demo traffic)."""
class FakeRequest:
def __init__(self):
self.headers = {}
class FakeClient:
def __init__(self):
self.event_hooks = {"request": []}
# Reset ContextVar to simulate a fresh request with no aimock headers
set_forwarded_headers({})
client = FakeClient()
install_httpx_hook(client)
request = FakeRequest()
client.event_hooks["request"][0](request)
assert request.headers == {}
def test_no_event_hooks_warns(self):
"""Client without event_hooks emits a warning."""
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
install_httpx_hook(object())
assert len(w) == 1
assert "event_hooks" in str(w[0].message)
class TestInstallHttpxHookNestedChain:
"""install_httpx_hook walks the ``._client`` chain (modern OpenAI SDK shape)."""
def test_walks_chain_to_find_event_hooks(self):
"""Modern OpenAI SDK: ChatOpenAI.client -> Resource -> openai.OpenAI
-> ._client = httpx wrapper (event_hooks here). The hook installer
must walk past intermediate ``._client`` hops that do NOT expose
event_hooks, and attach to the first one that does.
"""
class HttpxWrapper:
def __init__(self):
self.event_hooks = {"request": []}
class OpenAIClient:
"""Mirrors openai.OpenAI / openai.AsyncOpenAI: holds an httpx
wrapper at ``._client`` but exposes no event_hooks itself."""
def __init__(self):
self._client = HttpxWrapper()
class Resource:
"""Mirrors openai resources (Completions, AsyncCompletions, etc):
holds the OpenAI client at ``._client``."""
def __init__(self):
self._client = OpenAIClient()
class LangChainWrapper:
"""Mirrors langchain_openai.ChatOpenAI.client: the resource."""
def __init__(self):
self.client = Resource()
outer = LangChainWrapper()
install_httpx_hook(outer.client)
# The hook MUST land on the deepest object that has event_hooks
hooks = outer.client._client._client.event_hooks["request"]
assert len(hooks) == 1, (
f"expected hook installed on deepest httpx wrapper, got "
f"{len(hooks)} hook(s) (chain walk likely stopped too shallow)"
)
def test_idempotent_double_install(self):
"""Calling install_httpx_hook twice must NOT register the hook twice."""
class FakeClient:
def __init__(self):
self.event_hooks = {"request": []}
client = FakeClient()
install_httpx_hook(client)
install_httpx_hook(client)
assert len(client.event_hooks["request"]) == 1, (
"install_httpx_hook must be idempotent — double install detected"
)
def test_header_agnostic_injection(self):
"""Hook must forward whatever headers are in the ContextVar, not just x-aimock-*."""
class FakeRequest:
def __init__(self):
self.headers = {}
class FakeClient:
def __init__(self):
self.event_hooks = {"request": []}
client = FakeClient()
install_httpx_hook(client)
set_forwarded_headers(
{
"x-trace-id": "abc",
"x-team-id": "ck",
"x-anything-custom": "v",
}
)
request = FakeRequest()
client.event_hooks["request"][0](request)
assert request.headers["x-trace-id"] == "abc"
assert request.headers["x-team-id"] == "ck"
assert request.headers["x-anything-custom"] == "v"
class TestInstallHttpxHookAsync:
"""For httpx.AsyncClient instances the installed hook MUST be an
async callable; httpx awaits async-client request hooks."""
def test_async_client_gets_async_hook(self):
"""httpx.AsyncClient -> the installed hook is a coroutine function."""
async def _run():
client = httpx.AsyncClient()
try:
install_httpx_hook(client)
hooks = client.event_hooks["request"]
assert len(hooks) == 1
hook = hooks[0]
assert inspect.iscoroutinefunction(hook), (
f"AsyncClient must receive an async hook, got sync "
f"callable {hook!r}"
)
finally:
await client.aclose()
asyncio.run(_run())
def test_async_hook_is_awaitable_and_injects_headers(self):
"""Awaiting the installed async hook must inject forwarded headers."""
class FakeRequest:
def __init__(self):
self.headers = {}
async def _run():
client = httpx.AsyncClient()
try:
install_httpx_hook(client)
set_forwarded_headers({"x-aimock-strict": "true"})
hook = client.event_hooks["request"][0]
request = FakeRequest()
# Must be awaitable without TypeError
result = hook(request)
assert inspect.isawaitable(result), (
"async-client hook must return an awaitable"
)
await result
assert request.headers["x-aimock-strict"] == "true"
finally:
await client.aclose()
asyncio.run(_run())
def test_sync_client_gets_sync_hook(self):
"""httpx.Client -> the installed hook is a plain sync callable
(httpx calls request hooks synchronously on a sync client)."""
client = httpx.Client()
try:
install_httpx_hook(client)
hooks = client.event_hooks["request"]
assert len(hooks) == 1
hook = hooks[0]
assert not inspect.iscoroutinefunction(hook), (
f"sync Client must receive a sync hook, got coroutine function {hook!r}"
)
finally:
client.close()
class TestInstallHttpxHookRegressions:
"""Regression tests for chain-depth, async/sync MRO heuristic, and
foreign-hook preservation."""
def test_chain_depth_exhausted_warns(self):
"""A ``._client`` chain LONGER than the walker's max depth where NO
node exposes event_hooks must emit a warning (loud-via-warning) and
not silently no-op."""
class ChainNode:
def __init__(self):
self._client: object | None = None # filled in below
# Build a chain of depth well beyond _MAX_CHAIN_DEPTH (5 hops).
# 10 nodes total, none of which carries event_hooks.
nodes = [ChainNode() for _ in range(10)]
for i in range(len(nodes) - 1):
nodes[i]._client = nodes[i + 1]
# Terminate the chain at the last node with a non-None, non-event_hooks
# object so the walker doesn't short-circuit on None.
nodes[-1]._client = object()
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
install_httpx_hook(nodes[0])
assert len(w) >= 1
assert any("event_hooks" in str(item.message) for item in w), (
f"expected a warning mentioning event_hooks, got {[str(i.message) for i in w]}"
)
def test_sync_client_with_async_named_mro_base_is_sync(self):
"""A SYNC duck-typed client whose MRO includes a base class named
``Async*`` (but whose own class name is neither ``AsyncClient`` nor
starts with ``Async``) must be classified as SYNC. An overbroad
``startswith("Async")`` MRO heuristic would misclassify it as async
and install an async hook that httpx calls synchronously -> the
coroutine is never awaited and headers are silently dropped."""
class AsyncMixin:
"""A base class whose NAME starts with ``Async`` but which does
not represent an async client (mirrors e.g. AsyncContextManager
appearing in MRO)."""
pass
class FakeSyncClient(AsyncMixin):
def __init__(self):
self.event_hooks = {"request": []}
client = FakeSyncClient()
install_httpx_hook(client)
hooks = client.event_hooks["request"]
assert len(hooks) == 1
hook = hooks[0]
assert not inspect.iscoroutinefunction(hook), (
f"sync duck-typed client (MRO contains Async*-named base) must "
f"receive a sync hook, got coroutine function {hook!r}"
)
def test_preexisting_foreign_hook_is_preserved(self):
"""If event_hooks['request'] already contains an unrelated callable
(not carrying our idempotency marker), install_httpx_hook must
APPEND ours alongside it — not skip installation, not replace the
foreign hook."""
def foreign_hook(request):
# Unrelated pre-existing hook; carries no marker.
return None
class FakeClient:
def __init__(self):
self.event_hooks = {"request": [foreign_hook]}
client = FakeClient()
install_httpx_hook(client)
hooks = client.event_hooks["request"]
assert len(hooks) == 2, (
f"expected foreign hook + ours = 2 hooks, got {len(hooks)}: {hooks!r}"
)
# Foreign hook still present and unchanged.
assert foreign_hook in hooks
# Exactly one of the two hooks is ours (carries the marker).
from copilotkit.header_propagation import _HOOK_MARKER
marked = [h for h in hooks if getattr(h, _HOOK_MARKER, False)]
assert len(marked) == 1, (
f"expected exactly one hook to carry our marker, got {len(marked)}"
)
class TestContextVarIsolation:
"""ContextVar provides proper isolation across contexts."""
def test_context_isolation(self):
"""Headers set in one context don't leak to another."""
results = {}
def task_a():
set_forwarded_headers({"x-aimock-task": "a"})
results["a"] = get_forwarded_headers()
def task_b():
set_forwarded_headers({"x-aimock-task": "b"})
results["b"] = get_forwarded_headers()
# Run in separate contexts to verify isolation
ctx_a = contextvars.copy_context()
ctx_b = contextvars.copy_context()
ctx_a.run(task_a)
ctx_b.run(task_b)
assert results["a"] == {"x-aimock-task": "a"}
assert results["b"] == {"x-aimock-task": "b"}
def test_child_context_does_not_pollute_parent(self):
"""Setting headers in a child context does not affect the parent."""
# Ensure clean state in a fresh context
def _run():
parent_before = get_forwarded_headers()
def child():
set_forwarded_headers({"x-aimock-child": "yes"})
ctx = contextvars.copy_context()
ctx.run(child)
parent_after = get_forwarded_headers()
assert parent_before == parent_after
contextvars.copy_context().run(_run)
+48
View File
@@ -0,0 +1,48 @@
"""Tests for #3096: copilotkit_interrupt crashes on non-list resume values."""
import pytest
import json
from unittest.mock import patch, MagicMock
class TestInterruptNonListResume:
"""Verify copilotkit_interrupt handles string/dict resume values without crashing."""
@patch("copilotkit.langgraph.interrupt")
def test_string_resume_value(self, mock_interrupt):
"""When LangGraph 1.x returns a string resume value, should not crash."""
mock_interrupt.return_value = "user approved"
from copilotkit.langgraph import copilotkit_interrupt
answer, response = copilotkit_interrupt(message="Do you approve?")
assert answer == "user approved"
assert response == "user approved"
@patch("copilotkit.langgraph.interrupt")
def test_dict_resume_value(self, mock_interrupt):
"""When LangGraph 1.x returns a dict resume value, should return JSON string."""
mock_interrupt.return_value = {"approved": True, "reason": "looks good"}
from copilotkit.langgraph import copilotkit_interrupt
answer, response = copilotkit_interrupt(message="Do you approve?")
parsed = json.loads(answer)
assert parsed["approved"] is True
assert parsed["reason"] == "looks good"
@patch("copilotkit.langgraph.interrupt")
def test_list_resume_value_still_works(self, mock_interrupt):
"""Existing behavior: list resume values should still use [-1].content."""
mock_msg = MagicMock()
mock_msg.content = "yes"
mock_interrupt.return_value = [mock_msg]
from copilotkit.langgraph import copilotkit_interrupt
answer, response = copilotkit_interrupt(message="Do you approve?")
assert answer == "yes"
assert response == [mock_msg]
@@ -0,0 +1,194 @@
"""Tests for multi-part content handling in langchain_messages_to_copilotkit.
Covers the fix in PR #3844 / issue #1748: when AIMessage.content is a list
of content blocks (e.g. Anthropic models), all text parts must be extracted
and concatenated — not just the first element.
"""
import pytest
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from copilotkit.langgraph import langchain_messages_to_copilotkit
class TestMultiPartContentList:
"""AIMessage.content as a list should concatenate all text parts."""
def test_list_of_text_dicts(self):
"""Multiple {"type": "text", "text": "..."} dicts are all concatenated."""
msg = AIMessage(
id="ai-1",
content=[
{"type": "text", "text": "Hello "},
{"type": "text", "text": "world"},
],
)
result = langchain_messages_to_copilotkit([msg])
assert len(result) == 1
assert result[0]["content"] == "Hello world"
assert result[0]["role"] == "assistant"
def test_list_of_strings(self):
"""Content list of plain strings should be concatenated."""
msg = AIMessage(
id="ai-2",
content=["Part A", " Part B"],
)
result = langchain_messages_to_copilotkit([msg])
assert len(result) == 1
assert result[0]["content"] == "Part A Part B"
def test_mixed_strings_and_text_dicts(self):
"""Mix of plain strings and text dicts should all be concatenated."""
msg = AIMessage(
id="ai-3",
content=[
"Start ",
{"type": "text", "text": "middle "},
{"text": "end"},
],
)
result = langchain_messages_to_copilotkit([msg])
assert len(result) == 1
assert result[0]["content"] == "Start middle end"
def test_non_text_parts_are_skipped(self):
"""Non-text content blocks (e.g. images) should be ignored."""
msg = AIMessage(
id="ai-4",
content=[
{"type": "text", "text": "Sample png file"},
{
"type": "image",
"image_data": {"data": "base64data", "format": "image/png"},
},
],
)
result = langchain_messages_to_copilotkit([msg])
assert len(result) == 1
assert result[0]["content"] == "Sample png file"
def test_empty_list_returns_empty_content(self):
"""Empty content list should produce assistant message with empty string."""
msg = AIMessage(
id="ai-5",
content=[],
)
result = langchain_messages_to_copilotkit([msg])
# Assistant messages are always emitted (even with empty content)
# so that tool call entries can reference them via parentMessageId.
assert len(result) == 1
assert result[0]["content"] == ""
assert result[0]["role"] == "assistant"
def test_single_text_dict_in_list(self):
"""Single text dict in a list should still be extracted."""
msg = AIMessage(
id="ai-6",
content=[{"type": "text", "text": "Only one part"}],
)
result = langchain_messages_to_copilotkit([msg])
assert len(result) == 1
assert result[0]["content"] == "Only one part"
def test_dict_without_type_but_with_text_key(self):
"""A dict with "text" key but no "type" should still have text extracted."""
msg = AIMessage(
id="ai-7",
content=[{"text": "no type field"}],
)
result = langchain_messages_to_copilotkit([msg])
assert len(result) == 1
assert result[0]["content"] == "no type field"
class TestSingleDictContent:
"""AIMessage.content as a single dict (Anthropic style) should extract text.
Note: langchain_core.messages.AIMessage validates content as str | list,
so a raw dict cannot be passed directly. We use a mock to exercise the
dict-handling code path in langchain_messages_to_copilotkit, which exists
to handle edge cases from deserialized or non-standard message objects.
"""
def test_dict_with_text_key(self):
"""A content dict with "text" key should have its text extracted."""
from unittest.mock import MagicMock
msg = MagicMock(spec=AIMessage)
msg.content = {"text": "dict content"}
msg.id = "ai-8"
msg.tool_calls = []
result = langchain_messages_to_copilotkit([msg])
assert len(result) == 1
assert result[0]["content"] == "dict content"
class TestPlainStringContent:
"""Standard string content should still work as before."""
def test_plain_string_content(self):
"""Normal string content passes through unchanged."""
msg = AIMessage(
id="ai-9",
content="Just a string",
)
result = langchain_messages_to_copilotkit([msg])
assert len(result) == 1
assert result[0]["content"] == "Just a string"
def test_human_message_string(self):
"""HumanMessage with string content still works."""
msg = HumanMessage(id="human-1", content="Hello")
result = langchain_messages_to_copilotkit([msg])
assert len(result) == 1
assert result[0]["role"] == "user"
assert result[0]["content"] == "Hello"
def test_system_message_string(self):
"""SystemMessage with string content still works."""
msg = SystemMessage(id="sys-1", content="System prompt")
result = langchain_messages_to_copilotkit([msg])
assert len(result) == 1
assert result[0]["role"] == "system"
assert result[0]["content"] == "System prompt"
class TestIssue1748Reproduction:
"""Directly reproduces the scenario from issue #1748.
The original bug: when content is a list of dicts including an image block,
only the first element was taken via `content[0]`, which was the dict itself,
not a string. This caused the message to be silently dropped or mangled.
"""
def test_text_and_image_content_preserves_text(self):
"""The exact scenario from issue #1748: text + image content blocks."""
msg = AIMessage(
id="ai-repro",
content=[
{"type": "text", "text": "Sample png file"},
{
"type": "image",
"image_data": {"data": "aW1hZ2VfZGF0YQ==", "format": "image/png"},
},
],
)
result = langchain_messages_to_copilotkit([msg])
assert len(result) == 1
assert result[0]["content"] == "Sample png file"
assert result[0]["role"] == "assistant"
def test_multiple_text_parts_are_not_truncated(self):
"""The core bug: only the first element was kept. All text must survive."""
msg = AIMessage(
id="ai-trunc",
content=[
{"type": "text", "text": "First part. "},
{"type": "text", "text": "Second part. "},
{"type": "text", "text": "Third part."},
],
)
result = langchain_messages_to_copilotkit([msg])
assert len(result) == 1
assert result[0]["content"] == "First part. Second part. Third part."
@@ -0,0 +1,138 @@
"""Tests for langchain_messages_to_copilotkit assistant message emission.
Covers the parentMessageId orphan bug where an `if content:` guard skipped
emitting the assistant message when content was empty. Tool call entries
reference their parent assistant message via parentMessageId, so the
assistant message must always be emitted — even when content is empty
(standard OpenAI behavior for tool-call-only responses).
"""
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from copilotkit.langgraph import langchain_messages_to_copilotkit
def _convert_and_split(messages):
"""Convert messages and split result into assistant vs tool-call entries."""
result = langchain_messages_to_copilotkit(messages)
assistant_msgs = [m for m in result if m.get("role") == "assistant"]
tool_call_msgs = [m for m in result if "parentMessageId" in m]
return result, assistant_msgs, tool_call_msgs
class TestAssistantMessageAlwaysEmitted:
"""The assistant message must always be present so tool call entries can
reference it via parentMessageId. Without it, tool calls are orphaned
and the frontend cannot reconstruct tool call rendering on reconnect."""
def test_ai_message_with_content_and_tool_calls(self):
"""AIMessage with both content and tool_calls emits assistant + tool calls."""
messages = [
AIMessage(
id="ai-1",
content="Let me help with that.",
tool_calls=[
{"id": "tc-1", "name": "get_help", "args": {"topic": "billing"}}
],
),
]
_, assistant_msgs, tool_call_msgs = _convert_and_split(messages)
assert len(assistant_msgs) == 1
assert assistant_msgs[0]["id"] == "ai-1"
assert assistant_msgs[0]["content"] == "Let me help with that."
assert len(tool_call_msgs) == 1
assert tool_call_msgs[0]["parentMessageId"] == "ai-1"
def test_ai_message_with_empty_content_and_tool_calls(self):
"""AIMessage with empty content (OpenAI-style) still emits the assistant message."""
messages = [
AIMessage(
id="ai-1",
content="",
tool_calls=[
{"id": "tc-1", "name": "get_help", "args": {"topic": "billing"}}
],
),
]
_, assistant_msgs, tool_call_msgs = _convert_and_split(messages)
assert len(assistant_msgs) == 1, (
"Assistant message must be emitted even with empty content"
)
assert assistant_msgs[0]["id"] == "ai-1"
assert assistant_msgs[0]["content"] == ""
assert len(tool_call_msgs) == 1
assert tool_call_msgs[0]["parentMessageId"] == "ai-1"
def test_ai_message_with_none_content_and_tool_calls(self):
"""AIMessage with None content still emits the assistant message."""
msg = AIMessage(
id="ai-1",
content="",
tool_calls=[{"id": "tc-1", "name": "get_help", "args": {}}],
)
# Simulate None content (some models/edge cases)
msg.content = None # type: ignore[assignment]
_, assistant_msgs, _ = _convert_and_split([msg])
assert len(assistant_msgs) == 1, (
"Assistant message must be emitted even with None content"
)
assert assistant_msgs[0]["content"] == ""
def test_no_orphaned_parent_message_ids(self):
"""Every parentMessageId must reference an existing assistant message."""
messages = [
HumanMessage(id="h-1", content="help me"),
AIMessage(
id="ai-1",
content="",
tool_calls=[
{"id": "tc-1", "name": "get_help", "args": {"topic": "billing"}},
{"id": "tc-2", "name": "search", "args": {"query": "docs"}},
],
),
ToolMessage(id="tm-1", content="done", tool_call_id="tc-1"),
ToolMessage(id="tm-2", content="found", tool_call_id="tc-2"),
]
result, _, tool_call_msgs = _convert_and_split(messages)
message_ids = {m["id"] for m in result if "role" in m}
for tc in tool_call_msgs:
assert tc["parentMessageId"] in message_ids, (
f"Tool call {tc['id']} has orphaned parentMessageId {tc['parentMessageId']}"
)
def test_ai_message_with_list_content_and_tool_calls(self):
"""AIMessage with empty list content (Anthropic-style) still emits the assistant message."""
msg = AIMessage(
id="ai-1",
content="",
tool_calls=[{"id": "tc-1", "name": "get_help", "args": {}}],
)
# Anthropic models can return content as a list; empty list is falsy
msg.content = [] # type: ignore[assignment]
_, assistant_msgs, tool_call_msgs = _convert_and_split([msg])
assert len(assistant_msgs) == 1, (
"Assistant message must be emitted even with empty list content"
)
assert assistant_msgs[0]["content"] == ""
assert len(tool_call_msgs) == 1
assert tool_call_msgs[0]["parentMessageId"] == "ai-1"
def test_ai_message_without_tool_calls(self):
"""Plain AIMessage (no tool calls) emits just the assistant message."""
messages = [AIMessage(id="ai-1", content="Hello!")]
result, _, _ = _convert_and_split(messages)
assert len(result) == 1
assert result[0]["role"] == "assistant"
assert result[0]["content"] == "Hello!"
+94
View File
@@ -0,0 +1,94 @@
version = 1
revision = 3
requires-python = ">=3.11"
[manifest]
[manifest.dependency-groups]
dev = [
{ name = "pytest", specifier = ">=9.0.2,<10.0.0" },
{ name = "pytest-asyncio", specifier = ">=1.0.0,<2.0.0" },
]
[[package]]
name = "colorama"
version = "0.4.6"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" },
]
[[package]]
name = "iniconfig"
version = "2.3.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" },
]
[[package]]
name = "packaging"
version = "26.2"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/d7/f1/e7a6dd94a8d4a5626c03e4e99c87f241ba9e350cd9e6d75123f992427270/packaging-26.2.tar.gz", hash = "sha256:ff452ff5a3e828ce110190feff1178bb1f2ea2281fa2075aadb987c2fb221661", size = 228134, upload-time = "2026-04-24T20:15:23.917Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/df/b2/87e62e8c3e2f4b32e5fe99e0b86d576da1312593b39f47d8ceef365e95ed/packaging-26.2-py3-none-any.whl", hash = "sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e", size = 100195, upload-time = "2026-04-24T20:15:22.081Z" },
]
[[package]]
name = "pluggy"
version = "1.6.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
]
[[package]]
name = "pygments"
version = "2.20.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/c3/b2/bc9c9196916376152d655522fdcebac55e66de6603a76a02bca1b6414f6c/pygments-2.20.0.tar.gz", hash = "sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f", size = 4955991, upload-time = "2026-03-29T13:29:33.898Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" },
]
[[package]]
name = "pytest"
version = "9.0.3"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "colorama", marker = "sys_platform == 'win32'" },
{ name = "iniconfig" },
{ name = "packaging" },
{ name = "pluggy" },
{ name = "pygments" },
]
sdist = { url = "https://files.pythonhosted.org/packages/7d/0d/549bd94f1a0a402dc8cf64563a117c0f3765662e2e668477624baeec44d5/pytest-9.0.3.tar.gz", hash = "sha256:b86ada508af81d19edeb213c681b1d48246c1a91d304c6c81a427674c17eb91c", size = 1572165, upload-time = "2026-04-07T17:16:18.027Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/d4/24/a372aaf5c9b7208e7112038812994107bc65a84cd00e0354a88c2c77a617/pytest-9.0.3-py3-none-any.whl", hash = "sha256:2c5efc453d45394fdd706ade797c0a81091eccd1d6e4bccfcd476e2b8e0ab5d9", size = 375249, upload-time = "2026-04-07T17:16:16.13Z" },
]
[[package]]
name = "pytest-asyncio"
version = "1.4.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "pytest" },
{ name = "typing-extensions", marker = "python_full_version < '3.13'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/43/7c/d36d04db312ecf4298932ef77e6e4a9e8ad017906e24e34f0b0c361a2473/pytest_asyncio-1.4.0.tar.gz", hash = "sha256:c6c0d2259945122819f171a32ecea2c349ead889ee28176caaf492143424be42", size = 58514, upload-time = "2026-05-26T09:56:04.083Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/03/e2/08a497ef684b88559c9cc5f4ad53a37e7b99e727094a86d6ea32536d5d3c/pytest_asyncio-1.4.0-py3-none-any.whl", hash = "sha256:933ca923a23075a87fb7070c0ec272a6848489824d887c85c812670932835aa1", size = 16930, upload-time = "2026-05-26T09:56:02.576Z" },
]
[[package]]
name = "typing-extensions"
version = "4.15.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" },
]