aixtools 0.0.0__py3-none-any.whl

aixtools/a2a/google_sdk/pydantic_ai_adapter/agent_executor.py

@@ -0,0 +1,199 @@
+from pathlib import Path
+
+from a2a.server.agent_execution import AgentExecutor, RequestContext
+from a2a.server.events import EventQueue
+from a2a.types import (
+    Artifact,
+    FilePart,
+    FileWithUri,
+    Message,
+    Part,
+    TaskArtifactUpdateEvent,
+    TaskState,
+    TaskStatus,
+    TaskStatusUpdateEvent,
+)
+from a2a.utils import get_file_parts, get_message_text, new_agent_text_message, new_task
+from pydantic import BaseModel
+from pydantic_ai import Agent, BinaryContent
+
+from aixtools.a2a.google_sdk.pydantic_ai_adapter.storage import InMemoryHistoryStorage
+from aixtools.a2a.google_sdk.remote_agent_connection import is_in_terminal_state
+from aixtools.a2a.google_sdk.utils import get_session_id_tuple
+from aixtools.agents import get_agent
+from aixtools.agents.prompt import build_user_input
+from aixtools.context import SessionIdTuple
+from aixtools.logging.logging_config import get_logger
+from aixtools.mcp.client import get_configured_mcp_servers
+
+logger = get_logger(__name__)
+
+
+class AgentParameters(BaseModel):
+    system_prompt: str
+    mcp_servers: list[str]
+
+
+class RunOutput(BaseModel):
+    is_task_failed: bool
+    is_task_in_progress: bool
+    is_input_required: bool
+    output: str
+    created_artifacts_paths: list[str]
+
+
+def _task_failed_event(text: str, context_id: str | None, task_id: str | None) -> TaskStatusUpdateEvent:
+    """Creates a TaskStatusUpdateEvent indicating task failure."""
+    return TaskStatusUpdateEvent(
+        status=TaskStatus(
+            state=TaskState.failed, message=new_agent_text_message(text=text, context_id=context_id, task_id=task_id)
+        ),
+        final=True,
+        context_id=context_id,
+        task_id=task_id,
+    )
+
+
+class PydanticAgentExecutor(AgentExecutor):
+    def __init__(self, agent_parameters: AgentParameters):
+        self._agent_parameters = agent_parameters
+        self.history_storage = InMemoryHistoryStorage()
+
+    def _convert_message_to_pydantic_parts(
+        self,
+        session_tuple: SessionIdTuple,
+        message: Message,
+    ) -> str | list[str | BinaryContent]:
+        """Convert A2A Message to a Pydantic AI agent input format"""
+        text_prompt = get_message_text(message)
+        file_parts = get_file_parts(message.parts)
+        if not file_parts:
+            return text_prompt
+        file_paths = [Path(part.uri) for part in file_parts if isinstance(part, FileWithUri)]
+
+        return build_user_input(session_tuple, text_prompt, file_paths)
+
+    async def execute(
+        self,
+        context: RequestContext,
+        event_queue: EventQueue,
+    ) -> None:
+        """
+        Execute the agent run.
+        Wraps pydantic ai agent execution with a2a protocol events
+        Args:
+            context (RequestContext): The request context containing the message and task information.
+            event_queue (EventQueue): The event queue to enqueue events.
+        """
+        session_tuple = get_session_id_tuple(context)
+        agent = self._build_agent(session_tuple)
+        if context.message is None:
+            raise ValueError("No message provided")
+
+        task = context.current_task
+        message = context.message
+        if not task:
+            task = new_task(message)
+            await event_queue.enqueue_event(task)
+
+        if is_in_terminal_state(task):
+            raise RuntimeError("Can not perform a task as it is in a terminal state: %s", task.status.state)
+
+        prompt = self._convert_message_to_pydantic_parts(session_tuple, message)
+        history_message = self.history_storage.get(task.id)
+
+        try:
+            result = await agent.run(
+                user_prompt=prompt,
+                message_history=history_message,
+            )
+        except Exception as e:
+            await event_queue.enqueue_event(
+                _task_failed_event(
+                    text=f"Agent execution error: {e}",
+                    context_id=context.context_id,
+                    task_id=task.id,
+                )
+            )
+            return
+
+        self.history_storage.store(task.id, result.all_messages())
+
+        run_output: RunOutput = result.output
+        if run_output.is_task_failed:
+            await event_queue.enqueue_event(
+                _task_failed_event(
+                    text=f"Task failed: {run_output.output}",
+                    context_id=context.context_id,
+                    task_id=task.id,
+                )
+            )
+            return
+
+        if run_output.is_input_required:
+            await event_queue.enqueue_event(
+                TaskStatusUpdateEvent(
+                    status=TaskStatus(
+                        state=TaskState.input_required,
+                        message=new_agent_text_message(
+                            text=run_output.output, context_id=context.context_id, task_id=task.id
+                        ),
+                    ),
+                    final=False,
+                    context_id=context.context_id,
+                    task_id=task.id,
+                )
+            )
+            return
+
+        if run_output.is_task_in_progress:
+            logger.error("Task hasn't been completed: %s", run_output.output)
+            await event_queue.enqueue_event(
+                _task_failed_event(
+                    text=f"Agent didn't manage complete the task: {run_output.output}",
+                    context_id=context.context_id,
+                    task_id=task.id,
+                )
+            )
+            return
+
+        for idx, artifact in enumerate(run_output.created_artifacts_paths):
+            image_file = FileWithUri(uri=str(artifact), name=f"image_{idx}")
+            await event_queue.enqueue_event(
+                TaskArtifactUpdateEvent(
+                    append=False,
+                    context_id=task.context_id,
+                    task_id=task.id,
+                    last_chunk=True,
+                    artifact=Artifact(
+                        artifact_id=f"image_{idx}",
+                        parts=[Part(root=FilePart(file=image_file))],
+                    ),
+                )
+            )
+        await event_queue.enqueue_event(
+            TaskStatusUpdateEvent(
+                status=TaskStatus(
+                    state=TaskState.completed,
+                    message=new_agent_text_message(
+                        text=run_output.output, context_id=context.context_id, task_id=task.id
+                    ),
+                ),
+                final=True,
+                context_id=context.context_id,
+                task_id=task.id,
+            )
+        )
+
+    async def cancel(self, ctx: RequestContext, event_queue: EventQueue) -> None:
+        """Cancel"""
+        raise Exception("cancel not supported")
+
+    def _build_agent(self, session_tuple: SessionIdTuple) -> Agent:
+        params = self._agent_parameters
+        mcp_servers = get_configured_mcp_servers(session_tuple, params.mcp_servers)
+        return get_agent(
+            system_prompt=params.system_prompt,
+            toolsets=mcp_servers,
+            output_type=RunOutput,
+        )

