How-to: Add Data Registry tools to the CrewAI template¶

import json
import os
from dotenv import load_dotenv
from typing import Any, Generator, Optional, Union
from urllib.parse import urljoin, urlparse

from datarobot_drum import RuntimeParameters

from crewai import LLM, Agent, Crew, Task
from crewai.tools import BaseTool
from crewai_event_listener import CrewAIEventListener
from openai.types.chat import CompletionCreateParams
from ragas import MultiTurnSample
from ragas.messages import AIMessage, HumanMessage, ToolMessage

from datarobot_genai.core.chat.client import ToolClient
from tool_data_registry_search import SearchDataRegistryTool
from tool_data_registry_read import ReadDataRegistryTool

load_dotenv()

class MyAgent:
    """MyAgent is a custom agent that uses CrewAI to search, read, and summarize data from the Data Registry."""

    def __init__(
        self,
        api_key: Optional[str] = None,
        api_base: Optional[str] = None,
        model: Optional[str] = None,
        verbose: Optional[Union[bool, str]] = True,
        timeout: Optional[int] = 90,
        **kwargs: Any,
    ):
        """Initializes the MyAgent class with API key, base URL, model, and verbosity settings.

        Args:
            api_key: Optional[str]: API key for authentication with DataRobot services.
                Defaults to None, in which case it will use the DATAROBOT_API_TOKEN environment variable.
            api_base: Optional[str]: Base URL for the DataRobot API.
                Defaults to None, in which case it will use the DATAROBOT_ENDPOINT environment variable.
            model: Optional[str]: The LLM model to use.
                Defaults to None.
            verbose: Optional[Union[bool, str]]: Whether to enable verbose logging.
                Accepts boolean or string values ("true"/"false"). Defaults to True.
            timeout: Optional[int]: How long to wait for the agent to respond.
                Defaults to 90 seconds.
            **kwargs: Any: Additional keyword arguments passed to the agent.
                Contains any parameters received in the CompletionCreateParams.

        Returns:
            None
        """
        self.api_key = api_key or os.environ.get("DATAROBOT_API_TOKEN")
        self.api_base = (
            api_base
            or os.environ.get("DATAROBOT_ENDPOINT")
            or "https://api.datarobot.com"
        )
        self.model = model
        self.timeout = timeout
        if isinstance(verbose, str):
            self.verbose = verbose.lower() == "true"
        elif isinstance(verbose, bool):
            self.verbose = verbose
        self.event_listener = CrewAIEventListener()

@property
def llm(self) -> LLM:
    """Returns a CrewAI LLM instance configured to use DataRobot's LLM Gateway or a specific deployment.

    For help configuring different LLM backends see:
    https://github.com/datarobot-community/datarobot-agent-templates/blob/main/docs/developing-agents-llm-providers.md
    """
    api_base = urlparse(self.api_base)
    if os.environ.get("LLM_DEPLOYMENT_ID"):
        path = api_base.path
        if "/api/v2/deployments" not in path and "api/v2/genai" not in path:
            # Ensure the API base ends with /api/v2/ for deployments
            if not path.endswith("/"):
                path += "/"
            if not path.endswith("api/v2/"):
                path = urljoin(path, "api/v2/")
        api_base = api_base._replace(path=path)
        api_base_str = api_base.geturl()
        deployment_id = os.environ.get("LLM_DEPLOYMENT_ID")
        return LLM(
            model="openai/gpt-4o-mini",
            api_base=f"{api_base_str}deployments/{deployment_id}/",
            api_key=self.api_key,
            timeout=self.timeout,
        )
    else:
        # Use LLM Gateway
        api_base_str = api_base.geturl()
        if api_base_str.endswith("api/v2/"):
            api_base_str = api_base_str[:-7]  # Remove 'api/v2/'
        elif api_base_str.endswith("api/v2"):
            api_base_str = api_base_str[:-6]  # Remove 'api/v2'
        return LLM(
            model="datarobot/azure/gpt-4o-mini",
            api_base=api_base_str,
            api_key=self.api_key,
            timeout=self.timeout,
        )

@property
def tools_client(self) -> ToolClient:
    return ToolClient(
        api_key=self.api_key,
        base_url=self.api_base,
    )

