aixtools 0.1.0__py3-none-any.whl

aixtools/__init__.py

@@ -0,0 +1,5 @@
+"""
+AiXtools - Tools for AI exploration and debugging
+"""
+
+__version__ = "0.1.0"

aixtools/a2a/__init__.py

@@ -0,0 +1,5 @@
+"""A2A (Agent-to-Agent) communication utilities."""
+
+from .utils import fetch_agent_card, task, text_message, tool2skill
+
+__all__ = ["fetch_agent_card", "task", "text_message", "tool2skill"]

aixtools/a2a/app.py

@@ -0,0 +1,126 @@
+"""
+This module provides functionality to convert a Pydantic AI Agent into a FastA2A application
+"""
+
+import json
+from functools import partial
+from typing import assert_never
+
+from fasta2a.applications import FastA2A
+from fasta2a.broker import InMemoryBroker
+from fasta2a.schema import Part, TaskSendParams
+from fasta2a.storage import InMemoryStorage
+from pydantic_ai import Agent
+from pydantic_ai._a2a import AgentWorker, worker_lifespan
+from pydantic_ai.messages import (
+    AudioUrl,
+    BinaryContent,
+    DocumentUrl,
+    ImageUrl,
+    ModelRequestPart,
+    UserPromptPart,
+    VideoUrl,
+)
+from starlette.applications import Starlette
+from starlette.exceptions import HTTPException
+from starlette.requests import Request
+from starlette.responses import RedirectResponse
+
+from aixtools.context import session_id_var, user_id_var
+
+
+class AgentWorkerWithMetadataParser(AgentWorker):
+    """Custom AgentWorker class that extracts the session metadata from message metadata."""
+
+    async def run_task(self, params: TaskSendParams) -> None:
+        """
+        Extract session metadata from message and store them in context variables,
+        then call the parent class's run_task method.
+        """
+        # Load the task to extract metadata
+        task = await self.storage.load_task(params["id"])
+        if task:
+            # Extract headers from message metadata if available
+            if message := (task.get("history") or [None])[-1]:
+                metadata = message.get("metadata", {})
+                # Store in context variables
+                user_id_var.set(metadata.get("user_id", ""))
+                session_id_var.set(metadata.get("session_id", ""))
+        # Call the parent class's run_task method
+        return await super().run_task(params)
+
+
+class AgentWorkerWithDataPartSupport(AgentWorkerWithMetadataParser):
+    """Custom agent worker that adds support for data parts in messages."""
+
+    def _request_parts_from_a2a(self, parts: list[Part]) -> list[ModelRequestPart]:
+        """
+        Clones underlying method with additional support for data parts.
+        TODO: remove once pydantic-ai supports data parts natively.
+        """
+        model_parts: list[ModelRequestPart] = []
+        for part in parts:
+            if part["kind"] == "text":
+                model_parts.append(UserPromptPart(content=part["text"]))
+            elif part["kind"] == "file":
+                file_content = part["file"]
+                if "bytes" in file_content:
+                    data = file_content["bytes"].encode("utf-8")
+                    mime_type = file_content.get("mime_type", "application/octet-stream")
+                    content = BinaryContent(data=data, media_type=mime_type)
+                    model_parts.append(UserPromptPart(content=[content]))
+                else:
+                    url = file_content["uri"]
+                    for url_cls in (DocumentUrl, AudioUrl, ImageUrl, VideoUrl):
+                        content = url_cls(url=url)
+                        try:
+                            content.media_type
+                        except ValueError:  # pragma: no cover
+                            continue
+                        else:
+                            break
+                    else:
+                        raise ValueError(f"Unsupported file type: {url}")  # pragma: no cover
+                    model_parts.append(UserPromptPart(content=[content]))
+            elif part["kind"] == "data":
+                content = json.dumps(part["data"])
+                model_parts.append(UserPromptPart(content=[content]))
+            else:
+                assert_never(part)
+        return model_parts
+
+
+def agent_to_a2a(
+    agent: Agent, name: str, description: str, skills: list[dict], worker_class=AgentWorkerWithMetadataParser
+) -> FastA2A:
+    """Convert the agent to an A2A application taking care of session metadata extraction."""
+    storage = InMemoryStorage()
+    broker = InMemoryBroker()
+    worker = worker_class(broker=broker, storage=storage, agent=agent)
+    return FastA2A(
+        storage=storage,
+        broker=broker,
+        name=name,
+        description=description,
+        skills=skills,
+        url=None,
+        lifespan=partial(worker_lifespan, worker=worker, agent=agent),
+    )
+
+
+def fix_a2a_docs_pages(app: Starlette) -> None:
+    """
+    Fix the FastA2A documentation to point to the correct path.
+    This is a workaround for the issue with the FastA2A docs not being served correctly
+    when mounted as a sub-path.
+    """
+
+    async def redirect_to_sub_agent(request: Request):
+        """Redirect to proper sub-app using the Referer header to determine the path prefix."""
+        referer = request.headers.get("referer", "")
+        if referer.endswith("/docs"):
+            return RedirectResponse(url=f"{referer.rsplit('/', 1)[0]}{request.url.path}")
+        raise HTTPException(status_code=404)
+
+    app.router.add_route("/.well-known/agent.json", redirect_to_sub_agent, methods=["GET"])
+    app.router.add_route("/", redirect_to_sub_agent, methods=["POST"])

