LLMs

API-based tool-calling LLMs

These pipeline elements use tool-calling LLMs via the official tool-calling API.

OpenAILLM

OpenAILLM(
    client: OpenAI,
    model: str,
    reasoning_effort: (
        ChatCompletionReasoningEffort | None
    ) = None,
    temperature: float | None = 0.0,
)

Bases: BasePipelineElement

LLM pipeline element that uses OpenAI's API.

Parameters:

• client (OpenAI) –

  The OpenAI client.

• model (str) –

  The model name.

• temperature (float | None, default: 0.0 ) –

  The temperature to use for generation.

Source code in src/agentdojo/agent_pipeline/llms/openai_llm.py

def __init__(
    self,
    client: openai.OpenAI,
    model: str,
    reasoning_effort: ChatCompletionReasoningEffort | None = None,
    temperature: float | None = 0.0,
) -> None:
    self.client = client
    self.model = model
    self.temperature = temperature
    self.reasoning_effort: ChatCompletionReasoningEffort | None = reasoning_effort

AnthropicLLM

AnthropicLLM(
    client: Anthropic,
    model: str,
    temperature: float | None = 0.0,
    thinking_budget_tokens: int | None = None,
    max_tokens: int = 1024,
)

AnthropicLLM(
    client: AsyncAnthropic,
    model: str,
    temperature: float | None = 0.0,
    thinking_budget_tokens: int | None = None,
    max_tokens: int = 1024,
)

AnthropicLLM(
    client: Anthropic | AsyncAnthropic,
    model: str,
    temperature: float | None = 0.0,
    thinking_budget_tokens: int | None = None,
    max_tokens: int = 1024,
)

Bases: BasePipelineElement

LLM pipeline element using Anthropic's API. Supports Claude Opus, Sonnet, and Haiku models.

Parameters:

• client (Anthropic | AsyncAnthropic) –

  Anthropic client instance.

• model (str) –

  Model name.

• temperature (float | None, default: 0.0 ) –

  Sampling temperature.

Source code in src/agentdojo/agent_pipeline/llms/anthropic_llm.py

def __init__(
    self,
    client: Anthropic | AsyncAnthropic,
    model: str,
    temperature: float | None = 0.0,
    thinking_budget_tokens: int | None = None,
    max_tokens: int = 1024,
) -> None:
    if isinstance(client, Anthropic):
        self.client = AsyncAnthropic(
            api_key=client.api_key,
            auth_token=client.auth_token,
            base_url=client.base_url,
            timeout=client.timeout,
            max_retries=client.max_retries,
        )
    else:
        self.client = client
    self.model = model
    self.thinking_budget_tokens = thinking_budget_tokens
    if self.thinking_budget_tokens is not None:
        # Temperature must be 1 when using thinking.
        self.temperature = 1
    else:
        self.temperature = temperature
    if thinking_budget_tokens is not None:
        self._MAX_TOKENS = max_tokens + thinking_budget_tokens
    else:
        self._MAX_TOKENS = max_tokens

GoogleLLM

GoogleLLM(
    model: str,
    client: None = None,
    temperature: float | None = 0.0,
)

GoogleLLM(
    model: str,
    client: Client,
    temperature: float | None = 0.0,
)

GoogleLLM(
    model: str,
    client: Client | None = None,
    temperature: float | None = 0.0,
)

Bases: BasePipelineElement

LLM pipeline element that uses Google Vertex AI (i.e. Gemini models).

Warning

In order to use Google LLMs with VertexAI and not the AI Studio API (by specifying an API key in the client), you need to run vertexai.init() before instantiating this class. This is done automatically if you create a pipeline with AgentPipeline.from_config. However, you need to have Google Cloud CLI installed on your machine and be authenticated. After being authnticated, you need to run gcloud auth application-default login in the terminal to let Python access your Google Cloud account.

Parameters:

• model (str) –

  The model name.

