aixtools 0.1.4__tar.gz → 0.1.5__tar.gz

aixtools-0.1.4/.github/workflows/lint.yml → aixtools-0.1.5/.github/workflows/lint-and-test.yml

@@ -35,3 +35,7 @@ jobs:
       - name: Run linting
         run: |
           ./scripts/lint.sh
+
+      - name: Test
+        run: |
+          ./scripts/test.sh

{aixtools-0.1.4 → aixtools-0.1.5}/PKG-INFO

@@ -1,9 +1,11 @@
 Metadata-Version: 2.4
 Name: aixtools
-Version: 0.1.4
+Version: 0.1.5
 Summary: Tools for AI exploration and debugging
 Requires-Python: >=3.11.2
 Description-Content-Type: text/markdown
+Requires-Dist: a2a-sdk>=0.3.1
+Requires-Dist: cachebox>=5.0.1
 Requires-Dist: chainlit>=2.5.5
 Requires-Dist: colorlog>=6.9.0
 Requires-Dist: fasta2a>=0.5.0

{aixtools-0.1.4 → aixtools-0.1.5}/aixtools/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.4'
-__version_tuple__ = version_tuple = (0, 1, 4)
+__version__ = version = '0.1.5'
+__version_tuple__ = version_tuple = (0, 1, 5)
 
-__commit_id__ = commit_id = 'g2ba64c16e'
+__commit_id__ = commit_id = 'g6b31b6684'

{aixtools-0.1.4 → aixtools-0.1.5}/aixtools/a2a/app.py

@@ -103,7 +103,7 @@ def agent_to_a2a(
         name=name,
         description=description,
         skills=skills,
-        url=None,
+        url="",
         lifespan=partial(worker_lifespan, worker=worker, agent=agent),
     )
 

aixtools-0.1.5/aixtools/a2a/google_sdk/card.py

@@ -0,0 +1,27 @@
+import httpx
+from a2a.client import A2ACardResolver
+from a2a.types import AgentCard
+
+from aixtools.logging.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+async def get_agent_card(httpx_client: httpx.AsyncClient, agent_url: str) -> AgentCard:
+    resolver = A2ACardResolver(
+        httpx_client=httpx_client,
+        base_url=agent_url,
+    )
+
+    try:
+        _public_card = await resolver.get_agent_card()  # Fetches from default public path
+        logger.info("Successfully fetched public agent card:")
+        logger.info(_public_card.model_dump_json(indent=2, exclude_none=True))
+        final_agent_card_to_use = _public_card
+    except Exception as e:
+        logger.error(f"Critical error fetching public agent card: {e}", exc_info=True)
+        raise RuntimeError("Failed to fetch the public agent card. Cannot continue.") from e
+
+    # Set the URL which is accessible from the container
+    final_agent_card_to_use.url = agent_url
+    return final_agent_card_to_use

aixtools-0.1.5/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-0.1.5/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-0.1.5/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-0.1.5/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-0.1.5/aixtools/agents/prompt.py

@@ -0,0 +1,97 @@
+"""Prompt building utilities for Pydantic AI agent, including file handling and context management."""
+
+import mimetypes
+from pathlib import Path, PurePosixPath
+
+from pydantic_ai import BinaryContent
+
+from aixtools.context import SessionIdTuple
+from aixtools.server import container_to_host_path
+from aixtools.utils.files import is_text_content
+
+CLAUDE_MAX_FILE_SIZE_IN_CONTEXT = 4 * 1024 * 1024  # Claude limit 4.5 MB for PDF files
+CLAUDE_IMAGE_MAX_FILE_SIZE_IN_CONTEXT = (
+    5 * 1024 * 1024
+)  # Claude limit 5 MB for images, to avoid large image files in context
+
+
+def should_be_included_into_context(
+    file_content: BinaryContent | str | None,
+    file_size: int,
+    *,
+    max_img_size_bytes: int = CLAUDE_IMAGE_MAX_FILE_SIZE_IN_CONTEXT,
+    max_file_size_bytes: int = CLAUDE_MAX_FILE_SIZE_IN_CONTEXT,
+) -> bool:
+    """Decide whether a file content should be included into the model context based on its type and size."""
+    if not isinstance(file_content, BinaryContent):
+        return False
+
+    if file_content.media_type.startswith("text/"):
+        return False
+
+    # Exclude archive files as they're not supported by OpenAI models
+    archive_types = {
+        "application/zip",
+        "application/x-tar",
+        "application/gzip",
+        "application/x-gzip",
+        "application/x-rar-compressed",
+        "application/x-7z-compressed",
+    }
+    if file_content.media_type in archive_types:
+        return False
+
+    if file_content.is_image and file_size < max_img_size_bytes:
+        return True
+
+    return file_size < max_file_size_bytes
+
+
+def file_to_binary_content(file_path: str | Path, mime_type: str = "") -> str | BinaryContent:
+    """
+    Read a file and return its content as either a UTF-8 string (for text files)
+    or BinaryContent (for binary files).
+    """
+    with open(file_path, "rb") as f:
+        data = f.read()
+
+    if not mime_type:
+        mime_type, _ = mimetypes.guess_type(file_path)
+        mime_type = mime_type or "application/octet-stream"
+
+    if is_text_content(data, mime_type):
+        return data.decode("utf-8")
+
+    return BinaryContent(data=data, media_type=mime_type)
+
+
+def build_user_input(
+    session_tuple: SessionIdTuple,
+    user_text: str,
+    file_paths: list[Path],
+) -> str | list[str | BinaryContent]:
+    """Build user input for the Pydantic AI agent, including file attachments if provided."""
+    if not file_paths:
+        return user_text
+
+    attachment_info_lines = []
+    binary_attachments = []
+
+    for workspace_path in file_paths:
+        host_path = container_to_host_path(PurePosixPath(workspace_path), ctx=session_tuple)
+        file_size = host_path.stat().st_size
+        mime_type, _ = mimetypes.guess_type(host_path)
+        mime_type = mime_type or "application/octet-stream"
+
+        attachment_info = f"* {workspace_path.name} (file_size={file_size} bytes) (path in workspace: {workspace_path})"
+        binary_content = file_to_binary_content(host_path, mime_type)
+
+        if should_be_included_into_context(binary_content, file_size):
+            binary_attachments.append(binary_content)
+            attachment_info += f" -- provided to model context at index {len(binary_attachments) - 1}"
+
+        attachment_info_lines.append(attachment_info)
+
+    full_prompt = user_text + "\nAttachments:\n" + "\n".join(attachment_info_lines)
+
+    return [full_prompt] + binary_attachments