aixtools/a2a/google_sdk/pydantic_ai_adapter/storage.py

@@ -0,0 +1,26 @@
+"""Storage interface and in-memory implementation for Pydantic AI agent history."""
+
+from abc import ABC, abstractmethod
+
+from pydantic_ai.messages import ModelRequest, ModelResponse
+
+
+class PydanticAiAgentHistoryStorage(ABC):
+    @abstractmethod
+    def get(self, task_id: str) -> list[ModelRequest | ModelResponse] | None:
+        pass
+
+    @abstractmethod
+    def store(self, task_id: str, messages: list[ModelRequest | ModelResponse]) -> None:
+        pass
+
+
+class InMemoryHistoryStorage(PydanticAiAgentHistoryStorage):
+    def __init__(self):
+        self.storage: dict[str, list[ModelRequest | ModelResponse]] = {}
+
+    def get(self, task_id: str) -> list[ModelRequest | ModelResponse] | None:
+        return self.storage.get(task_id, None)
+
+    def store(self, task_id: str, messages: list[ModelRequest | ModelResponse]) -> None:
+        self.storage[task_id] = messages

aixtools/a2a/google_sdk/remote_agent_connection.py

@@ -0,0 +1,88 @@
+import asyncio
+
+from a2a.client import Client
+from a2a.types import (
+    AgentCard,
+    Message,
+    Task,
+    TaskQueryParams,
+    TaskState,
+)
+
+from aixtools.logging.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+def is_in_terminal_state(task: Task) -> bool:
+    return task.status.state in [
+        TaskState.completed,
+        TaskState.canceled,
+        TaskState.failed,
+    ]
+
+
+def is_in_terminal_or_interrupted_state(task: Task) -> bool:
+    return task.status.state in [
+        TaskState.input_required,
+        TaskState.unknown,
+    ] or is_in_terminal_state(task)
+
+
+class RemoteAgentConnection:
+    def __init__(self, card: AgentCard, client: Client):
+        self._client = client
+        self._card = card
+
+    def get_agent_card(self) -> AgentCard:
+        """
+        Returns the agent card associated with this connection.
+        """
+        return self._card
+
+    async def send_message(self, message: Message) -> Task | Message | None:
+        """
+        Sends a message to the remote agent and returns either a Task, a Message, or None.
+        """
+        last_task: Task | None = None
+        try:
+            async for event in self._client.send_message(message):
+                if isinstance(event, Message):
+                    return event
+                if is_in_terminal_or_interrupted_state(event[0]):
+                    return event[0]
+                last_task = event[0]
+        except Exception as e:
+            logger.error("Exception found in send_message: %s", str(e))
+            raise e
+        return last_task
+
+    async def send_message_with_polling(
+        self,
+        message: Message,
+        *,
+        sleep_time: float = 0.2,
+        max_iter=1000,
+    ) -> Task | Message:
+        """
+        Sends a message to the remote agent and polls for the task status at regular intervals.
+        If the task reaches a terminal state or is interrupted, it returns the task.
+        If the task does not complete within the maximum number of iterations, it raises an exception.
+        """
+        last_task = await self.send_message(message)
+        if not last_task:
+            raise ValueError("No task or message returned from send_message")
+        if isinstance(last_task, Message):
+            return last_task
+
+        if is_in_terminal_or_interrupted_state(last_task):
+            return last_task
+        task_id = last_task.id
+        for _ in range(max_iter):
+            await asyncio.sleep(sleep_time)
+            task = await self._client.get_task(TaskQueryParams(id=task_id))
+            if is_in_terminal_or_interrupted_state(task):
+                return task
+
+        timeout_seconds = max_iter * sleep_time
+        raise Exception(f"Task did not complete in {timeout_seconds} seconds")  # pylint: disable=broad-exception-raised