• client (Client | None, default: None ) –

  The google-genai client to use. Not passing this argument (or passing None) is deprecated.

• temperature (float | None, default: 0.0 ) –

  The temperature to use for generation.

Source code in src/agentdojo/agent_pipeline/llms/google_llm.py

def __init__(self, model: str, client: genai.Client | None = None, temperature: float | None = 0.0):
    self.model = model
    if client is not None:
        self.client = client
    else:
        warnings.warn(
            "Not specifying a client for `GoogleLLM` is deprecated and will be removed in a future release."
        )
        self.client = genai.Client(
            vertexai=True, project=os.getenv("GCP_PROJECT"), location=os.getenv("GCP_LOCATION")
        )
    self.temperature = temperature

CohereLLM

CohereLLM(
    client: Client,
    model: str,
    temperature: float = 0.0,
    seed: int | None = None,
)

Bases: BasePipelineElement

LLM pipeline element using Cohere's API.

Parameters:

• client (Client) –

  Cohere client

• model (str) –

  Model name

• temperature (float, default: 0.0 ) –

  Temperature for sampling

• seed (int | None, default: None ) –

  Seed for sampling

Source code in src/agentdojo/agent_pipeline/llms/cohere_llm.py

def __init__(self, client: cohere.Client, model: str, temperature: float = 0.0, seed: int | None = None) -> None:
    self.client = client
    self.model = model
    self.temperature = temperature
    self.seed = seed

Prompt-based tool-calling LLMs

These pipeline elements use tool-calling LLMs via a system prompt that instructs them to use tools, and provides them documentation as part of the system prompt.

BasePromptingLLM

BasePromptingLLM(
    client: OpenAI,
    model: str,
    reasoning_effort: (
        ChatCompletionReasoningEffort | None
    ) = None,
    temperature: float | None = 0.0,
)

Bases: OpenAILLM

OpenAI client-based LLM that can be prompted with a different system prompt. We use this in our paper to run Llama 3 70B via TogetherAI's API.

Source code in src/agentdojo/agent_pipeline/llms/openai_llm.py

def __init__(
    self,
    client: openai.OpenAI,
    model: str,
    reasoning_effort: ChatCompletionReasoningEffort | None = None,
    temperature: float | None = 0.0,
) -> None:
    self.client = client
    self.model = model
    self.temperature = temperature
    self.reasoning_effort: ChatCompletionReasoningEffort | None = reasoning_effort

_tool_calling_prompt instance-attribute

_tool_calling_prompt: str

The system prompt that instructs the LLM on how to use the tools. It should also provide an implementation. Must be defined by subclasses

_MAX_ATTEMPTS class-attribute instance-attribute

_MAX_ATTEMPTS = 3

Maximum number of attempts at parsing the model output.

_make_tools_prompt abstractmethod

_make_tools_prompt(
    system_message: ChatSystemMessage | None,
    tools: Sequence[Function],
) -> ChatSystemMessage | None

Builds the system prompt by putting the original system message, the tools documentation and the instruction prompt together.

See PromptingLLM's implementation for an example. Must be implemented by subclass.

Parameters:

• system_message (ChatSystemMessage | None) –

  The original system message.

• tools (Sequence[Function]) –

  The tools that the LLM can use.

Returns:

• ChatSystemMessage | None –

  The system message with the tools documentation and the instruction prompt.

Source code in src/agentdojo/agent_pipeline/llms/prompting_llm.py

@abc.abstractmethod
def _make_tools_prompt(
    self, system_message: ChatSystemMessage | None, tools: Sequence[Function]
) -> ChatSystemMessage | None:
    """Builds the system prompt by putting the original system message, the tools documentation and
    the instruction prompt together.

    See [`PromptingLLM`][agentdojo.agent_pipeline.PromptingLLM._make_tools_prompt]'s implementation
    for an example. Must be implemented by subclass.

    Args:
        system_message: The original system message.
        tools: The tools that the LLM can use.

    Returns:
        The system message with the tools documentation and the instruction prompt.
    """
    raise NotImplementedError()

