crewai_agent.py

"""
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)