wehub-resource-sync
679 lines
21 KiB
Python
679 lines
21 KiB
Python
"""
|
|
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
|