_parse_model_output abstractmethod

_parse_model_output(
    message: ChatCompletionMessage,
) -> ChatAssistantMessage

Parses the model output by identifying text content and tool-calling content.

See PromptingLLM's implementation for an example. Must be implemented by subclass.

Parameters:

• message (ChatCompletionMessage) –

  The model output to be parsed.

Returns:

• ChatAssistantMessage –

  The assistant message with the text content and the tool-calling content.

Source code in src/agentdojo/agent_pipeline/llms/prompting_llm.py

@abc.abstractmethod
def _parse_model_output(self, message: ChatCompletionMessage) -> ChatAssistantMessage:
    """Parses the model output by identifying text content and tool-calling content.

    See [`PromptingLLM`][agentdojo.agent_pipeline.PromptingLLM._parse_model_output]'s
    implementation for an example. Must be implemented by subclass.

    Args:
        message: The model output to be parsed.

    Returns:
        The assistant message with the text content and the tool-calling content.
    """
    raise NotImplementedError()

_tool_message_to_user_message abstractmethod

_tool_message_to_user_message(
    tool_message: ChatToolResultMessage,
) -> ChatUserMessage

Parses the model output by identifying text content and tool-calling content.

See PromptingLLM's implementation for an example. Must be implemented by subclass.

Parameters:

• tool_message (ChatToolResultMessage) –

  The tool message whose output should be converted to a user message.

Returns:

• ChatUserMessage –

  The user message containing the tool result.

Source code in src/agentdojo/agent_pipeline/llms/prompting_llm.py

@abc.abstractmethod
def _tool_message_to_user_message(self, tool_message: ChatToolResultMessage) -> ChatUserMessage:
    """Parses the model output by identifying text content and tool-calling content.

    See [`PromptingLLM`][agentdojo.agent_pipeline.PromptingLLM._tool_message_to_user_message]'s
    implementation for an example. Must be implemented by subclass.

    Args:
        tool_message: The tool message whose output should be converted to a user message.

    Returns:
        The user message containing the tool result.
    """
    ...

PromptingLLM

PromptingLLM(
    client: OpenAI,
    model: str,
    reasoning_effort: (
        ChatCompletionReasoningEffort | None
    ) = None,
    temperature: float | None = 0.0,
)

Bases: BasePromptingLLM

Function calling for Llama 3 based on hamelsmu's prompt.

Source code in src/agentdojo/agent_pipeline/llms/openai_llm.py

def __init__(
    self,
    client: openai.OpenAI,
    model: str,
    reasoning_effort: ChatCompletionReasoningEffort | None = None,
    temperature: float | None = 0.0,
) -> None:
    self.client = client
    self.model = model
    self.temperature = temperature
    self.reasoning_effort: ChatCompletionReasoningEffort | None = reasoning_effort

_make_tools_prompt

_make_tools_prompt(
    system_message: ChatSystemMessage | None,
    tools: Sequence[Function],
) -> ChatSystemMessage | None

It creates a system message by instructing the model on how to use tools and provides the tools' documentation, placed within the tags. The tool calling prompt is stored in the _tool_calling_prompt attribute.

It preserves the original system message, by placing it at the end of the tool calling prompt.

Parameters:

• system_message (ChatSystemMessage | None) –

  The system message to be adapted.

Returns:

• ChatSystemMessage | None –

  The adapted system message.

Source code in src/agentdojo/agent_pipeline/llms/prompting_llm.py

