chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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",
|
||||
]
|
||||
@@ -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}
|
||||
"""
|
||||
@@ -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)),
|
||||
}
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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",
|
||||
}
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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]
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user