wehub-resource-sync
518 lines
15 KiB
Python
518 lines
15 KiB
Python
"""
|
|
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)
|