## Data Registry tools
@property
def tool_data_registry_search(self) -> BaseTool:
    deployment_id = os.environ.get("DATA_REGISTRY_SEARCH_TOOL_DEPLOYMENT_ID")
    if not deployment_id:
        deployment_id = RuntimeParameters.get("DATA_REGISTRY_SEARCH_TOOL_DEPLOYMENT_ID")

    print(
        f"Initializing Data Registry Search Tool with deployment ID: {deployment_id}"
    )
    return SearchDataRegistryTool(
        tool_client=self.tools_client,
        deployment_id=deployment_id
    )

@property
def tool_data_registry_read(self) -> BaseTool:
    deployment_id = os.environ.get("DATA_REGISTRY_READ_TOOL_DEPLOYMENT_ID")
    if not deployment_id:
        deployment_id = RuntimeParameters.get("DATA_REGISTRY_READ_TOOL_DEPLOYMENT_ID")

    print(
        f"Initializing Data Registry Read Tool with deployment ID: {deployment_id}"
    )
    return ReadDataRegistryTool(
        tool_client=self.tools_client,
        deployment_id=deployment_id,
    )

## Data Registry agents and tasks
@property
def agent_data_registry_searcher(self) -> Agent:
    return Agent(
        role="Expert Data Registry Searcher",
        goal="Search for and retrieve relevant files from Data Registry.",
        backstory="You are a meticulous analyst that is skilled at examining lists of files and "
        "determining the most appropriate file based on the context.",
        verbose=self.verbose,
        allow_delegation=False,
        llm=self.llm,
    )

@property
def task_data_registry_search(self) -> Task:
    return Task(
        description=(
            "You should search for a relevant dataset ID in the Data Registry based on the provided dataset topic: "
            "{dataset_topic}."
        ),
        expected_output=(
            "Search for a list of relevant files in the Data Registry and determine the most "
            "relevant dataset ID that matches the given topic. You should return the entire dataset ID."
        ),
        agent=self.agent_data_registry_searcher,
        tools=[self.tool_data_registry_search],
    )

@property
def agent_data_registry_reader(self) -> Agent:
    return Agent(
        role="Expert Data Registry Data Reader and Analyst",
        goal=(
            "Read the data from a file in the Data Registry and summarize the csv data to answer a question."
            "The file is likely relatively large, so you should read it in chunks if necessary to answer the question."
        ),
        backstory="When provided with a dataset name and a dataset ID, you are an expert at using your available"
        "tools to read the data. You always return your responses in a csv style format."
        "You are an expert at analyzing and understanding csv data, especially when comma separated "
        "data is presented to you as a string. When provided with a dataset ID you always use your tools "
        "to read the data, especially data from the Data Registry as you have tools available to you that "
        "can read the data directly from the Data Registry. You take care to not read too much data at once, "
        "and read it in chunks, especially when you don't know the data structure and size in advance. "
        "If necessary, you call available tools many times to read the data in chunks, analyzing the data as you go. "
        "You can quickly analyze the data and provide informative insights to help users understand "
        "the data better. You can also answer questions about the data in a concise manner.",
        verbose=self.verbose,
        allow_delegation=False,
        llm=self.llm,
    )

@property
def task_data_registry_read(self) -> Task:
    return Task(
        description=(
            "Given a dataset ID you should use the dataset ID to read the data from the Data Registry. You then use this "
            "data to answer the question: {question}. "
            "It is critically important that you answer this question. The accuracy of your answer is "
            "paramount. You should always provide a comprehensive and verbose answer."
        ),
        expected_output=(
            "You should return a comprehensive and verbose answer to the question."
        ),
        agent=self.agent_data_registry_reader,
        tools=[self.tool_data_registry_read],
    )

@property
def agent_response_editor(self) -> Agent:
    return Agent(
        role="Expert Editor",
        goal="Summarize a verbose response provided by an Expert Data Reader and Analyst into a concise "
        "and clear answer.",
        backstory="You are an expert editor with a keen eye for detail. When provided with a verbose response from "
        "an Expert Data Reader and Analyst, your job is to summarize the response into a concise "
        "and clear answer. You should focus on clarity and brevity, ensuring that the final output "
        "is easy to understand and directly answers the question posed. You are skilled at distilling "
        "complex information into simple, digestible formats. You always ensure that the final output "
        "is accurate and reflects the key points from the original response.",
        verbose=self.verbose,
        llm=self.llm,
    )

