How-to: Modify the CrewAI template¶

git clone --branch 11.3.5 --depth 1 https://github.com/datarobot-community/datarobot-agent-templates.git
cd datarobot-agent-templates

cp .env.template .env

task start

task start
task: [start] uv run quickstart.py
*****           *          ****        *             *  
*    *  ***   *****  ***   *   *  ***  ****   ***  *****
*    * *   *    *   *   *  ****  *   * *  *  *   *   *  
*****   *** *   *    *** * *   *  ***  ****   ***    *  

--------------------------------------------------------
           Quickstart for DataRobot AI Agents           
--------------------------------------------------------
Checking environment setup for required pre-requisites...
All pre-requisites are installed.


You will now select an agentic framework to use for this project.
For more information on the different agentic frameworks please go to:
  https://github.com/datarobot-community/datarobot-agent-templates/blob/main/docs/getting-started.md

Please select an agentic framework to use:
1. agent_crewai
2. agent_generic_base
3. agent_langgraph
4. agent_llamaindex
5. agent_nat
Enter your choice (1-5): 1

Would you like to setup the uv python environments and install pre-requisites now?
(y/n): y

task install

task agent:cli START_DEV=1 -- execute --user_prompt 'Hi, how are you?'

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

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


class MyAgent:
    """MyAgent is a custom marketing crew that uses CrewAI to research, write, and edit marketing content.
    It utilizes DataRobot's LLM Gateway or a specific deployment for language model interactions.
    This example illustrates 3 agents that handle marketing tasks: market research, copywriting, and editing.
    """

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

        # Load SERPER_API_KEY from runtime parameters if available
        try:
            from datarobot_drum import RuntimeParameters
            serper_key = RuntimeParameters.get("SERPER_API_KEY")
            if serper_key and serper_key != "SET_VIA_PULUMI_OR_MANUALLY":
                os.environ["SERPER_API_KEY"] = serper_key
        except (ValueError, ImportError):
            pass  # Runtime parameters not available in local development

        # Initialize the search tool
        self.search_tool = SerperDevTool()

    @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 agent_researcher(self) -> Agent:
        """Market Researcher Agent - finds and analyzes target audience and marketing angles."""
        return Agent(
            role='Senior Market Researcher',
            goal='To find and analyze the target audience and marketing angles for a given product',
            backstory="""You are a seasoned market researcher with a knack for uncovering consumer insights.
            You excel at identifying demographics, psychographics, and the key pain points of a target audience.
            Your analysis forms the foundation of any successful marketing campaign.""",
            verbose=self.verbose,
            allow_delegation=False,
            tools=[self.search_tool],
            llm=self.llm
        )

    @property
    def agent_copywriter(self) -> Agent:
        """Marketing Copywriter Agent - writes compelling and persuasive marketing copy."""
        return Agent(
            role='Expert Marketing Copywriter',
            goal='To write compelling and persuasive marketing copy based on market research',
            backstory="""You are a world-class copywriter, known for your ability to weave words that sell.
            You take research and transform it into captivating headlines, engaging social media posts,
            and persuasive product descriptions. Your copy is clear, concise, and converts.""",
            verbose=self.verbose,
            allow_delegation=False,
            llm=self.llm
        )

    @property
    def agent_editor(self) -> Agent:
        """Review and Editing Specialist Agent - reviews and refines marketing copy."""
        return Agent(
            role='Chief Editor and Marketing Strategist',
            goal='To review and refine marketing copy for clarity, tone, brand alignment, and effectiveness',
            backstory="""You are a meticulous editor with a deep understanding of marketing strategy.
            You have a keen eye for detail and ensure that all copy is not only grammatically perfect
            but also aligns with the brand's voice and strategic goals. You provide the final polish.""",
            verbose=self.verbose,
            allow_delegation=False,
            llm=self.llm
        )

    @property
    def task_research(self) -> Task:
        """Research task - analyzes the product and identifies target audience and USPs."""
        return Task(
            description="""Analyze the provided product: '{product_description}'.
            Identify the primary target audience, their key pain points, and desires.
            Research competitors and identify unique selling propositions (USPs).
            The final output should be a concise report with actionable insights.""",
            expected_output="""A report containing:
            1. Target Audience Profile (demographics, interests, online behavior).
            2. Key Pain Points and Desires of the audience.
            3. A list of 3-5 Unique Selling Propositions (USPs) for the product.
            4. Competitor analysis summary.""",
            agent=self.agent_researcher
        )

    @property
    def task_copywriting(self) -> Task:
        """Copywriting task - writes three distinct pieces of marketing copy."""
        return Task(
            description="""Using the market research report, write three distinct pieces of marketing copy for the product: '{product_description}'.
            The copy should be tailored to the identified target audience.
            The pieces are:
            1. A catchy Facebook ad post (approx. 50-80 words).
            2. An engaging email subject line and a short preview text.
            3. A persuasive product description for an e-commerce website (approx. 100-150 words).""",
            expected_output="""A document with three sections:
            1. Facebook Ad Copy.
            2. Email Subject Line and Preview.
            3. E-commerce Product Description.""",
            agent=self.agent_copywriter,
            context=[self.task_research]
        )

    @property
    def task_review(self) -> Task:
        """Review task - reviews and refines the marketing copy."""
        return Task(
            description="""Review the marketing copy provided. Check for grammatical errors, clarity, and brand voice alignment.
            Ensure the copy is persuasive and directly addresses the target audience's pain points identified in the research.
            Provide a final, polished version of all three copy pieces with constructive feedback.""",
            expected_output="""The final, production-ready marketing copy for all three formats,
            along with a brief note on any significant changes made and why.""",
            agent=self.agent_editor,
            context=[self.task_copywriting]
        )

    def crew(self) -> Crew:
        """Creates the marketing crew with all agents and tasks."""
        return Crew(
            agents=[self.agent_researcher, self.agent_copywriter, self.agent_editor],
            tasks=[self.task_research, self.task_copywriting, self.task_review],
            process=Process.sequential,
            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]],
    ]:
        """Run the marketing crew 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 product description 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"]
            # You can use other roles as needed (e.g. "system", "assistant")
            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 dict inputs
        if isinstance(user_prompt_content, str):
            try:
                inputs = json.loads(user_prompt_content)
                if isinstance(inputs, dict) and "product_description" not in inputs:
                    # If it's a dict but doesn't have product_description, use the first value
                    if inputs:
                        inputs = {"product_description": list(inputs.values())[0]}
                    else:
                        inputs = {"product_description": ""}
                elif not isinstance(inputs, dict):
                    inputs = {"product_description": inputs}
            except json.JSONDecodeError:
                inputs = {"product_description": user_prompt_content}
        else:
            inputs = {"product_description": str(user_prompt_content)}

        # Print commands may need flush=True to ensure they are displayed in real-time.
        print("Running marketing crew with inputs:", inputs, flush=True)

        # 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

        pipeline_interactions = self.create_pipeline_interactions_from_events(events)

        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 datarobot_genai.core.chat import initialize_authorization_context

def chat(
    completion_create_params: CompletionCreateParams
    | CompletionCreateParamsNonStreaming
    | CompletionCreateParamsStreaming,
    model: str,
    **kwargs,
) -> Union[CustomModelChatResponse, Iterator[CustomModelStreamingResponse]]:
    # Extract all headers from kwargs
    headers = kwargs.get("headers", {})

    # Access specific X-Untrusted-* headers
    authorization_header = headers.get("X-Untrusted-Authorization")
    request_id = headers.get("X-Untrusted-Request-ID")

    # Use headers in your agent logic
    if authorization_header:
        # Pass to downstream services or tools
        print(f"Processing request with auth: {authorization_header}")

    # Continue with normal agent execution
    initialize_authorization_context(completion_create_params, **kwargs)
    agent = MyAgent(**completion_create_params)
    # ...

task agent:cli START_DEV=1 -- execute --user_prompt "A revolutionary AI-powered fitness tracker that monitors your health"