Skip to content

AI Agent

Install packages

1
2
3
4
!uv pip install -q \
    litellm==1.78.5 \
    python-dotenv==1.1.1 \
    pydantic==2.12.3

Import packages

1
2
3
4
5
6
7
8
import json
from typing import Any, Callable, Dict, List, Optional

import litellm  # type: ignore
from dotenv import load_dotenv  # type: ignore
from pydantic import BaseModel  # type: ignore

load_dotenv()

True

Define tool model

1
2
3
4
5
class Tool(BaseModel):
    name: str
    description: str
    parameters: Dict[str, Any]
    strict: bool

Define agent class

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
class AIAgent:
    def __init__(self, max_iterations: int = 10):
        self.messages: List[Dict[str, Any]] = []
        self.tools: List[Tool] = []
        self.max_iterations: int = max_iterations
        self.current_iteration: int = 0
        self.tool_map: Dict[str, Callable[..., str]] = {}

    def add_tool_definition(self, tool: Tool) -> None:
        self.tools.append(tool)

    def add_tool_function(self, name: str, func: Callable[..., str]) -> None:
        self.tool_map[name] = func

    def execute_tool_function(self, tool_name: str, **kwargs: Any) -> str:
        return self.tool_map[tool_name](**kwargs)

    def chat(self, user_input: str):
        self.messages.append({"role": "user", "content": user_input})

        tool_schemas = [
            {
                "type": "function",
                "function": {
                    "name": tool.name,
                    "description": tool.description,
                    "parameters": tool.parameters,
                },
            }
            for tool in self.tools
        ]

        self.current_iteration = 0

        while self.current_iteration < self.max_iterations:
            self.current_iteration += 1

            try:
                completion = litellm.completion(
                    model="gemini/gemini-2.0-flash",
                    messages=self.messages,
                    tools=tool_schemas if tool_schemas else None,
                )

                if not completion.choices:
                    raise Exception("Model returned an empty response.")

                choice = completion.choices[0].message
                tool_calls = getattr(choice, "tool_calls", None)

                if not tool_calls:
                    final_content = choice.content

                    if not final_content:
                        raise Exception(
                            "Agent did not return a response content."
                        )

                    self.messages.append(
                        {"role": "assistant", "content": final_content}
                    )

                    return final_content

                self.messages.append(
                    {
                        "role": "assistant",
                        "content": choice.content,
                        "tool_calls": [
                            {
                                "id": tool_call.id,
                                "type": "function",
                                "function": {
                                    "name": tool_call.function.name,
                                    "arguments": tool_call.function.arguments,
                                },
                            }
                            for tool_call in tool_calls
                        ],
                    }
                )

                for tool_call in tool_calls:
                    name = tool_call.function.name
                    kwargs = json.loads(tool_call.function.arguments)

                    result = self.execute_tool_function(name, **kwargs)
                    result_str = (
                        result if isinstance(result, str) else str(result)
                    )

                    self.messages.append(
                        {
                            "role": "tool",
                            "tool_call_id": tool_call.id,
                            "content": result_str,
                        }
                    )

            except Exception as error:
                print(f"Error: {error}")
                return f"An error occurred: {error}"

        return "Maximum iterations reached without completion."

Define greeting tool

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
greeting_tool = Tool(
    name="greeting",
    description="greet a user",
    parameters={
        "type": "object",
        "properties": {
            "name": {
                "type": "string",
                "description": "Name of the person to greet",
            },
        },
        "required": ["name"],
        "additionalProperties": False,
    },
    strict=True,
)


def greeting(name: str):
    return f"Hello, {name}"

Add tools to agent

1
2
3
4
agent = AIAgent()
agent.add_tool_definition(greeting_tool)
agent.add_tool_function("greeting", greeting)
agent.execute_tool_function("greeting", name="John")

'Hello, John'

Run agent loop

1
2
3
4
5
6
7
8
9
while True:
    user_input = input("🧑 You: ")
    if user_input.lower() in ["exit", "quit", "bye"]:
        print("👋 Goodbye!")
        break

    response = agent.chat(user_input)
    print(f"🧑 User: {user_input}")
    print(f"🤖 Agent: {response}")

🧑 User: Hello

🤖 Agent: I need your name to greet you. What is your name?



🧑 User: My name is John

🤖 Agent: Hello, John



👋 Goodbye!