python-base-agent 2026.2.13__py3-none-any.whl

base_agent/tools/tool_registry_client.py

@@ -0,0 +1,607 @@
+"""Tool registry client for registering agent tools with the platform.
+
+This is the Python equivalent of Java's ToolRegistryClient.
+Reference: /base-agent/src/main/java/one/ai/platform/baseagent/agent/tools/ToolRegistryClient.java
+"""
+
+import logging
+import threading
+import time
+
+import httpx
+
+from base_agent.tools.models import (
+    OutputSchemaConfig,
+    OutputSchemaResponse,
+    ToolRegistration,
+    ToolResponse,
+)
+from base_agent.utils.logger import sanitize_for_logging
+
+logger = logging.getLogger(__name__)
+
+
+class ToolRegistryClient:
+    """
+    Client for registering tools with the platform tool registry.
+
+    Provides feature parity with Java ToolRegistryClient:
+    - Hash-based schema registration
+    - Smart update logic (only updates when fields change)
+    - Thread-safe operations
+    - Health check integration via is_live property
+    - Automatic deactivation of registered tools on close()
+
+    Reference: Java ToolRegistryClient.java
+    """
+
+    API_PATH = "/api/v1/ai-agent-tools"
+
+    def __init__(self, admin_gateway_url: str, agent_id: int):
+        """
+        Initialize the tool registry client.
+
+        Args:
+            admin_gateway_url: Base URL of the admin gateway (e.g., "http://localhost:8888")
+            agent_id: ID of the agent registering tools
+        """
+        if not admin_gateway_url:
+            raise ValueError("admin_gateway_url is required")
+        if agent_id <= 0:
+            raise ValueError("agent_id must be positive")
+
+        self._admin_gateway_url = admin_gateway_url.rstrip("/")
+        self._agent_id = agent_id
+        self._registered_tool_ids: set[int] = set()
+        self._is_live = True
+        self._lock = threading.Lock()
+        self._client = httpx.Client(
+            base_url=self._admin_gateway_url,
+            timeout=30.0,
+        )
+
+        logger.info(
+            f"ToolRegistryClient initialized for agent {agent_id} "
+            f"at {sanitize_for_logging(admin_gateway_url)}"
+        )
+
+    @property
+    def is_live(self) -> bool:
+        """Check if the client is live (healthy)."""
+        return self._is_live
+
+    @property
+    def agent_id(self) -> int:
+        """Get the agent ID."""
+        return self._agent_id
+
+    @property
+    def registered_tool_ids(self) -> set[int]:
+        """Get the set of registered tool IDs."""
+        with self._lock:
+            return self._registered_tool_ids.copy()
+
+    def _request_with_retry(
+        self,
+        method: str,
+        path: str,
+        max_retries: int = 3,
+        **kwargs,
+    ) -> httpx.Response | None:
+        """
+        Make an HTTP request with retry logic.
+
+        Args:
+            method: HTTP method (GET, POST, PUT, PATCH, DELETE)
+            path: API path
+            max_retries: Maximum number of retries on connection errors
+            **kwargs: Additional arguments to pass to httpx
+
+        Returns:
+            Response object or None if all retries failed
+        """
+        for attempt in range(max_retries):
+            try:
+                response = self._client.request(method, path, **kwargs)
+                return response
+            except (httpx.ConnectError, httpx.TimeoutException) as e:
+                if attempt == max_retries - 1:
+                    self._is_live = False
+                    logger.error(
+                        f"Connection failed after {max_retries} attempts: "
+                        f"{sanitize_for_logging(str(e))}"
+                    )
+                    return None
+                wait_time = 0.5 * (attempt + 1)
+                logger.warning(
+                    f"Request failed (attempt {attempt + 1}/{max_retries}), "
+                    f"retrying in {wait_time}s: {sanitize_for_logging(str(e))}"
+                )
+                time.sleep(wait_time)
+        return None
+
+    def _get_existing_tool(self, tool_technical_name: str) -> ToolResponse | None:
+        """
+        Check if a tool exists for this agent.
+
+        Args:
+            tool_technical_name: Technical name of the tool
+
+        Returns:
+            ToolResponse if found, None otherwise
+        """
+        response = self._request_with_retry(
+            "GET",
+            self.API_PATH,
+            params={
+                "agent_id": self._agent_id,
+                "tool_technical_name": tool_technical_name,
+            },
+        )
+
+        if response is None:
+            return None
+
+        if response.status_code == 200:
+            data = response.json()
+            if data and len(data) > 0:
+                return ToolResponse(**data[0])
+        elif response.status_code == 404:
+            return None
+        else:
+            logger.warning(
+                f"Unexpected status {response.status_code} checking for tool: "
+                f"{sanitize_for_logging(response.text)}"
+            )
+
+        return None
+
+    def _needs_update(
+        self,
+        existing: ToolResponse,
+        name: str,
+        description: str | None,
+        ai_model_prompt_guidance: str | None,
+        input_schema_hash: str | None,
+        output_schemas: list[OutputSchemaConfig] | None,
+        metadata: dict | None,
+    ) -> bool:
+        """
+        Check if an existing tool needs to be updated.
+
+        Args:
+            existing: Existing tool from API
+            name: New name
+            description: New description
+            ai_model_prompt_guidance: New AI guidance
+            input_schema_hash: New input schema hash
+            output_schemas: New output schemas
+            metadata: New metadata
+
+        Returns:
+            True if any field differs and update is needed
+        """
+        if existing.name != name:
+            logger.debug(f"Tool name changed: {existing.name} -> {name}")
+            return True
+
+        if existing.description != description:
+            logger.debug("Tool description changed")
+            return True
+
+        if existing.ai_model_prompt_guidance != ai_model_prompt_guidance:
+            logger.debug("Tool AI guidance changed")
+            return True
+
+        if existing.input_schema_hash != input_schema_hash:
+            logger.debug(
+                f"Tool input schema hash changed: "
+                f"{existing.input_schema_hash} -> {input_schema_hash}"
+            )
+            return True
+
+        if existing.tool_metadata != metadata:
+            logger.debug("Tool metadata changed")
+            return True
+
+        # Compare output schemas by hash
+        existing_hashes = set()
+        if existing.output_schemas:
+            existing_hashes = {s.schema_hash for s in existing.output_schemas if s.schema_hash}
+
+        desired_hashes = set()
+        if output_schemas:
+            desired_hashes = {s.schema_hash for s in output_schemas}
+
+        if existing_hashes != desired_hashes:
+            logger.debug("Tool output schemas changed")
+            return True
+
+        return False
+
+    def _sync_output_schemas(
+        self,
+        tool_id: int,
+        existing_schemas: list[OutputSchemaResponse] | None,
+        desired_schemas: list[OutputSchemaConfig] | None,
+    ) -> bool:
+        """
+        Sync output schemas - add missing, remove stale.
+
+        Args:
+            tool_id: Tool ID to update
+            existing_schemas: Current output schemas
+            desired_schemas: Desired output schemas
+
+        Returns:
+            True if sync was successful
+        """
+        existing_by_hash: dict[str, OutputSchemaResponse] = {}
+        if existing_schemas:
+            for schema in existing_schemas:
+                if schema.schema_hash:
+                    existing_by_hash[schema.schema_hash] = schema
+
+        desired_by_hash: dict[str, OutputSchemaConfig] = {}
+        if desired_schemas:
+            for schema in desired_schemas:
+                desired_by_hash[schema.schema_hash] = schema
+
+        existing_hashes = set(existing_by_hash.keys())
+        desired_hashes = set(desired_by_hash.keys())
+
+        # Remove stale schemas FIRST (before adding new ones)
+        # This is important when is_primary=True, as only one primary can exist
+        for hash_value in existing_hashes - desired_hashes:
+            schema = existing_by_hash[hash_value]
+            if not self.remove_output_schema(tool_id, schema.id):
+                return False
+
+        # Add new schemas
+        for hash_value in desired_hashes - existing_hashes:
+            schema = desired_by_hash[hash_value]
+            if not self.add_output_schema(tool_id, schema):
+                return False
+
+        return True
+
+    def register_tool_with_hash(
+        self,
+        name: str,
+        tool_technical_name: str,
+        description: str | None = None,
+        ai_model_prompt_guidance: str | None = None,
+        input_schema_hash: str | None = None,
+        output_schemas: list[OutputSchemaConfig] | None = None,
+        metadata: dict | None = None,
+    ) -> bool:
+        """
+        Register a tool with hash-based schema references.
+
+        This method implements smart update logic:
+        1. Checks if tool exists
+        2. If not exists, creates new tool
+        3. If exists, compares all fields
+        4. Only updates if fields have changed
+
+        Args:
+            name: Display name of the tool
+            tool_technical_name: Technical identifier (unique per agent+version)
+            description: Tool description
+            ai_model_prompt_guidance: Guidance for AI models using this tool
+            input_schema_hash: SHA-256 hash of input schema
+            output_schemas: List of output schema configurations
+            metadata: Additional metadata (can include icon)
+
+        Returns:
+            True if registration succeeded, False otherwise
+
+        Reference: Java ToolRegistryClient.registerToolWithHash()
+        """
+        with self._lock:
+            logger.info(
+                f"Registering tool '{sanitize_for_logging(name)}' "
+                f"(technical: {sanitize_for_logging(tool_technical_name)})"
+            )
+
+            # Check if tool exists
+            existing = self._get_existing_tool(tool_technical_name)
+
+            if existing is None:
+                # Create new tool
+                return self._create_tool(
+                    name=name,
+                    tool_technical_name=tool_technical_name,
+                    description=description,
+                    ai_model_prompt_guidance=ai_model_prompt_guidance,
+                    input_schema_hash=input_schema_hash,
+                    output_schemas=output_schemas,
+                    metadata=metadata,
+                )
+
+            # Tool exists - check if update needed
+            if not self._needs_update(
+                existing=existing,
+                name=name,
+                description=description,
+                ai_model_prompt_guidance=ai_model_prompt_guidance,
+                input_schema_hash=input_schema_hash,
+                output_schemas=output_schemas,
+                metadata=metadata,
+            ):
+                logger.info(
+                    f"Tool '{sanitize_for_logging(tool_technical_name)}' unchanged, skipping update"
+                )
+                self._registered_tool_ids.add(existing.tool_id)
+                return True
+
+            # Update existing tool
+            return self._update_tool(
+                tool_id=existing.tool_id,
+                name=name,
+                description=description,
+                ai_model_prompt_guidance=ai_model_prompt_guidance,
+                input_schema_hash=input_schema_hash,
+                output_schemas=output_schemas,
+                existing_schemas=existing.output_schemas,
+                metadata=metadata,
+            )
+
+    def _create_tool(
+        self,
+        name: str,
+        tool_technical_name: str,
+        description: str | None,
+        ai_model_prompt_guidance: str | None,
+        input_schema_hash: str | None,
+        output_schemas: list[OutputSchemaConfig] | None,
+        metadata: dict | None,
+    ) -> bool:
+        """Create a new tool."""
+        registration = ToolRegistration(
+            agent_id=self._agent_id,
+            name=name,
+            tool_technical_name=tool_technical_name,
+            description=description,
+            ai_model_prompt_guidance=ai_model_prompt_guidance,
+            input_schema_hash=input_schema_hash,
+            output_schemas=output_schemas,
+            version_number=1,
+            status="ACTIVE",
+            tool_metadata=metadata,
+        )
+
+        response = self._request_with_retry(
+            "POST",
+            self.API_PATH,
+            json=registration.model_dump(exclude_none=True),
+        )
+
+        if response is None:
+            logger.error(f"Failed to create tool '{sanitize_for_logging(tool_technical_name)}'")
+            return False
+
+        if response.status_code in (200, 201):
+            data = response.json()
+            tool_id = data.get("tool_id")
+            if tool_id:
+                self._registered_tool_ids.add(tool_id)
+            logger.info(
+                f"Successfully created tool '{sanitize_for_logging(tool_technical_name)}' "
+                f"(ID: {tool_id})"
+            )
+            return True
+        elif response.status_code == 409:
+            # Conflict - tool already exists (race condition), treat as success
+            logger.info(f"Tool '{sanitize_for_logging(tool_technical_name)}' already exists (409)")
+            # Try to fetch the existing tool to track its ID
+            existing = self._get_existing_tool(tool_technical_name)
+            if existing:
+                self._registered_tool_ids.add(existing.tool_id)
+            return True
+        else:
+            logger.error(
+                f"Failed to create tool '{sanitize_for_logging(tool_technical_name)}': "
+                f"{response.status_code} - {sanitize_for_logging(response.text)}"
+            )
+            return False
+
+    def _update_tool(
+        self,
+        tool_id: int,
+        name: str,
+        description: str | None,
+        ai_model_prompt_guidance: str | None,
+        input_schema_hash: str | None,
+        output_schemas: list[OutputSchemaConfig] | None,
+        existing_schemas: list[OutputSchemaResponse] | None,
+        metadata: dict | None,
+    ) -> bool:
+        """Update an existing tool."""
+        update_payload = {
+            "name": name,
+            "description": description,
+            "ai_model_prompt_guidance": ai_model_prompt_guidance,
+            "input_schema_hash": input_schema_hash,
+            "tool_metadata": metadata,
+        }
+
+        # Remove None values
+        update_payload = {k: v for k, v in update_payload.items() if v is not None}
+
+        response = self._request_with_retry(
+            "PUT",
+            f"{self.API_PATH}/{tool_id}",
+            json=update_payload,
+        )
+
+        if response is None:
+            logger.error(f"Failed to update tool {tool_id}")
+            return False
+
+        if response.status_code not in (200, 204):
+            logger.error(
+                f"Failed to update tool {tool_id}: "
+                f"{response.status_code} - {sanitize_for_logging(response.text)}"
+            )
+            return False
+
+        # Sync output schemas
+        if not self._sync_output_schemas(tool_id, existing_schemas, output_schemas):
+            logger.warning(f"Failed to sync output schemas for tool {tool_id}")
+
+        self._registered_tool_ids.add(tool_id)
+        logger.info(f"Successfully updated tool {tool_id}")
+        return True
+
+    def add_output_schema(self, tool_id: int, schema: OutputSchemaConfig) -> bool:
+        """
+        Add an output schema to a tool.
+
+        Args:
+            tool_id: Tool ID
+            schema: Output schema configuration
+
+        Returns:
+            True if successful
+        """
+        response = self._request_with_retry(
+            "POST",
+            f"{self.API_PATH}/{tool_id}/output-schemas",
+            json={
+                "schema_hash": schema.schema_hash,
+                "schema_name": schema.schema_name,
+                "is_primary": schema.is_primary,
+                "description": schema.description,
+            },
+        )
+
+        if response is None:
+            return False
+
+        if response.status_code in (200, 201):
+            logger.debug(f"Added output schema '{schema.schema_name}' to tool {tool_id}")
+            return True
+        elif response.status_code == 409:
+            # Already exists
+            logger.debug(f"Output schema '{schema.schema_name}' already exists for tool {tool_id}")
+            return True
+        else:
+            logger.error(
+                f"Failed to add output schema to tool {tool_id}: "
+                f"{response.status_code} - {sanitize_for_logging(response.text)}"
+            )
+            return False
+
+    def remove_output_schema(self, tool_id: int, schema_id: int) -> bool:
+        """
+        Remove an output schema from a tool.
+
+        Args:
+            tool_id: Tool ID
+            schema_id: Output schema ID to remove
+
+        Returns:
+            True if successful
+        """
+        response = self._request_with_retry(
+            "DELETE",
+            f"{self.API_PATH}/{tool_id}/output-schemas/{schema_id}",
+        )
+
+        if response is None:
+            return False
+
+        if response.status_code in (200, 204):
+            logger.debug(f"Removed output schema {schema_id} from tool {tool_id}")
+            return True
+        elif response.status_code == 404:
+            # Already removed
+            logger.debug(f"Output schema {schema_id} already removed from tool {tool_id}")
+            return True
+        else:
+            logger.error(
+                f"Failed to remove output schema {schema_id} from tool {tool_id}: "
+                f"{response.status_code} - {sanitize_for_logging(response.text)}"
+            )
+            return False
+
+    def activate_tool(self, tool_id: int) -> bool:
+        """
+        Activate a tool.
+
+        Args:
+            tool_id: Tool ID to activate
+
+        Returns:
+            True if successful
+        """
+        response = self._request_with_retry(
+            "PATCH",
+            f"{self.API_PATH}/{tool_id}/activate",
+        )
+
+        if response is None:
+            return False
+
+        if response.status_code in (200, 204):
+            logger.info(f"Activated tool {tool_id}")
+            return True
+        else:
+            logger.error(
+                f"Failed to activate tool {tool_id}: "
+                f"{response.status_code} - {sanitize_for_logging(response.text)}"
+            )
+            return False
+
+    def deactivate_tool(self, tool_id: int) -> bool:
+        """
+        Deactivate a tool.
+
+        Args:
+            tool_id: Tool ID to deactivate
+
+        Returns:
+            True if successful
+        """
+        response = self._request_with_retry(
+            "PATCH",
+            f"{self.API_PATH}/{tool_id}/deactivate",
+        )
+
+        if response is None:
+            return False
+
+        if response.status_code in (200, 204):
+            logger.info(f"Deactivated tool {tool_id}")
+            return True
+        else:
+            logger.error(
+                f"Failed to deactivate tool {tool_id}: "
+                f"{response.status_code} - {sanitize_for_logging(response.text)}"
+            )
+            return False
+
+    def close(self) -> None:
+        """
+        Close the client, deactivating all registered tools.
+
+        Reference: Java ToolRegistryClient deactivates all tools on shutdown.
+        """
+        with self._lock:
+            logger.info(
+                f"Closing ToolRegistryClient, deactivating {len(self._registered_tool_ids)} tools"
+            )
+
+            for tool_id in self._registered_tool_ids:
+                try:
+                    self.deactivate_tool(tool_id)
+                except Exception as e:
+                    logger.warning(
+                        f"Failed to deactivate tool {tool_id} during close: "
+                        f"{sanitize_for_logging(str(e))}"
+                    )
+
+            self._registered_tool_ids.clear()
+            self._client.close()
+            logger.info("ToolRegistryClient closed")

base_agent/utils/__init__.py

@@ -0,0 +1,5 @@
+"""Utility modules for BaseAgent."""
+
+from .version_utils import get_version_info
+
+__all__ = ["get_version_info"]