{aixtools-0.1.4 → aixtools-0.1.5}/aixtools/context.py

@@ -10,3 +10,8 @@ from contextvars import ContextVar
 # 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)
+
+DEFAULT_USER_ID = "default_user"
+DEFAULT_SESSION_ID = "default_session"
+
+SessionIdTuple = tuple[str, str]

aixtools-0.1.5/aixtools/google/client.py

@@ -0,0 +1,25 @@
+import os
+from pathlib import Path
+
+from google import genai
+
+from aixtools.logging.logging_config import get_logger
+from aixtools.utils.config import GOOGLE_CLOUD_LOCATION, GOOGLE_CLOUD_PROJECT, GOOGLE_GENAI_USE_VERTEXAI
+
+logger = get_logger(__name__)
+
+
+def get_genai_client(service_account_key_path: Path | None = None) -> genai.Client:
+    """Initialize and return a Google GenAI client using Vertex AI / Gemini Developer API."""
+    assert GOOGLE_CLOUD_PROJECT, "GOOGLE_CLOUD_PROJECT is not set"
+    assert GOOGLE_CLOUD_LOCATION, "GOOGLE_CLOUD_LOCATION is not set"
+    if service_account_key_path:
+        if not service_account_key_path.exists():
+            raise FileNotFoundError(f"Service account key file not found: {service_account_key_path}")
+        logger.info(f"✅ GCP Service Account Key File: {service_account_key_path}")
+        os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = str(service_account_key_path)
+    return genai.Client(
+        vertexai=GOOGLE_GENAI_USE_VERTEXAI,
+        project=GOOGLE_CLOUD_PROJECT,
+        location=GOOGLE_CLOUD_LOCATION,
+    )

{aixtools-0.1.4 → aixtools-0.1.5}/aixtools/logging/logging_config.py

@@ -6,6 +6,7 @@ import json
 import logging
 import logging.config
 import os
+import sys
 import time
 from pathlib import Path
 
@@ -72,7 +73,51 @@ def configure_logging():
     2. logging.yaml in the current working directory.
     3. logging.json in the current working directory.
     4. Hardcoded default configuration.
+
+    Special handling for pytest: If running under pytest without explicit
+    log flags, console logging is suppressed to avoid interfering with
+    pytest's own log capture mechanism.
     """
+    # Detect if running under pytest and suppress console logging unless explicitly requested
+    is_pytest = "pytest" in sys.modules or "pytest" in sys.argv[0] if sys.argv else False
+
+    # Check for live log flags - handle both separate and combined flags
+    wants_live_logs = False
+    if sys.argv:
+        for arg in sys.argv:
+            # Check for explicit --log-cli
+            if arg == "--log-cli":
+                wants_live_logs = True
+                break
+            # Check for -s either standalone or combined with other flags (like -vsk, -vs, etc.)
+            if arg.startswith("-") and not arg.startswith("--") and "s" in arg:
+                wants_live_logs = True
+                break
+
+    if is_pytest and not wants_live_logs:
+        # Use a minimal configuration that doesn't output to console during pytest
+        pytest_config = {
+            "version": 1,
+            "disable_existing_loggers": False,
+            "formatters": {
+                "simple": {
+                    "format": "%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+                }
+            },
+            "handlers": {
+                # Only use NullHandler to suppress console output but allow pytest log capture
+                "null": {
+                    "class": "logging.NullHandler",
+                }
+            },
+            "root": {
+                "handlers": ["null"],
+                "level": "INFO",
+            },
+        }
+        logging.config.dictConfig(pytest_config)
+        return
+
     config_path_str = os.environ.get("LOGGING_CONFIG_PATH")
 
     if config_path_str: