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

base_agent/schemas/schema_registry_client.py

@@ -0,0 +1,493 @@
+"""Schema registry client for registering Pydantic models with the platform.
+
+This is the Python equivalent of Java's SchemaRegistryClient.
+Reference: base-agent/src/main/java/one/ai/platform/baseagent/agent/schemas/SchemaRegistryClient.java
+"""
+
+import json
+import logging
+import random
+import threading
+import time
+from collections.abc import Callable
+from typing import TypeVar
+
+import httpx
+from pydantic import BaseModel
+
+from base_agent.schemas.source_type import SourceType
+from base_agent.schemas.technical_name_validator import to_technical_name_label
+from base_agent.utils.logger import sanitize_for_logging
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+class SchemaRegistryClient:
+    """
+    Client for registering schemas with the platform schema registry.
+
+    Provides feature parity with Java SchemaRegistryClient:
+    - JSON schema generation from Pydantic models
+    - Schema registration with admin-gateway
+    - Schema hash lookup by class or name
+    - Thread-safe operations
+    - Health check integration via is_live property
+    - Retry logic with exponential backoff
+
+    Reference: Java SchemaRegistryClient.java
+    """
+
+    API_PATH = "/api/v1/schema-registry"
+    MAX_RETRY_DURATION_SECONDS = 30
+    INITIAL_BACKOFF_MS = 500
+    BACKOFF_MULTIPLIER = 1.5
+
+    def __init__(self, admin_gateway_url: str):
+        """
+        Initialize the schema registry client.
+
+        Args:
+            admin_gateway_url: Base URL of the admin gateway (e.g., "http://localhost:8888")
+        """
+        if not admin_gateway_url:
+            raise ValueError("admin_gateway_url is required")
+
+        self._admin_gateway_url = admin_gateway_url.rstrip("/")
+        self._is_live = True
+        self._lock = threading.Lock()
+        self._client = httpx.Client(
+            base_url=self._admin_gateway_url,
+            timeout=30.0,
+        )
+
+        logger.info(
+            f"SchemaRegistryClient initialized at {sanitize_for_logging(admin_gateway_url)}"
+        )
+
+    @property
+    def is_live(self) -> bool:
+        """Check if the client is live (healthy)."""
+        return self._is_live
+
+    @staticmethod
+    def get_schema_name_for_class(model: type[BaseModel]) -> str:
+        """
+        Convert a Pydantic model class name to its RFC1035 compliant schema name.
+
+        Args:
+            model: The Pydantic model class
+
+        Returns:
+            The RFC1035 compliant schema name
+        """
+        return to_technical_name_label(model.__name__)
+
+    def register_schemas(
+        self,
+        models: list[type[BaseModel]],
+        group_name: str | None = None,
+        source_type: SourceType = SourceType.PLATFORM_AGENT,
+    ) -> dict[str, bool]:
+        """
+        Register schemas for a list of Pydantic model classes.
+
+        Args:
+            models: List of Pydantic model classes to register schemas for
+            group_name: Group name for the schemas (optional)
+            source_type: The source type of the schema
+
+        Returns:
+            Dict mapping schema names to their registration status
+        """
+        results: dict[str, bool] = {}
+
+        for model in models:
+            try:
+                registered = self.register_schema(model, group_name, source_type)
+                schema_name = to_technical_name_label(model.__name__)
+                results[schema_name] = registered
+            except Exception as e:
+                logger.error(
+                    f"Failed to register schema for class: "
+                    f"{sanitize_for_logging(model.__name__)}: {sanitize_for_logging(str(e))}"
+                )
+                schema_name = to_technical_name_label(model.__name__)
+                results[schema_name] = False
+
+        return results
+
+    def register_schema(
+        self,
+        model: type[BaseModel],
+        group_name: str | None = None,
+        source_type: SourceType = SourceType.PLATFORM_AGENT,
+    ) -> bool:
+        """
+        Register a schema for a Pydantic model class.
+
+        Args:
+            model: The Pydantic model class to generate a schema for
+            group_name: Group name for the schema (optional)
+            source_type: The source type of the schema
+
+        Returns:
+            True if a new schema was registered, False if no changes were needed
+        """
+        with self._lock:
+            # Convert class name to RFC1035 compliant schema name
+            schema_name = to_technical_name_label(model.__name__)
+            logger.debug(
+                f"Converting class name '{sanitize_for_logging(model.__name__)}' "
+                f"to RFC1035 compliant schema name '{sanitize_for_logging(schema_name)}'"
+            )
+
+            # Generate JSON schema for the model
+            json_schema = model.model_json_schema()
+            schema_text = json.dumps(json_schema, sort_keys=True)
+
+            # Check if schema already exists and compare content
+            existing_schema = self._get_existing_schema(schema_name, group_name)
+
+            if existing_schema is not None:
+                # Compare normalized schemas (to ignore formatting differences)
+                try:
+                    existing_schema_obj = json.loads(existing_schema)
+                    new_schema_obj = json.loads(schema_text)
+
+                    if existing_schema_obj == new_schema_obj:
+                        logger.debug(
+                            f"Schema for {sanitize_for_logging(schema_name)} is unchanged, "
+                            "skipping registration"
+                        )
+                        return False
+                except json.JSONDecodeError:
+                    # If we can't parse, assume they're different
+                    pass
+
+            # Schema is new or different, register it
+            return self._create_or_update_schema(schema_name, group_name, schema_text, source_type)
+
+    def get_schema_hash_for_class(
+        self,
+        model: type[BaseModel],
+        group_name: str | None = None,
+    ) -> str | None:
+        """
+        Get the schema hash for a given Pydantic model class.
+
+        Args:
+            model: The Pydantic model class to get the schema hash for
+            group_name: Group name for the schema (optional)
+
+        Returns:
+            The SHA-256 hash of the schema, or None if not found
+        """
+        # Convert class name to RFC1035 compliant schema name
+        schema_name = to_technical_name_label(model.__name__)
+        logger.debug(
+            f"Getting schema hash for class: {sanitize_for_logging(model.__name__)} "
+            f"-> schema: {sanitize_for_logging(schema_name)} "
+            f"(group: {sanitize_for_logging(group_name)})"
+        )
+
+        return self.get_schema_hash_by_name(schema_name, group_name)
+
+    def get_schema_hash_by_name(
+        self,
+        name: str,
+        group_name: str | None = None,
+    ) -> str | None:
+        """
+        Get the schema hash for a given schema name and group.
+
+        Args:
+            name: The schema name to look up
+            group_name: Group name for the schema (optional, defaults to "default")
+
+        Returns:
+            The SHA-256 hash of the schema, or None if not found
+        """
+        logger.info(
+            f"Getting schema hash by name: {sanitize_for_logging(name)} "
+            f"(group: {sanitize_for_logging(group_name)})"
+        )
+
+        def operation() -> str | None:
+            encoded_group = group_name if group_name else "default"
+
+            response = self._client.get(
+                self.API_PATH,
+                params={
+                    "name": name,
+                    "group": encoded_group,
+                    "latest_only": "true",
+                },
+            )
+
+            if response.status_code == 200:
+                schemas = response.json()
+                if schemas:
+                    schema_hash = schemas[0].get("schema_hash")
+                    logger.debug(
+                        f"Found schema hash for {sanitize_for_logging(name)}: "
+                        f"{sanitize_for_logging(schema_hash)}"
+                    )
+                    return schema_hash
+                else:
+                    logger.warning(
+                        f"No schema found for name: {sanitize_for_logging(name)} "
+                        f"(group: {sanitize_for_logging(group_name)})"
+                    )
+                    return None
+            else:
+                logger.warning(
+                    f"Failed to get schema hash for name {sanitize_for_logging(name)} "
+                    f"(group {sanitize_for_logging(group_name)}): "
+                    f"HTTP {response.status_code}"
+                )
+                return None
+
+        try:
+            return self._execute_with_retries(operation)
+        except Exception as e:
+            logger.error(f"Failed to get schema hash: {sanitize_for_logging(str(e))}")
+            return None
+
+    def get_schema_by_hash(
+        self,
+        schema_hash: str,
+        latest_only: bool = True,
+    ) -> str | None:
+        """
+        Get a schema by its hash value.
+
+        Args:
+            schema_hash: The hash value of the schema
+            latest_only: Whether to return only the latest version with this hash
+
+        Returns:
+            Schema text content, or None if not found
+        """
+        logger.debug(f"Getting schema by hash: {sanitize_for_logging(schema_hash)}")
+
+        def operation() -> str | None:
+            response = self._client.get(
+                f"{self.API_PATH}/by-hash/{schema_hash}",
+                params={"latest_only": str(latest_only).lower()},
+            )
+
+            if response.status_code == 200:
+                data = response.json()
+                return data.get("schema_text")
+            else:
+                logger.warning(
+                    f"Failed to get schema by hash {sanitize_for_logging(schema_hash)}: "
+                    f"HTTP {response.status_code}"
+                )
+                return None
+
+        try:
+            return self._execute_with_retries(operation)
+        except Exception as e:
+            logger.error(f"Failed to get schema by hash: {sanitize_for_logging(str(e))}")
+            return None
+
+    def _is_retryable_exception(self, e: Exception) -> bool:
+        """Check if an exception is retryable."""
+        if isinstance(e, (httpx.ConnectError, httpx.ConnectTimeout, httpx.ReadTimeout)):
+            return True
+        if e.__cause__ is not None and isinstance(e.__cause__, Exception):
+            return self._is_retryable_exception(e.__cause__)
+        return False
+
+    def _execute_with_retries(self, operation: Callable[[], T]) -> T:
+        """
+        Execute an operation with retries.
+
+        Args:
+            operation: The operation to execute
+
+        Returns:
+            The result of the operation
+
+        Raises:
+            Exception: If all retries are exhausted
+        """
+        start_time = time.time()
+        end_time = start_time + self.MAX_RETRY_DURATION_SECONDS
+
+        attempt = 0
+        backoff_ms = self.INITIAL_BACKOFF_MS
+        last_exception: Exception | None = None
+
+        while time.time() < end_time:
+            attempt += 1
+            try:
+                return operation()
+            except Exception as e:
+                last_exception = e
+                if not self._is_retryable_exception(e):
+                    logger.warning(
+                        f"Non-retryable exception occurred, abandoning retry: "
+                        f"{sanitize_for_logging(str(e))}"
+                    )
+                    break
+
+                if time.time() + (backoff_ms / 1000) < end_time:
+                    logger.info(
+                        f"Attempt {attempt} failed with retryable error: "
+                        f"{sanitize_for_logging(str(e))}. Retrying in {backoff_ms}ms..."
+                    )
+                    time.sleep(backoff_ms / 1000)
+                    # Apply exponential backoff with jitter
+                    jitter = 0.9 + random.random() * 0.2
+                    backoff_ms = int(backoff_ms * self.BACKOFF_MULTIPLIER * jitter)
+                else:
+                    logger.warning(f"Retry time budget exceeded after {attempt} attempts")
+                    break
+
+        # If we got here, all retries failed
+        logger.error(f"Schema registry operation failed after {attempt} attempts")
+        self._is_live = False
+
+        if last_exception is not None:
+            raise last_exception
+        raise RuntimeError(f"Failed after {attempt} attempts with unknown error")
+
+    def _get_existing_schema(
+        self,
+        name: str,
+        group_name: str | None,
+    ) -> str | None:
+        """
+        Get an existing schema from the registry.
+
+        Args:
+            name: Schema name
+            group_name: Group name
+
+        Returns:
+            The schema text, or None if not found
+        """
+        logger.debug(
+            f"Getting existing schema: {sanitize_for_logging(name)} "
+            f"({sanitize_for_logging(group_name)})"
+        )
+
+        try:
+
+            def operation() -> str | None:
+                encoded_group = group_name if group_name else "default"
+
+                response = self._client.get(
+                    self.API_PATH,
+                    params={
+                        "name": name,
+                        "group": encoded_group,
+                        "latest_only": "true",
+                    },
+                )
+
+                if response.status_code == 200:
+                    schemas = response.json()
+                    if schemas:
+                        schema_id = schemas[0].get("schema_id")
+                        # Get the schema text content
+                        return self._get_schema_by_id(schema_id)
+                return None
+
+            return self._execute_with_retries(operation)
+        except Exception as e:
+            logger.warning(
+                f"Failed to get existing schema: {sanitize_for_logging(name)} "
+                f"({sanitize_for_logging(str(e))})"
+            )
+            return None
+
+    def _get_schema_by_id(self, schema_id: int) -> str | None:
+        """
+        Get a schema by its ID.
+
+        Args:
+            schema_id: Schema ID
+
+        Returns:
+            Schema text content, or None if not found
+        """
+        logger.debug(f"Getting schema by ID: {schema_id}")
+
+        def operation() -> str | None:
+            response = self._client.get(f"{self.API_PATH}/{schema_id}")
+
+            if response.status_code == 200:
+                data = response.json()
+                return data.get("schema_text")
+            else:
+                logger.warning(
+                    f"Failed to get schema by ID {schema_id}: HTTP {response.status_code}"
+                )
+                return None
+
+        return self._execute_with_retries(operation)
+
+    def _create_or_update_schema(
+        self,
+        name: str,
+        group_name: str | None,
+        schema_text: str,
+        source_type: SourceType,
+    ) -> bool:
+        """
+        Create or update a schema in the registry.
+
+        Args:
+            name: Schema name
+            group_name: Group name
+            schema_text: Schema text content
+            source_type: The source type of the schema
+
+        Returns:
+            True if successful, False otherwise
+        """
+
+        def operation() -> bool:
+            payload = {
+                "schema_data": {
+                    "name": name,
+                    "group": group_name if group_name else "default",
+                    "format": "JSON_SCHEMA",
+                    "compatibility": "FULL",
+                    "source_type": source_type.value,
+                },
+                "schema_text": {
+                    "schema_text": schema_text,
+                },
+            }
+
+            response = self._client.post(
+                f"{self.API_PATH}/create-with-text",
+                json=payload,
+            )
+
+            if response.status_code == 201:
+                logger.info(f"Successfully registered schema for {sanitize_for_logging(name)}")
+                return True
+            else:
+                logger.error(
+                    f"Failed to register schema: HTTP {response.status_code} - "
+                    f"{sanitize_for_logging(response.text)}"
+                )
+                return False
+
+        try:
+            return self._execute_with_retries(operation)
+        except Exception as e:
+            logger.error(f"Failed to create/update schema: {sanitize_for_logging(str(e))}")
+            return False
+
+    def close(self) -> None:
+        """Close the HTTP client."""
+        self._client.close()
+        logger.info("SchemaRegistryClient closed")