aixtools/a2a/google_sdk/utils.py

@@ -0,0 +1,59 @@
+import asyncio
+
+import httpx
+from a2a.client import A2ACardResolver, ClientConfig, ClientFactory
+from a2a.server.agent_execution import RequestContext
+from a2a.types import AgentCard
+from a2a.utils import AGENT_CARD_WELL_KNOWN_PATH, PREV_AGENT_CARD_WELL_KNOWN_PATH
+
+from aixtools.a2a.google_sdk.remote_agent_connection import RemoteAgentConnection
+from aixtools.context import DEFAULT_SESSION_ID, DEFAULT_USER_ID, SessionIdTuple
+
+
+class AgentCardLoadFailedError(Exception):
+    pass
+
+
+class _AgentCardResolver:
+    def __init__(self, client: httpx.AsyncClient):
+        self._httpx_client = client
+        self._a2a_client_factory = ClientFactory(ClientConfig(httpx_client=self._httpx_client))
+        self.clients: dict[str, RemoteAgentConnection] = {}
+
+    def register_agent_card(self, card: AgentCard):
+        remote_connection = RemoteAgentConnection(card, self._a2a_client_factory.create(card))
+        self.clients[card.name] = remote_connection
+
+    async def retrieve_card(self, address: str):
+        for card_path in [AGENT_CARD_WELL_KNOWN_PATH, PREV_AGENT_CARD_WELL_KNOWN_PATH]:
+            try:
+                card_resolver = A2ACardResolver(self._httpx_client, address, card_path)
+                card = await card_resolver.get_agent_card()
+                card.url = address
+                self.register_agent_card(card)
+                return
+            except Exception as e:
+                print(f"Error retrieving agent card from {address} at path {card_path}: {e}")
+
+        raise AgentCardLoadFailedError(f"Failed to load agent card from {address}")
+
+    async def get_a2a_clients(self, agent_hosts: list[str]) -> dict[str, RemoteAgentConnection]:
+        async with asyncio.TaskGroup() as task_group:
+            for address in agent_hosts:
+                task_group.create_task(self.retrieve_card(address))
+
+        return self.clients
+
+
+async def get_a2a_clients(ctx: SessionIdTuple, agent_hosts: list[str]) -> dict[str, RemoteAgentConnection]:
+    headers = {
+        "user-id": ctx[0],
+        "session-id": ctx[1],
+    }
+    httpx_client = httpx.AsyncClient(headers=headers, timeout=60.0)
+    return await _AgentCardResolver(httpx_client).get_a2a_clients(agent_hosts)
+
+
+def get_session_id_tuple(context: RequestContext) -> SessionIdTuple:
+    headers = context.call_context.state.get("headers", {})
+    return headers.get("user-id", DEFAULT_USER_ID), headers.get("session-id", DEFAULT_SESSION_ID)

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