@property
def task_response_edit(self) -> Task:
    return Task(
        description=(
            "You are given a verbose summary that is supposed to answer the question: {question}. You must ensure "
            "that the final output is concise and directly answers the question. It is critically important that "
            "you summarize the verbose response into a clear and concise answer."
        ),
        expected_output="You should return a concise summary answer to the question.",
        agent=self.agent_response_editor,
    )

## Data Registry crew
def crew(self) -> Crew:
    return Crew(
        agents=[self.agent_data_registry_searcher, self.agent_data_registry_reader, self.agent_response_editor],
        tasks=[self.task_data_registry_search, self.task_data_registry_read, self.task_response_edit],
        verbose=self.verbose,
    )

async def invoke(
    self, completion_create_params: CompletionCreateParams
) -> Union[
    Generator[tuple[str, Any | None, dict[str, int]], None, None],
    tuple[str, Any | None, dict[str, int]],
]:
    """Invoke the agent with the provided completion parameters.

    [THIS METHOD IS REQUIRED FOR THE AGENT TO WORK WITH DRUM SERVER]

    Args:
        completion_create_params: The completion request parameters including input topic and settings.
    Returns:
        Union[
            Generator[tuple[str, Any | None, dict[str, int]], None, None],
            tuple[str, Any | None, dict[str, int]],
        ]: For streaming requests, returns a generator yielding tuples of (response_text, pipeline_interactions, usage_metrics).
           For non-streaming requests, returns a single tuple of (response_text, pipeline_interactions, usage_metrics).

    """
    # Retrieve the starting user prompt from the CompletionCreateParams
    user_messages = [
        msg
        for msg in completion_create_params["messages"]
        if msg.get("role") == "user"
    ]
    user_prompt: Any = user_messages[0] if user_messages else {}
    user_prompt_content = user_prompt.get("content", "")

    # Handle both string and JSON inputs
    if isinstance(user_prompt_content, str):
        try:
            inputs = json.loads(user_prompt_content)
            if isinstance(inputs, dict) and "dataset_topic" not in inputs:
                # If it's a dict but doesn't have our expected keys, use the first value
                if inputs:
                    first_key, first_value = next(iter(inputs.items()))
                    inputs = {
                        "dataset_topic": first_value,
                        "question": first_value
                    }
                else:
                    # Fallback: use user_prompt_content or empty string
                    inputs = {
                        "dataset_topic": user_prompt_content if user_prompt_content else "",
                        "question": user_prompt_content if user_prompt_content else ""
                    }
            elif not isinstance(inputs, dict):
                inputs = {
                    "dataset_topic": inputs,
                    "question": inputs
                }
        except json.JSONDecodeError:
            inputs = {
                "dataset_topic": user_prompt_content,
                "question": user_prompt_content
            }
    else:
        inputs = {
            "dataset_topic": str(user_prompt_content),
            "question": str(user_prompt_content)
        }

    print("Running agent with inputs:", inputs)

    # Run the crew with the inputs
    crew_output = self.crew().kickoff(inputs=inputs)

    # Extract the response text from the crew output
    response_text = str(crew_output.raw)

    # Create a list of events from the event listener
    events = self.event_listener.messages
    if len(events) > 0:
        last_message = events[-1].content
        if last_message != response_text:
            events.append(AIMessage(content=response_text))
    else:
        events = None
    # Create pipeline interactions from events
    pipeline_interactions = self.create_pipeline_interactions_from_events(events) if events else None

    # Create usage metrics from crew output
    usage_metrics = {
        "completion_tokens": crew_output.token_usage.completion_tokens,
        "prompt_tokens": crew_output.token_usage.prompt_tokens,
        "total_tokens": crew_output.token_usage.total_tokens,
    }

    return response_text, pipeline_interactions, usage_metrics

@staticmethod
def create_pipeline_interactions_from_events(
    events: list[Union[HumanMessage, AIMessage, ToolMessage]],
) -> MultiTurnSample | None:
    """Convert a list of events into a MultiTurnSample.

    Creates the pipeline interactions for moderations and evaluation
    (e.g. Task Adherence, Agent Goal Accuracy, Tool Call Accuracy)
    """
    if not events:
        return None
    return MultiTurnSample(user_input=events)