base_agent/schemas/source_type.py

@@ -0,0 +1,24 @@
+"""Source type enum for schema registry.
+
+This module provides the SourceType enum matching Java's SourceType.java.
+Reference: base-agent/src/main/java/one/ai/platform/baseagent/agent/schemas/SourceType.java
+"""
+
+from enum import Enum
+
+
+class SourceType(str, Enum):
+    """
+    Schema source types for the schema registry.
+
+    These values correspond to the source_type field in the schema registry API.
+    """
+
+    NAMED_QUERY = "NAMED_QUERY"
+    PDF = "PDF"
+    WORKFLOW = "WORKFLOW"
+    TASK_AGENT = "TASK_AGENT"
+    UNSTRUCTURED_DATA_METADATA = "UNSTRUCTURED_DATA_METADATA"
+    PLATFORM_AGENT = "PLATFORM_AGENT"
+    PYTHON_SANDBOX_AGENT = "PYTHON_SANDBOX_AGENT"
+    CUSTOM_AGENT = "CUSTOM_AGENT"

base_agent/schemas/technical_name_validator.py

@@ -0,0 +1,147 @@
+"""Technical name validation and conversion utilities.
+
+This module provides RFC1035-compliant name validation and conversion,
+porting Java's TechnicalNameValidator.
+Reference: base-agent/src/main/java/one/ai/platform/baseagent/utils/TechnicalNameValidator.java
+"""
+
+import re
+
+MAX_LABEL_LENGTH = 63
+
+# Regex patterns matching Java implementation
+START_WITH_LETTER = re.compile(r"^[a-zA-Z]")
+END_WITH_LETTER_OR_DIGIT = re.compile(r"[a-zA-Z0-9]$")
+VALID_CHARACTERS = re.compile(r"^[a-zA-Z0-9\-_]+$")
+
+
+def is_valid_label(label: str) -> bool:
+    """
+    Validate if a string is a valid technical name.
+
+    Args:
+        label: The label to validate
+
+    Returns:
+        True if the label is valid, False otherwise
+    """
+    if not label or len(label) > MAX_LABEL_LENGTH:
+        return False
+
+    # Must start with a letter
+    if not START_WITH_LETTER.match(label):
+        return False
+
+    # Must end with a letter or digit
+    if not END_WITH_LETTER_OR_DIGIT.search(label):
+        return False
+
+    # Can only contain letters, digits, hyphens, and underscores
+    return bool(VALID_CHARACTERS.match(label))
+
+
+def to_technical_name_label(class_name: str) -> str:
+    """
+    Convert a class name to a technical name compliant label.
+
+    This method converts camelCase/PascalCase to kebab-case and ensures
+    the result is RFC1035 compliant.
+
+    Args:
+        class_name: The class name to convert (e.g., 'GetIssueInput')
+
+    Returns:
+        Technical name compliant label (e.g., 'get-issue-input')
+
+    Raises:
+        ValueError: If class_name is empty or cannot be converted
+
+    Examples:
+        >>> to_technical_name_label('GetIssueInput')
+        'get-issue-input'
+        >>> to_technical_name_label('XMLParser')
+        'xml-parser'
+        >>> to_technical_name_label('Test123Class')
+        'test-123-class'
+    """
+    if not class_name:
+        raise ValueError("Class name cannot be null or empty")
+
+    result = class_name
+
+    # Replace underscores with hyphens first
+    result = result.replace("_", "-")
+
+    # Insert hyphen before uppercase letters (except at the beginning)
+    result = re.sub(r"([a-z0-9])([A-Z])", r"\1-\2", result)
+
+    # Handle consecutive uppercase letters (e.g., XMLParser -> XML-Parser)
+    result = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1-\2", result)
+
+    # Insert hyphen between letters and numbers
+    result = re.sub(r"([a-zA-Z])([0-9])", r"\1-\2", result)
+    result = re.sub(r"([0-9])([a-zA-Z])", r"\1-\2", result)
+
+    # Convert to lowercase
+    result = result.lower()
+
+    # Replace other special characters with hyphens
+    result = re.sub(r"[^a-z0-9\-]", "-", result)
+
+    # Remove leading/trailing hyphens
+    result = re.sub(r"^-+|-+$", "", result)
+
+    # Replace multiple consecutive hyphens with a single hyphen
+    result = re.sub(r"-+", "-", result)
+
+    # Ensure it starts with a letter
+    if result and not result[0].isalpha():
+        result = "schema-" + result
+
+    # Ensure it ends with a letter or digit
+    while result and result.endswith("-"):
+        result = result[:-1]
+
+    # Truncate if too long
+    if len(result) > MAX_LABEL_LENGTH:
+        result = result[:MAX_LABEL_LENGTH]
+        # Ensure it still ends with a letter or digit after truncation
+        while result and not result[-1].isalnum():
+            result = result[:-1]
+
+    # Final validation
+    if not is_valid_label(result):
+        raise ValueError(
+            f"Unable to convert class name '{class_name}' to technical name "
+            f"compliant label. Result: '{result}'"
+        )
+
+    return result
+
+
+def get_validation_error(label: str) -> str | None:
+    """
+    Get a validation error message for a label.
+
+    Args:
+        label: The label to validate
+
+    Returns:
+        Error message if invalid, None if valid
+    """
+    if not label:
+        return "Label cannot be null or empty"
+
+    if len(label) > MAX_LABEL_LENGTH:
+        return "Label must be 63 characters or less"
+
+    if not START_WITH_LETTER.match(label):
+        return "Label must start with a letter"
+
+    if not END_WITH_LETTER_OR_DIGIT.search(label):
+        return "Label must end with a letter or digit"
+
+    if not VALID_CHARACTERS.match(label):
+        return "Label can only contain letters, digits, hyphens, and underscores"
+
+    return None

base_agent/state/__init__.py

@@ -0,0 +1,13 @@
+"""State management module for BaseAgent."""
+
+from .logical_clock_tracker import LogicalClockTracker
+from .session_tracker import SessionTracker
+from .store import InMemoryStateStore, RedisStateStore, StateStore
+
+__all__ = [
+    "StateStore",
+    "InMemoryStateStore",
+    "RedisStateStore",
+    "SessionTracker",
+    "LogicalClockTracker",
+]