aixtools/a2a/utils.py

@@ -0,0 +1,115 @@
+"""Utilities for Agent-to-Agent (A2A) communication and task management."""
+
+import asyncio
+import uuid
+from typing import Callable
+
+from fasta2a import Skill
+from fasta2a.client import A2AClient
+from fasta2a.schema import GetTaskResponse, Message, Part, TextPart
+from fastapi import status
+
+from ..server import get_session_id_tuple
+
+SLEEP_TIME = 0.2
+MAX_ITER = 1000
+HTTP_OK = 200
+
+
+def card2description(card):
+    """Convert agent card to a description string."""
+    descr = f"{card['name']}: {card['description']}\n"
+    skills = card.get("skills", [])
+    for skill in skills:
+        descr += f"\t - {skill['name']}: {skill['description']}\n"
+    return descr
+
+
+async def fetch_agent_card(client: A2AClient) -> dict:
+    """Request the Agent's card"""
+    server_url = str(client.http_client.base_url).rstrip("/")
+    agent_card_url = f"{server_url}/.well-known/agent.json"
+    response = await client.http_client.get(agent_card_url, timeout=10)
+    if response.status_code == status.HTTP_200_OK:
+        card_data = response.json()
+        return card_data
+    raise Exception(f"Failed to retrieve agent card from {agent_card_url}. Status code: {response.status_code}")  # pylint: disable=broad-exception-raised
+
+
+def get_result_text(ret: GetTaskResponse) -> str | None:
+    """Extract the result text from the task result"""
+    if "result" not in ret:
+        return None
+    result = ret["result"]
+    if "artifacts" not in result:
+        return None
+    artifacts = result["artifacts"]
+    for artifact in artifacts:
+        if "parts" not in artifact:
+            continue
+        parts = artifact["parts"]
+        for part in parts:
+            if part["kind"] == "text":
+                return part["text"]
+    return None
+
+
+async def poll_task(client: A2AClient, task_id: str) -> GetTaskResponse:
+    """Polls the task status until it is completed or failed."""
+    state = None
+    for _ in range(MAX_ITER):
+        ret = await client.get_task(task_id=task_id)
+        # Check the state of the task
+        state = ret["result"]["status"]["state"] if "result" in ret and "status" in ret["result"] else None
+        if state == "completed":
+            return ret
+        if state == "failed":
+            raise Exception("Task failed")  # pylint: disable=broad-exception-raised
+        # Sleep for a while before checking again
+        await asyncio.sleep(SLEEP_TIME)
+    timeout_seconds = MAX_ITER * SLEEP_TIME
+    raise Exception(f"Task did not complete in {timeout_seconds} seconds")  # pylint: disable=broad-exception-raised
+
+
+async def submit_task(client: A2AClient, message: Message) -> str:
+    """Send a message to the client and return task id."""
+    user_id, session_id = get_session_id_tuple()
+    msg = message.copy()
+    msg["metadata"] = {
+        **msg.get("metadata", {}),
+        "user_id": client.http_client.headers.get("user-id", user_id),
+        "session_id": client.http_client.headers.get("session-id", session_id),
+    }
+    ret = await client.send_message(message=msg)
+    task_id = ret["result"]["id"] if "result" in ret and "id" in ret["result"] else ""
+    return task_id
+
+
+def multipart_message(parts: list[Part]) -> Message:
+    """Create a message object"""
+    message = Message(kind="message", role="user", parts=parts, message_id=str(uuid.uuid4()))
+    return message
+
+
+def text_message(text: str) -> Message:
+    """Create a message object with a text part."""
+    text_part = TextPart(kind="text", text=text, metadata={})
+    return multipart_message([text_part])
+
+
+async def task(client: A2AClient, text: str) -> GetTaskResponse:
+    """Send a text message to the client and wait for task completion."""
+    msg = text_message(text)
+    task_id = await submit_task(client, msg)
+    print(f"Task ID: {task_id}")
+    ret = await poll_task(client, task_id)
+    return ret
+
+
+def tool2skill(tool: Callable) -> Skill:
+    """Convert a tool to a skill."""
+    return Skill(
+        id=tool.__name__,
+        name=tool.__name__,
+        description=tool.__doc__ or "",
+    )  # type: ignore