from abc import ABC
from crewai.tools import BaseTool
from datarobot_genai.core.chat.client import ToolClient

class BaseToolWithDeployment(BaseTool, ABC):
    model_config = {
        "arbitrary_types_allowed": True
    }
    """Adds support for arbitrary types in Pydantic models, needed for the ToolClient."""

    tool_client: ToolClient
    """The tool client initialized by the agent, which has access to its authorization context."""

    deployment_id: str
    """The DataRobot deployment ID of the custom model that will execute tool logic."""

import json
from typing import Dict, Type
from typing import List

from pydantic import BaseModel as PydanticBaseModel, Field

from tool_deployment import BaseToolWithDeployment

class SearchDataRegistryArgs(PydanticBaseModel):
    search_terms: str = Field(
        default="",
        description="Terms for the search. Leave blank to return all datasets."
    )
    limit: int = Field(
        default=2,
        description="The maximum number of datasets to return. Set to -1 to return all."
    )

class SearchDataRegistryTool(BaseToolWithDeployment):
    name: str = "Search Data Registry"
    description: str = (
        "This tool provides a list of all available dataset names and their associated IDs from the Data Registry. "
        "You should always check to see if the dataset you are looking for can be found here. "
        "For future queries, you should use the associated dataset ID instead of the name to avoid ambiguity."
    )
    args_schema: Type[PydanticBaseModel] = SearchDataRegistryArgs

    def _run(self, search_terms: str = "", limit: int = 2) -> List[Dict[str, str]]:

        request_payload = {"search_terms": search_terms, "limit": limit}

        result = self.tool_client.call(
            self.deployment_id,
            payload=request_payload,
        )

        return json.loads(result.data).get("datasets", [])

from io import BytesIO
from typing import Type, Union

import pandas as pd
from pydantic import BaseModel as PydanticBaseModel, Field
from tool_deployment import BaseToolWithDeployment

class ReadDataRegistryArgs(PydanticBaseModel):
    dataset_id: str = Field(
        ...,
        description="The ID of the dataset to read from the Data Registry."
    )
    offset: int = Field(
        default=0,
        description="The first row to start reading the dataset from."
    )
    limit: int | None = Field(
        default=20,
        description=(
            "The number of rows to read from the dataset. If not set, will return all rows."
        )
    )

class ReadDataRegistryTool(BaseToolWithDeployment):
    name: str = "Read Data Registry Dataset"
    description: str = (
        "This tool reads the contents of a Data Registry dataset when given the `dataset_id` input parameter. "
        "Use this tool whenever you need to read the contents of a dataset. "
        "Optionally, you can specify `offset` and `limit` to read a slice of the dataset rows."
    )
    args_schema: Type[PydanticBaseModel] = ReadDataRegistryArgs

    def _get_dataframe_from_response(
        self, data: Union[bytes, pd.DataFrame], headers: dict[str, str]
    ) -> pd.DataFrame:
        """Convert response data into a DataFrame.

        Args:
            data (Union[bytes, pd.DataFrame]): The response data, either as bytes or a DataFrame.
            headers (dict[str, str]): The response headers.

        Returns:
            pd.DataFrame: The parsed DataFrame.

        Raises:
            ValueError: If the response data is neither bytes nor a DataFrame.
        """
        if isinstance(data, bytes) and headers.get("Content-Type", "").lower().startswith("text/csv"):
            return pd.read_csv(BytesIO(data))
        if isinstance(data, pd.DataFrame):
            return data
        raise ValueError("The response data must be either bytes or a DataFrame.")

    def _run(self, dataset_id: str, offset: int=0, limit: int | None = None) -> str:
        if not dataset_id:
            raise ValueError("dataset_id is required but was not provided")

        rv = self.tool_client.call(self.deployment_id, payload={
            "dataset_id": dataset_id,
            "offset": offset,
            "limit": limit,
        }, )
        try:
            return self._get_dataframe_from_response(rv.data, rv.response_headers).to_csv()
        except Exception as e:
            raise ValueError(
                f"Could not read dataset with dataset_id '{dataset_id}'. "
                f"Please verify that the dataset_id exists and you have access to it. Error: {e}"
            )