def _make_tools_prompt(
    self, system_message: ChatSystemMessage | None, tools: Sequence[Function]
) -> ChatSystemMessage | None:
    """It creates a system message by instructing the model on how to use tools and
    provides the tools' documentation, placed within the <function-definitions> tags.
    The tool calling prompt is stored in the
    [`_tool_calling_prompt`][agentdojo.agent_pipeline.PromptingLLM._tool_calling_prompt] attribute.

    It preserves the original system message,  by placing it at the end of the tool calling prompt.

    Args:
        system_message: The system message to be adapted.

    Returns:
        The adapted system message.
    """
    if len(tools) == 0:
        return system_message
    tools_docs = ""
    system_prompt = get_text_content_as_str(system_message["content"]) if system_message is not None else ""
    for index, tool in enumerate(tools, start=1):
        tool_dict = {
            "name": tool.name,
            "description": tool.description,
            "parameters": tool.parameters.model_json_schema(),
        }
        tools_docs += f"<function-{index}>\n"
        tools_docs += json.dumps(tool_dict, indent=4)
        tools_docs += f"\n</function-{index}>\n\n"
    tool_calling_prompt = self._tool_calling_prompt.format(funcs=tools_docs)
    message_content = f"{tool_calling_prompt}\n{system_prompt}"
    return ChatSystemMessage(role="system", content=[text_content_block_from_string(message_content)])

_parse_model_output

_parse_model_output(
    message: ChatCompletionMessage,
) -> ChatAssistantMessage

Parses the model output by extracting text and/or tool call contents from the message.

It looks for the function call content within the <function-call> tags and extracts it. Each function call is expected to look like a python function call with parameters specified by name. For example, calling the function func1 with parameters a=1 and b=3 would look like:

<function-call>func1(a=1, b=3)</function-call>

Content related to the LLM's thoughts are expected to be in the <function-thoughts> tags and are returned as part of the assistant message's content.

If no function call is done, the answer is expected to be in the <answer> tags.

Parameters:

• message (ChatCompletionMessage) –

  The model output message in OpenAI format.

Returns:

• ChatAssistantMessage –

  The assistant message with the extracted text and tool calls.

Source code in src/agentdojo/agent_pipeline/llms/prompting_llm.py

def _parse_model_output(self, message: ChatCompletionMessage) -> ChatAssistantMessage:
    """Parses the model output by extracting text and/or tool call contents from the message.

    It looks for the function call content within the `<function-call>` tags and extracts it. Each
    function call is expected to look like a python function call with parameters specified by name.
    For example, calling the function `func1` with parameters `a=1` and `b=3` would look like:

        <function-call>func1(a=1, b=3)</function-call>

    Content related to the LLM's thoughts are expected to be in the `<function-thoughts>` tags and are
    returned as part of the assistant message's `content`.

    If no function call is done, the answer is expected to be in the `<answer>` tags.

    Args:
        message: The model output message in OpenAI format.

    Returns:
        The assistant message with the extracted text and tool calls.
    """
    if message.content is None:
        return ChatAssistantMessage(
            role="assistant",
            content=[text_content_block_from_string("")],
            tool_calls=None,
        )
    tool_call_pattern = re.compile(r"<function-call>(.*?)</function-call>", re.DOTALL)
    tool_call_match = tool_call_pattern.search(message.content)

    # Extract the function call content
    tool_call_content = tool_call_match.group(1) if tool_call_match else "[]"
    # Remove the function call section from the original text
    outside_content = (
        re.sub(r"<function-call>.*?</function-call>", "", message.content, flags=re.DOTALL)
        .replace("<function-thoughts>", "")
        .replace("</function-thoughts>", "")
        .strip()
    )
    outside_content = re.sub(r"\n\s*\n", "\n\n", outside_content)
    tool_calls = parse_tool_calls_from_python_function(tool_call_content)
    for tool_call in tool_calls:
        args = {
            arg_name: ("..." if arg_value == Ellipsis else arg_value)
            for arg_name, arg_value in tool_call.args.items()
        }
        tool_call.args = args
    if len(tool_calls) == 0:
        answer_pattern = re.compile(r"<answer>(.*?)</answer>", re.DOTALL)
        answer_match = answer_pattern.search(outside_content)
        if answer_match is None:
            raise InvalidModelOutputError("The answer should be in <answer> tags if no tool calls are provided.")
        outside_content = answer_match.group(1)

    return ChatAssistantMessage(
        role="assistant",
        content=[text_content_block_from_string(outside_content)],
        tool_calls=tool_calls,
    )