aixtools/agents/__init__.py

@@ -0,0 +1,12 @@
+"""Agent utilities for running and managing AI agents."""
+
+from .agent import get_agent, get_model, run_agent
+from .agent_batch import AgentQueryParams, run_agent_batch
+
+__all__ = [
+    "get_agent",
+    "get_model",
+    "run_agent",
+    "AgentQueryParams",
+    "run_agent_batch",
+]

aixtools/agents/agent.py

@@ -0,0 +1,164 @@
+"""
+Core agent implementation providing model selection and configuration for AI agents.
+"""
+
+from types import NoneType
+from typing import Any
+
+from openai import AsyncAzureOpenAI
+from pydantic_ai import Agent
+from pydantic_ai.models.bedrock import BedrockConverseModel
+from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.providers.bedrock import BedrockProvider
+from pydantic_ai.providers.openai import OpenAIProvider
+from pydantic_ai.settings import ModelSettings
+from pydantic_ai.usage import UsageLimits
+
+from aixtools.logging.log_objects import ObjectLogger
+from aixtools.logging.logging_config import get_logger
+from aixtools.logging.model_patch_logging import model_patch_logging
+from aixtools.utils.config import (
+    AWS_PROFILE,
+    AWS_REGION,
+    AZURE_MODEL_NAME,
+    AZURE_OPENAI_API_KEY,
+    AZURE_OPENAI_API_VERSION,
+    AZURE_OPENAI_ENDPOINT,
+    BEDROCK_MODEL_NAME,
+    MODEL_FAMILY,
+    MODEL_TIMEOUT,
+    OLLAMA_MODEL_NAME,
+    OLLAMA_URL,
+    OPENAI_API_KEY,
+    OPENAI_MODEL_NAME,
+    OPENROUTER_API_KEY,
+    OPENROUTER_API_URL,
+    OPENROUTER_MODEL_NAME,
+)
+
+logger = get_logger(__name__)
+
+
+def _get_model_bedrock(model_name=BEDROCK_MODEL_NAME, aws_region=AWS_REGION):
+    assert model_name, "BEDROCK_MODEL_NAME is not set"
+    assert aws_region, "AWS_REGION is not set"
+
+    if AWS_PROFILE is not None:
+        return BedrockConverseModel(model_name=model_name)
+
+    provider = BedrockProvider(region_name=aws_region)
+    return BedrockConverseModel(model_name=model_name, provider=provider)
+
+
+def _get_model_ollama(model_name=OLLAMA_MODEL_NAME, ollama_url=OLLAMA_URL):
+    assert ollama_url, "OLLAMA_URL is not set"
+    assert model_name, "Model name is not set"
+    provider = OpenAIProvider(base_url=ollama_url)
+    return OpenAIModel(model_name=model_name, provider=provider)
+
+
+def _get_model_openai(model_name=OPENAI_MODEL_NAME, openai_api_key=OPENAI_API_KEY):
+    assert openai_api_key, "OPENAI_API_KEY is not set"
+    assert model_name, "Model name is not set"
+    provider = OpenAIProvider(api_key=openai_api_key)
+    return OpenAIModel(model_name=model_name, provider=provider)
+
+
+def _get_model_openai_azure(
+    model_name=AZURE_MODEL_NAME,
+    azure_openai_api_key=AZURE_OPENAI_API_KEY,
+    azure_openai_endpoint=AZURE_OPENAI_ENDPOINT,
+    azure_openai_api_version=AZURE_OPENAI_API_VERSION,
+):
+    assert azure_openai_endpoint, "AZURE_OPENAI_ENDPOINT is not set"
+    assert azure_openai_api_key, "AZURE_OPENAI_API_KEY is not set"
+    assert azure_openai_api_version, "AZURE_OPENAI_API_VERSION is not set"
+    assert model_name, "Model name is not set"
+    client = AsyncAzureOpenAI(
+        azure_endpoint=azure_openai_endpoint, api_version=azure_openai_api_version, api_key=azure_openai_api_key
+    )
+    return OpenAIModel(model_name=model_name, provider=OpenAIProvider(openai_client=client))
+
+
+def _get_model_open_router(
+    model_name=OPENROUTER_MODEL_NAME, openrouter_api_url=OPENROUTER_API_URL, openrouter_api_key=OPENROUTER_API_KEY
+):
+    assert openrouter_api_url, "OPENROUTER_API_URL is not set"
+    assert openrouter_api_key, "OPENROUTER_API_KEY is not set"
+    assert model_name, "Model name is not set, missing 'OPENROUTER_MODEL_NAME' environment variable?"
+    provider = OpenAIProvider(base_url=openrouter_api_url, api_key=openrouter_api_key)
+    return OpenAIModel(model_name, provider=provider)
+
+
+def get_model(model_family=MODEL_FAMILY, model_name=None, **kwargs):
+    """Create and return appropriate model instance based on specified family and name."""
+    assert model_family is not None and model_family != "", f"Model family '{model_family}' is not set"
+    match model_family:
+        case "azure":
+            return _get_model_openai_azure(model_name=model_name or AZURE_MODEL_NAME, **kwargs)
+        case "bedrock":
+            return _get_model_bedrock(model_name=model_name or BEDROCK_MODEL_NAME, **kwargs)
+        case "ollama":
+            return _get_model_ollama(model_name=model_name or OLLAMA_MODEL_NAME, **kwargs)
+        case "openai":
+            return _get_model_openai(model_name=model_name or OPENAI_MODEL_NAME, **kwargs)
+        case "openrouter":
+            return _get_model_open_router(model_name=model_name or OPENROUTER_MODEL_NAME, **kwargs)
+        case _:
+            raise ValueError(f"Model family '{model_family}' not supported")
+
+
+def get_agent(  # noqa: PLR0913, pylint: disable=too-many-arguments,too-many-positional-arguments
+    model=None,
+    *,
+    instructions=None,
+    system_prompt=(),
+    tools=(),
+    toolsets=(),
+    model_settings=None,
+    output_type: Any = str,
+    deps_type=NoneType,
+) -> Agent:
+    """Get a PydanticAI agent"""
+    if model_settings is None:
+        model_settings = ModelSettings(timeout=MODEL_TIMEOUT)
+    if model is None:
+        model = get_model()
+    agent = Agent(
+        model=model,
+        output_type=output_type,
+        instructions=instructions,
+        system_prompt=system_prompt,
+        deps_type=deps_type,
+        model_settings=model_settings,
+        tools=tools,
+        toolsets=toolsets,
+        instrument=True,
+    )
+    return agent
+
+
+async def run_agent(  # noqa: PLR0913, pylint: disable=too-many-arguments,too-many-positional-arguments
+    agent: Agent,
+    prompt: str | list[str],
+    usage_limits: UsageLimits | None = None,
+    verbose: bool = False,
+    debug: bool = False,
+    log_model_requests: bool = False,
+    parent_logger: ObjectLogger | None = None,
+):
+    """Query the LLM"""
+    # Results
+    nodes, result = [], None
+    async with agent.iter(prompt, usage_limits=usage_limits) as agent_run:
+        # Create a new log file for each run
+        with ObjectLogger(parent_logger=parent_logger, verbose=verbose, debug=debug) as agent_logger:
+            # Patch the model with the logger
+            if log_model_requests:
+                agent.model = model_patch_logging(agent.model, agent_logger)
+            # Run the agent
+            async for node in agent_run:
+                agent_logger.log(node)
+                nodes.append(node)
+            result = agent_run.result
+    return result.output if result else None, nodes

