Skip to main content

Command Palette

Search for a command to run...

Agentic workflow Example - with litellm

Published
4 min read

Introduction:

Create two agents which talks to each other and comes to conclusion

  1. Native Tool Usage: It uses litellm's built-in tool handling (compatible with OpenAI, Anthropic, etc.), rather than manual string parsing.

  2. Two Agents, Two Tools: As requested, Agent A and Agent B each have their own unique tool.

  3. Shared History: The agents "talk" by appending messages to a shared conversation list.

  4. Clear Conclusion: The loop ends automatically when Agent A decides the plan is complete.

Prerequisites

You need litellm installed and an API key (e.g., OpenAI).

pip install litellm
export OPENAI_API_KEY="your-api-key"

Python Code Example

In this scenario, Agent A (Planner) wants to plan a road trip and uses a tool to check distances. Agent B (Accountant) calculates costs and uses a calculator tool.

import json
import os
from litellm import completion

# --- 1. Define the Tools ---

# Tool for Agent A
def get_distance(city_a, city_b):
    """Returns mock distance in miles between two cities."""
    # In a real app, you'd call a maps API here
    return json.dumps({"distance": 300, "unit": "miles"})

# Tool for Agent B
def calculate_fuel_cost(distance, mpg, gas_price):
    """Calculates fuel cost."""
    cost = (float(distance) / float(mpg)) * float(gas_price)
    return json.dumps({"total_cost": round(cost, 2), "currency": "USD"})

# Tool Schemas (OpenAI format)
tools_a = [{
    "type": "function",
    "function": {
        "name": "get_distance",
        "description": "Get distance between two cities",
        "parameters": {
            "type": "object",
            "properties": {
                "city_a": {"type": "string"},
                "city_b": {"type": "string"}
            },
            "required": ["city_a", "city_b"]
        }
    }
}]

tools_b = [{
    "type": "function",
    "function": {
        "name": "calculate_fuel_cost",
        "description": "Calculate fuel cost based on distance and vehicle stats",
        "parameters": {
            "type": "object",
            "properties": {
                "distance": {"type": "number"},
                "mpg": {"type": "number"},
                "gas_price": {"type": "number"}
            },
            "required": ["distance", "mpg", "gas_price"]
        }
    }
}]

# Map names to actual python functions for execution
available_functions = {
    "get_distance": get_distance,
    "calculate_fuel_cost": calculate_fuel_cost
}

# --- 2. Agent Helper Function ---

def run_agent_turn(agent_name, system_prompt, history, tools):
    """
    Runs a single turn for an agent.
    1. Calls the LLM.
    2. Checks if the LLM wants to run a tool.
    3. If yes, runs the tool and sends result back to LLM to interpret.
    4. Returns the final text response.
    """

    # Prepare messages: System prompt + Conversation History
    messages = [{"role": "system", "content": system_prompt}] + history

    response = completion(
        model="gpt-4o",  # or "claude-3-opus", "gpt-3.5-turbo", etc.
        messages=messages,
        tools=tools,
        tool_choice="auto"
    )

    message = response.choices[0].message

    # Check if the agent wants to use a tool
    if message.tool_calls:
        print(f"   🛠️  {agent_name} is using a tool...")

        # Add the "assistant" message with tool_calls to history so the model knows it asked
        messages.append(message)

        for tool_call in message.tool_calls:
            function_name = tool_call.function.name
            function_args = json.loads(tool_call.function.arguments)

            # Execute the python function
            function_response = available_functions[function_name](**function_args)

            # Add the tool result to messages
            messages.append({
                "tool_call_id": tool_call.id,
                "role": "tool",
                "name": function_name,
                "content": function_response,
            })

        # Get the agent's interpretation of the tool result
        second_response = completion(
            model="gpt-4o",
            messages=messages
        )
        return second_response.choices[0].message.content

    return message.content

# --- 3. Main Conversation Loop ---

def main():
    # Shared history of the conversation
    history = []

    # Task to start the conversation
    task = "Plan a budget for a road trip from New York to Boston. Assume 25 MPG and gas is $3.50/gallon."
    history.append({"role": "user", "content": task})

    print(f"🎯 Task: {task}\n")

    # System Prompts
    prompt_a = (
        "You are Agent A (Planner). You determine distances. "
        "Ask Agent B for cost calculations once you have the distance. "
        "When you have the final budget, say 'CONCLUSION: <summary>' to end."
    )

    prompt_b = (
        "You are Agent B (Accountant). You calculate costs using your tool. "
        "Provide the numbers to Agent A."
    )

    # Loop for a maximum of 6 turns
    for i in range(6):
        # --- Agent A's Turn ---
        response_a = run_agent_turn("Agent A", prompt_a, history, tools_a)
        print(f"🤖 Agent A: {response_a}")
        history.append({"role": "assistant", "content": f"Agent A says: {response_a}"})

        # Check for conclusion
        if "CONCLUSION:" in response_a:
            print("\n✅ Conversation Finished.")
            break

        # --- Agent B's Turn ---
        response_b = run_agent_turn("Agent B", prompt_b, history, tools_b)
        print(f"🤖 Agent B: {response_b}")
        history.append({"role": "user", "content": f"Agent B says: {response_b}"})

if __name__ == "__main__":
    main()

Explanation of Key Concepts

  1. Shared History: Both agents see the same history list.

    • When Agent A speaks, we add it as an assistant message.

    • When Agent B speaks, we add it as a user message (from A's perspective, B is just another user providing info). This keeps the role distinction clear for the LLM.

  2. Tool Execution Flow: Notice the logic inside run_agent_turn:

    • Call 1: Model sees prompt -> Returns tool_calls.

    • Execution: Python runs the function (e.g., get_distance).

    • Call 2: Model sees prompt + tool call + tool result -> Returns "The distance is 300 miles."

  3. Termination: We instructed Agent A in the system prompt to use the keyword CONCLUSION: when the job is done. The Python loop checks for this string to break the loop