_tool_message_to_user_message

_tool_message_to_user_message(
    tool_message: ChatToolResultMessage,
) -> ChatUserMessage

It places the output of the tool call in the tags and the tool call in the tags. If an error is returned, it places the error in the tags and the tool call in the tags.""

Source code in src/agentdojo/agent_pipeline/llms/prompting_llm.py

def _tool_message_to_user_message(self, tool_message: ChatToolResultMessage) -> ChatUserMessage:
    """It places the output of the tool call in the <function-result> tags and the tool call in the <function-call> tags.
    If an error is returned, it places the error in the <function-error> tags and the tool call in the <function-call> tags.""
    """
    function_call_signature = create_python_function_from_tool_call(tool_message["tool_call"])
    function_call = f"<function-call>{function_call_signature}</function-call>"
    if tool_message["error"] is None:
        tool_result = f"<function-result>{get_text_content_as_str(tool_message['content'])}</function-result>"
    else:
        tool_result = f"<function-error>{tool_message['error']}</function-error>"
    return ChatUserMessage(
        role="user",
        content=[text_content_block_from_string(f"{function_call}{tool_result}")],
    )

_tool_calling_prompt class-attribute instance-attribute

_tool_calling_prompt = 'You are a hepful assistant. You are given a task and a set of possible functions inside <function-definitions> tags.\nCalling these functions are optional. Carefully consider the question and determine if one or more functions can be used to answer the question. Place your thoughts and reasoning behind your decision in <function-thoughts> tags.\nIf the given question lacks the parameters required by the function, point it out in <function-thoughts> tags. Below is a list of function definitions:\n<function-definitions>\n{funcs}\n</function-definitions>\n\nIf you wish to call a particular function, specify the name of the function and any arguments in a way that conforms to that function\'s schema inside <function-call> tags.\nFunction calls should be in this format: <function-thoughts>Calling func1 would be helpful because of ...</function-thoughts><function-call>[function1(a="1", b=3), function2(foo=["a", "b"], bar=None)]</function-call>, WITHOUT any answer. Pass the arguments in correct format, i.e., strings should be enclosed in quotes, lists should be enclosed in square brackets, integers should have no quotes, etc. If you do not wish to call any functions, say so in the <function-thoughts> tags followed by <function-call>[]</function-call><answer>...</answer>\nIf no tools are provided, act as if no tools are available to you, but still provide <function-call>[]</function-call> as part of your output.\n\nIf and only if NO function calls are made, answer the question to the best of your ability inside <answer> tags.  If you are unsure of the answer, say so in <answer> tags.\n\nThe user will provide the output of the function calls in the <function-result> tags. The function call that output the given result is provided in <function-call> tags. Give your reply based on the output of the function calls. If no function calls are made, the user will provide the answer without any tags.\nIf the tool returned an error, then the user will provide the error message in <function-error> tags. The function call that returned the error is provided in <function-call> tags. If the tool returns an error, call the tool again with the correct parameters. If the tool returns an error again, without asking the user. You might need to call other tools to solve the task. Be persistent!\n\nSometimes, you may need to call multiple functions in multiple rounds to get the desired output, based on the results of the previous function calls. When this is the case, then use the <function-thoughts> tags to explain your reasoning for calling the next function, and the <function-call> tags to specify the next function call. \nOnce you receive the result of one tool, think if you need to call another tool in the <function-thoughts> tags. If you do, then call the next tool. If you are done and do not need to call another tool, then do explicitly give the answer in <answer> tags. Tell the answer to the user in the message if you don\'t need other tools.\n\nIf you think you will need to call multiple tools in multiple stages, but you don\'t have yet the information to call all of them, then wait to receive the information back from the tools you can already call.'