aixtools/agents/agent_batch.py

@@ -0,0 +1,74 @@
+"""
+Batch processing functionality for running multiple agent queries in parallel.
+"""
+
+import asyncio
+from typing import Any
+
+from pydantic import BaseModel
+
+from aixtools.agents.agent import get_agent, run_agent
+
+
+class AgentQueryParams(BaseModel):
+    """Parameters for configuring agent queries in batch processing."""
+
+    class Config:  # pylint: disable=too-few-public-methods
+        """Configuration for the model."""
+
+        arbitrary_types_allowed = True
+
+    id: str = ""  # Unique identifier for the query
+    prompt: str | list[str]
+    agent: Any = None
+    model: Any = None
+    debug: bool = False
+    output_type: Any = str
+    tools: list | None = []
+
+    async def run(self):
+        """Query the LLM"""
+        agent = self.agent
+        if agent is None:
+            agent = get_agent(
+                system_prompt=self.prompt, model=self.model, tools=self.tools, output_type=self.output_type
+            )
+        return await run_agent(agent=agent, prompt=self.prompt, debug=self.debug)
+
+
+async def run_agent_batch(query_parameters: list[AgentQueryParams], batch_size=10):
+    """
+    Run multiple queries simultanously in batches of at most batch_size
+    and yield the results as they come in.
+
+        Usage example:
+        query_parameters = [
+            AgentQueryParams(prompt="What is the meaning of life")
+            AgentQueryParams(prompt="Who is the prime minister of Canada")
+        ]
+
+        async for result in agent_batch(query_parameters):
+            print(result)
+    """
+    tasks = []
+    batch_num, total = 1, len(query_parameters)
+    for i, qp in enumerate(query_parameters):
+        tasks.append(qp.run())
+        if len(tasks) >= batch_size:
+            # Run a batch of tasks
+            print(f"Running batch {batch_num}, {i + 1} / {total}")
+            tasks_results = await asyncio.gather(
+                *tasks
+            )  # Returns a list of results, each one is a tuple (result, nodes)
+            # Yield the results
+            for r, _ in tasks_results:
+                yield r
+            tasks = []
+            batch_num += 1
+    # Run the last batch of tasks
+    if tasks:
+        print(f"Running final batch {batch_num}")
+        tasks_results = await asyncio.gather(*tasks)
+        for r, _ in tasks_results:
+            yield r
+    print("Done")

aixtools/app.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+
+"""
+Simple Chainlit app example
+"""
+
+import traceback
+
+import chainlit as cl
+from pydantic_graph import End
+
+from aixtools.agents.agent import get_agent
+from aixtools.logging.logging_config import get_logger
+from aixtools.utils.chainlit import cl_agent_show
+
+logger = get_logger(__name__)
+
+HISTORY = "history"
+
+SYSTEM_PROMPT = """
+You are a helpful assistant.
+"""
+
+
+@cl.step
+async def greet_tool(msg: str) -> str:
+    """A simple greeting tool"""
+    return f"Hello! You said: {msg}"
+
+
+async def parse_user_message(message):
+    """Parse user message and check if it is a command"""
+    # When we type something that starts with ':', we are using a "command" (i.e. it does not go to the agent)
+    command = str(message.content).strip().lower()
+    if command.startswith(":"):
+        logger.debug("Received command: %s", command)
+        match command:
+            case ":clear":
+                # Clear the history
+                cl.user_session.set(HISTORY, [])
+                return None
+            case ":help":
+                # Show help
+                help_message = """
+                Available commands:
+                - :clear: Clear the chat history
+                - :help: Show this help message
+                """
+                await cl.Message(content=help_message).send()
+                return None
+            case ":history":
+                # Show history
+                history = cl.user_session.get(HISTORY)
+                if history:
+                    history_message = "\n".join(history)
+                    await cl.Message(content=f"Chat history:\n{history_message}").send()
+                else:
+                    await cl.Message(content="No history available.").send()
+                return None
+            case _:
+                # Unknown command
+                await cl.Message(content=f"Unknown command: {command}").send()
+                return None
+    else:
+        user_message = message.content
+        logger.debug("User message: %s", user_message)
+    return user_message
+
+
+async def run_agent(messages):
+    """Run the agent with the given messages"""
+    agent = get_agent(system_prompt=SYSTEM_PROMPT, tools=[greet_tool])
+    ret = ""
+    msg = cl.Message(content="")
+    await msg.send()
+    try:
+        ret = await cl_agent_show.show_run(agent=agent, prompt=messages, msg=msg, debug=False)
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        msg.elements.append(cl.Text(name="Error", content=f"Error: {e}", type="error"))  # pylint: disable=unexpected-keyword-arg
+        logger.error("Error: %s", e)
+        # Log the full stack trace for debugging
+        stack_trace = traceback.format_exc()
+        logger.error("Stack tarace:\n%s", stack_trace)
+        logger.error("Stack trace:\n%s", stack_trace)
+        msg.elements.append(cl.Text(name="Stack Trace", content=stack_trace, language="python"))
+        ret = f"Internal server error: {e}"
+    await msg.send()
+    return ret
+
+
+def update_history(history, user_message=None, run_return=None):
+    """Update history with user message and model run output"""
+    assert user_message is not None or run_return is not None, "Either user message or run return must be provided"
+    if user_message is not None:
+        logger.debug("Updating history: Got user message type %s: %s", type(user_message), user_message)
+        assert isinstance(user_message, str)
+        history.append(user_message)
+    if run_return is not None:
+        logger.debug("Updating history: Got agent output type %s: %s", type(run_return), run_return)
+        latest_item = ""
+        if isinstance(run_return, list):
+            # If it is a list of 'node' items, the last element is the 'end_message' with the final result
+            end_message: End = run_return[-1]
+            final_result = end_message.data
+            latest_item = str(final_result.data)
+        else:
+            latest_item = str(run_return)
+        # Update history and store it
+        logger.debug("Updating history: Adding to history type %s: %s", type(latest_item), latest_item)
+        history.append(latest_item)
+    return history
+
+
+@cl.set_starters
+async def set_starters():
+    """Set the starters"""
+    return [
+        cl.Starter(label="Message", message="Hello world!"),
+    ]
+
+
+@cl.on_chat_start
+async def on_chat_start():
+    """Initialize chat session by resetting history when a new chat starts."""
+    # Reset history
+    logger.debug("On chat start")
+    cl.user_session.set(HISTORY, [])
+
+
+@cl.on_message
+async def on_message(message: cl.Message):
+    """Process incoming chat messages and generate responses using the agent."""
+    history = cl.user_session.get(HISTORY)  # Get user message and history
+    user_message = await parse_user_message(message)  # Parse user message
+    # Check if user message is None (e.g. if it is a command)
+    if user_message is None:
+        return
+    messages = update_history(history, user_message=user_message)  # Update history with user message
+    # Run the agent
+    run_return = await run_agent(messages)
+    # Update history and store it
+    history = update_history(history, run_return=run_return)
+    cl.user_session.set(HISTORY, messages)

aixtools/context.py

@@ -0,0 +1,12 @@
+"""
+This module defines global context variables for request-specific information
+that can be used for logging, tracing, and other purposes across applications
+that use aixtools.
+"""
+
+from contextvars import ContextVar
+
+# Define context variables with default values.
+# These can be populated by middleware or where they are initialized
+session_id_var: ContextVar[str | None] = ContextVar("session_id", default=None)
+user_id_var: ContextVar[str | None] = ContextVar("user_id", default=None)

aixtools/db/__init__.py

@@ -0,0 +1,17 @@
+"""
+Database module for vector storage and retrieval.
+"""
+
+from aixtools.db.database import DatabaseError, SqliteDb
+from aixtools.db.vector_db import get_vdb_embedding, get_vector_db, vdb_add, vdb_get_by_id, vdb_has_id, vdb_query
+
+__all__ = [
+    "DatabaseError",
+    "SqliteDb",
+    "get_vdb_embedding",
+    "get_vector_db",
+    "vdb_add",
+    "vdb_get_by_id",
+    "vdb_has_id",
+    "vdb_query",
+]