aip-agents-binary 0.0.0__py3-none-any.whl

aip_agents/schema/agent.py

@@ -0,0 +1,117 @@
+"""Schema definitions for agent configuration types.
+
+Authors:
+    Christian Trisno Sen Long Chen (christian.t.s.l.chen@gdplabs.id)
+    Reinhart Linanda (reinhart.linanda@gdplabs.id)
+"""
+
+from enum import StrEnum
+from typing import Any
+
+from a2a.types import AgentCard
+from gllm_core.utils.retry import RetryConfig
+from pydantic import BaseModel, Field
+
+try:
+    # Pydantic v2 preferred config
+    from pydantic import ConfigDict  # type: ignore
+except Exception:  # pragma: no cover
+    ConfigDict = None  # type: ignore
+
+__all__ = [
+    "CredentialType",
+    "StreamMode",
+    "HttpxClientOptions",
+    "A2AClientConfig",
+    "BaseAgentConfig",
+    "AgentConfig",
+    "LangflowAgentConfig",
+]
+
+
+class CredentialType(StrEnum):
+    """Credential type enumeration for type safety and better developer experience."""
+
+    API_KEY = "api_key"
+    FILE = "file"
+    DICT = "dict"
+
+
+class StreamMode(StrEnum):
+    """LangGraph stream modes for astream operations."""
+
+    VALUES = "values"
+    CUSTOM = "custom"
+    MESSAGES = "messages"
+
+
+class HttpxClientOptions(BaseModel):
+    """Options for the HTTP client."""
+
+    timeout: float = 360.0
+    trust_env: bool = False
+    follow_redirects: bool = True
+
+    if ConfigDict is not None:
+        # Pydantic v2 style config
+        model_config = ConfigDict(extra="allow")  # type: ignore[assignment]
+    else:  # pragma: no cover
+
+        class Config:  # type: ignore[no-redef]
+            """Pydantic v1 fallback config for HttpxClientOptions."""
+
+            extra = "allow"
+
+
+class A2AClientConfig(BaseModel):
+    """Configuration for A2A client."""
+
+    discovery_urls: list[str] | None = None
+    known_agents: dict[str, AgentCard] = Field(default_factory=dict)
+    httpx_client_options: HttpxClientOptions | None = HttpxClientOptions()
+
+
+class BaseAgentConfig(BaseModel):
+    """Base configuration for agent implementations."""
+
+    tools: list[Any] | None = None
+    default_hyperparameters: dict[str, Any] | None = None
+
+    if ConfigDict is not None:
+        model_config = ConfigDict(extra="allow")  # type: ignore[assignment]
+    else:  # pragma: no cover
+
+        class Config:  # type: ignore[no-redef]
+            """Pydantic v1 fallback config for BaseAgentConfig."""
+
+            extra = "allow"
+
+
+class AgentConfig(BaseAgentConfig):
+    """Configuration for agent implementations with language model settings."""
+
+    lm_name: str | None = None
+    lm_hyperparameters: dict[str, Any] | None = None
+    lm_provider: str | None = None
+    lm_base_url: str | None = None
+    lm_api_key: str | None = None
+    lm_credentials: str | dict[str, Any] | None = None
+    lm_retry_config: RetryConfig | None = None
+    a2ui: dict[str, Any] | None = None
+
+
+class LangflowAgentConfig(BaseAgentConfig):
+    """Configuration for Langflow agent implementations."""
+
+    flow_id: str
+    base_url: str | None = None
+    api_key: str | None = None
+
+    if ConfigDict is not None:
+        model_config = ConfigDict(extra="allow")  # type: ignore[assignment]
+    else:  # pragma: no cover
+
+        class Config:  # type: ignore[no-redef]
+            """Pydantic v1 fallback config for LangflowAgentConfig."""
+
+            extra = "allow"

aip_agents/schema/artifacts.py

@@ -0,0 +1,75 @@
+"""Shared typed contracts for artifact payloads exchanged across tool and agent boundaries.
+
+These contracts codify the dict shapes already produced by
+``aip_agents.utils.artifact_helpers`` and ``aip_agents.utils.artifacts.tracker``
+so downstream typing slices can import named structures instead of repeating
+``dict[str, Any]``. They describe the *serialized* artifact payload (the dict
+form of ``aip_agents.a2a.types.Artifact`` plus an optional internal ``path``
+used by the tracker); they do not replace the Pydantic ``Artifact`` model.
+
+Authors:
+    Putu Ravindra Wiguna (putu.r.wiguna@gdplabs.id)
+"""
+
+from __future__ import annotations
+
+from typing import Any, Required, TypedDict
+
+__all__ = [
+    "ArtifactPayload",
+    "TrackedArtifact",
+    "ArtifactBuildError",
+    "ToolArtifactResponse",
+    "ToolMultiArtifactResponse",
+]
+
+
+class ArtifactPayload(TypedDict, total=False):
+    """Serialized artifact payload produced by artifact helpers.
+
+    Mirrors ``aip_agents.a2a.types.Artifact.model_dump()``.
+
+    Wire-mandatory keys (always present when emitted by the A2A pipeline) are
+    marked ``Required``: ``artifact_type``, ``data``, ``name``, ``mime_type``.
+    ``description`` defaults to an empty string and ``metadata`` is omitted when
+    absent.
+    """
+
+    artifact_type: Required[str]
+    data: Required[str]
+    name: Required[str]
+    description: str
+    mime_type: Required[str]
+    metadata: dict[str, Any] | None
+
+
+class TrackedArtifact(ArtifactPayload, total=False):
+    """``ArtifactPayload`` plus the tracker-only ``path`` provenance key.
+
+    Emitted by ``ArtifactTracker.build_artifacts``; ``path`` is internal and
+    must be stripped before serialization to the A2A wire contract
+    (``ArtifactPayload``).
+    """
+
+    path: str
+
+
+class ArtifactBuildError(TypedDict):
+    """Error entry returned by ``ArtifactTracker.build_artifacts`` for a file that could not be turned into an artifact."""
+
+    path: str
+    reason: str
+
+
+class ToolArtifactResponse(TypedDict):
+    """Single-artifact tool response shape returned by ``ArtifactHandler.create_artifact``."""
+
+    result: str
+    artifact: ArtifactPayload
+
+
+class ToolMultiArtifactResponse(TypedDict):
+    """Multi-artifact tool response shape returned by artifact command helpers."""
+
+    result: str
+    artifacts: list[ArtifactPayload]

aip_agents/schema/attachments.py

@@ -0,0 +1,425 @@
+"""Runtime attachment contracts for user inputs and generated artifacts.
+
+Authors:
+    Raymond Christopher (raymond.christopher@gdplabs.id)
+"""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import os
+from enum import StrEnum
+from typing import Any, Required, TypedDict
+from urllib.parse import urlparse
+
+from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator
+
+from aip_agents.a2a.types import MimeType, get_mime_type_from_filename
+
+
+class AttachmentKind(StrEnum):
+    """Supported first-class attachment kinds."""
+
+    IMAGE = "image"
+    DOCUMENT = "document"
+
+
+class AttachmentSource(StrEnum):
+    """Where an attachment reference originated."""
+
+    USER_INPUT = "user_input"
+    LOCAL_TOOL_OUTPUT = "local_tool_output"
+    SUB_AGENT_ARTIFACT = "sub_agent_artifact"
+    REMOTE_UPLOAD = "remote_upload"
+    PLATFORM_ARTIFACT = "platform_artifact"
+
+
+_HTTP_SCHEME = "http"
+_HTTP_PROTOCOL_PREFIXES = (f"{_HTTP_SCHEME}://", f"{_HTTP_SCHEME}s://")
+SUPPORTED_IMAGE_MIME_TYPES = ("image/png", "image/jpeg", "image/webp")
+SUPPORTED_DOCUMENT_MIME_TYPES = (
+    MimeType.APPLICATION_PDF,
+    MimeType.APPLICATION_MSWORD,
+    MimeType.APPLICATION_EXCEL,
+    MimeType.APPLICATION_DOCX,
+    MimeType.APPLICATION_XLSX,
+)
+MAX_IMAGE_ATTACHMENT_BYTES = 20 * 1024 * 1024
+_UNSAFE_TRACE_METADATA_KEYS = {
+    "base64",
+    "bytes",
+    "data",
+    "data_ref",
+    "file_data",
+    "file_uri",
+    "local_path",
+    "path",
+    "presigned_url",
+    "raw",
+    "raw_data",
+    "s3_url",
+    "signed_url",
+    "url",
+}
+
+
+def _normalize_trace_metadata_key(key: object) -> str:
+    """Normalize metadata keys before trace-safety blocklist checks."""
+    return str(key).strip().lower().replace("-", "_").replace(" ", "_")
+
+
+def _is_likely_local_path_ref(ref: str) -> bool:
+    """Return whether a ref explicitly names a local filesystem path."""
+    if os.path.isabs(ref):
+        return True
+    normalized_ref = ref.lower()
+    if normalized_ref.startswith(("file://", "./", "../", "~/")):
+        return True
+    if "\\" in ref:
+        return True
+    return len(ref) > 2 and ref[1] == ":" and ref[0].isalpha()
+
+
+class AttachmentEntryPayload(TypedDict, total=False):
+    """Serialized attachment entry — the dict form of ``RuntimeAttachment``.
+
+    Represents the shape stored in metadata ``attachments`` lists and
+    ``attachment_registry`` dicts, and merged by
+    ``utils.attachment_helpers.merge_input_attachments``. Mirrors
+    ``RuntimeAttachment.model_dump()`` so registry/list consumers can use a
+    typed contract instead of ``dict[str, Any]``.
+
+    Identity fields always produced by ``model_dump()`` are marked ``Required``;
+    ``size_bytes`` and ``metadata`` are genuinely optional.
+    """
+
+    id: Required[str]
+    kind: Required[str]
+    source: Required[str]
+    filename: Required[str]
+    mime_type: Required[str]
+    data_ref: Required[str]
+    size_bytes: int | None
+    metadata: dict[str, Any] | None
+
+
+class RuntimeAttachment(BaseModel):
+    """Safe runtime reference for an input attachment or generated artifact."""
+
+    model_config = ConfigDict(use_enum_values=True)
+
+    id: str
+    kind: AttachmentKind
+    source: AttachmentSource
+    filename: str
+    mime_type: str
+    size_bytes: int | None = Field(default=None, ge=0)
+    data_ref: str
+    metadata: dict[str, Any] | None = None
+
+    @field_validator("mime_type")
+    @classmethod
+    def validate_mime_type(cls, value: str, info: ValidationInfo) -> str:
+        """Apply MIME validation based on attachment kind."""
+        raw_kind = info.data.get("kind") if isinstance(info.data, dict) else None
+        kind = AttachmentKind(raw_kind) if raw_kind is not None else None
+        normalized_mime = normalize_mime_type(value)
+
+        if normalized_mime is None:
+            raise ValueError("attachment MIME type must be a non-empty string")  # pragma: no cover
+        if kind == AttachmentKind.IMAGE and normalized_mime not in SUPPORTED_IMAGE_MIME_TYPES:
+            raise ValueError("image attachments must use a supported image MIME type")
+        if kind == AttachmentKind.DOCUMENT and normalized_mime not in SUPPORTED_DOCUMENT_MIME_TYPES:
+            raise ValueError("document attachments must use a supported document MIME type")
+        return normalized_mime
+
+    @field_validator("data_ref")
+    @classmethod
+    def validate_safe_data_ref(cls, value: str) -> str:
+        """Reject inline bytes, base64 data URIs, and likely signed URLs."""
+        normalized_value = value.lower()
+        if normalized_value.startswith("data:"):
+            raise ValueError("data_ref must not contain inline data")
+        if normalized_value.startswith(_HTTP_PROTOCOL_PREFIXES):
+            parsed = urlparse(value)
+            if parsed.query:
+                raise ValueError("data_ref must not contain signed URL query parameters")
+            if "@" in parsed.netloc:
+                raise ValueError("data_ref must not contain credentials")
+        return value
+
+
+class AttachmentRegistry(BaseModel):
+    """Lookup container shared by user inputs and generated artifacts."""
+
+    attachments: dict[str, RuntimeAttachment] = Field(default_factory=dict)
+
+    def register(self, attachment: RuntimeAttachment) -> RuntimeAttachment:
+        """Register an attachment by ID and return the stored value."""
+        self.attachments[attachment.id] = attachment
+        return attachment
+
+    def resolve(self, attachment_id: str) -> RuntimeAttachment | None:
+        """Resolve an attachment by ID."""
+        return self.attachments.get(attachment_id)
+
+    def as_list(self) -> list[dict[str, Any]]:
+        """Return attachment entries as serializable dictionaries."""
+        return [attachment.model_dump() for attachment in self.attachments.values()]
+
+
+def normalize_mime_type(mime_type: str | None) -> str | None:
+    """Normalize MIME strings for MIME support checks."""
+    if not isinstance(mime_type, str):
+        return None
+    normalized = mime_type.split(";", 1)[0].strip().lower()
+    return normalized or None
+
+
+def normalize_image_mime_type(mime_type: str | None) -> str | None:
+    """Backward-compatible image MIME normalization helper."""
+    return normalize_mime_type(mime_type)
+
+
+def is_supported_image_mime(mime_type: str | None) -> bool:
+    """Return whether the MIME type is supported for image attachments."""
+    return normalize_mime_type(mime_type) in SUPPORTED_IMAGE_MIME_TYPES
+
+
+def is_supported_document_mime(mime_type: str | None) -> bool:
+    """Return whether the MIME type is supported for document attachments."""
+    return normalize_mime_type(mime_type) in SUPPORTED_DOCUMENT_MIME_TYPES
+
+
+def create_runtime_attachment(
+    *,
+    attachment_id: str,
+    source: AttachmentSource | str,
+    filename: str,
+    mime_type: str,
+    data_ref: str,
+    kind: AttachmentKind | str = AttachmentKind.IMAGE,
+    size_bytes: int | None = None,
+    metadata: dict[str, Any] | None = None,
+    max_size_bytes: int = MAX_IMAGE_ATTACHMENT_BYTES,
+) -> RuntimeAttachment:
+    """Create a runtime attachment from caller-owned safe metadata."""
+    resolved_kind = AttachmentKind(kind)
+    normalized_mime = normalize_mime_type(mime_type)
+
+    if resolved_kind == AttachmentKind.IMAGE:
+        if normalized_mime is None or normalized_mime not in SUPPORTED_IMAGE_MIME_TYPES:
+            raise ValueError(f"unsupported image MIME type: {mime_type}")
+        if size_bytes is not None and size_bytes > max_size_bytes:
+            raise ValueError(f"image attachment {filename} exceeds {max_size_bytes} bytes")
+    elif resolved_kind == AttachmentKind.DOCUMENT:
+        if normalized_mime is None or normalized_mime not in SUPPORTED_DOCUMENT_MIME_TYPES:
+            raise ValueError(f"unsupported document MIME type: {mime_type}")
+
+    return RuntimeAttachment(
+        id=attachment_id,
+        kind=resolved_kind,
+        source=source,
+        filename=filename,
+        mime_type=normalized_mime,
+        size_bytes=size_bytes,
+        data_ref=data_ref,
+        metadata=metadata,
+    )
+
+
+create_runtime_image_attachment = create_runtime_attachment
+
+
+def merge_attachment_registry(
+    metadata: dict[str, Any] | None,
+    attachments: list[RuntimeAttachment | dict[str, Any]],
+) -> dict[str, Any]:
+    """Merge image attachments into runtime metadata and registry fields."""
+    merged = dict(metadata or {})
+    existing_attachments = merged.get("attachments") if isinstance(merged.get("attachments"), list) else []
+    existing_registry = merged.get("attachment_registry") if isinstance(merged.get("attachment_registry"), dict) else {}
+    attachment_list = [*existing_attachments]
+    registry = dict(existing_registry)
+
+    for attachment in attachments:
+        attachment_dict = attachment.model_dump() if isinstance(attachment, RuntimeAttachment) else dict(attachment)
+        attachment_id = attachment_dict.get("id")
+        if not isinstance(attachment_id, str) or not attachment_id:
+            continue
+        if attachment_id not in registry:
+            attachment_list.append(attachment_dict)
+        registry[attachment_id] = attachment_dict
+
+    merged["attachments"] = attachment_list
+    merged["attachment_registry"] = registry
+    return merged
+
+
+def safe_image_data_ref(value: object) -> str | None:
+    """Return safe non-inline image refs while excluding signed URLs and local paths."""
+    if not isinstance(value, str):
+        return None
+    ref = value.strip()
+    normalized_ref = ref.lower()
+    if not ref or normalized_ref.startswith("data:"):
+        return None
+    if normalized_ref.startswith(_HTTP_PROTOCOL_PREFIXES):
+        parsed = urlparse(ref)
+        if parsed.query or "@" in parsed.netloc:
+            return None
+        return ref
+    if _is_likely_local_path_ref(ref):
+        return None
+    return ref
+
+
+def _extract_trace_base_fields(attachment_dict: dict[str, Any]) -> dict[str, Any]:
+    """Copy stable attachment identity fields into a trace-safe payload."""
+    trace: dict[str, Any] = {}
+    for key in ("id", "kind", "source", "filename", "mime_type", "size_bytes"):
+        if key in attachment_dict:
+            trace[key] = attachment_dict[key]
+    return trace
+
+
+def _sanitize_trace_metadata(metadata: dict[str, Any]) -> dict[str, Any]:
+    """Remove unsafe metadata keys and values before trace emission."""
+    safe_metadata: dict[str, Any] = {}
+    for key, value in metadata.items():
+        if _normalize_trace_metadata_key(key) in _UNSAFE_TRACE_METADATA_KEYS:
+            continue
+        if isinstance(value, str):
+            safe_value = safe_image_data_ref(value.strip())
+            if safe_value is None:
+                continue
+            value = safe_value
+        safe_metadata[key] = value
+    return safe_metadata
+
+
+def sanitize_image_attachment_for_trace(attachment: RuntimeAttachment | dict[str, Any]) -> dict[str, Any]:
+    """Return trace-safe image metadata without local paths, signed URLs, or bytes."""
+    attachment_dict = attachment.model_dump() if isinstance(attachment, RuntimeAttachment) else dict(attachment)
+    trace = _extract_trace_base_fields(attachment_dict)
+
+    safe_ref = safe_image_data_ref(attachment_dict.get("data_ref"))
+    if safe_ref is not None:
+        trace["data_ref"] = safe_ref
+
+    metadata = attachment_dict.get("metadata")
+    if isinstance(metadata, dict):
+        safe_metadata = _sanitize_trace_metadata(metadata)
+        if safe_metadata:
+            trace["metadata"] = safe_metadata
+    return trace
+
+
+def create_input_image_attachment(
+    *,
+    attachment_id: str,
+    path: str,
+    filename: str | None = None,
+    mime_type: str | None = None,
+    size_bytes: int | None = None,
+    data_ref: str | None = None,
+    source: AttachmentSource = AttachmentSource.USER_INPUT,
+) -> RuntimeAttachment:
+    """Create a safe image attachment reference for user-provided files."""
+    resolved_filename = filename or os.path.basename(path)
+    resolved_mime_type = mime_type or str(get_mime_type_from_filename(resolved_filename))
+    resolved_size = size_bytes if size_bytes is not None else _get_file_size(path)
+    return create_runtime_attachment(
+        attachment_id=attachment_id,
+        source=source,
+        filename=resolved_filename,
+        mime_type=resolved_mime_type,
+        data_ref=data_ref or path,
+        size_bytes=resolved_size,
+    )
+
+
+def create_generated_image_attachment(
+    *,
+    artifact_name: str,
+    mime_type: str,
+    artifact_data: str,
+    source: AttachmentSource = AttachmentSource.LOCAL_TOOL_OUTPUT,
+    metadata: dict[str, Any] | None = None,
+) -> RuntimeAttachment:
+    """Create a safe image attachment reference for a generated artifact."""
+    digest = hashlib.sha256(f"{artifact_name}:{artifact_data}".encode()).hexdigest()
+    return create_runtime_attachment(
+        attachment_id=f"img_artifact_{digest[:12]}",
+        source=source,
+        filename=artifact_name,
+        mime_type=mime_type,
+        data_ref=f"artifact://sha256/{digest}",
+        size_bytes=_decoded_base64_size(artifact_data),
+        metadata=metadata,
+    )
+
+
+def create_generated_document_attachment(
+    *,
+    artifact_name: str,
+    mime_type: str,
+    artifact_data: str,
+    source: AttachmentSource = AttachmentSource.LOCAL_TOOL_OUTPUT,
+    metadata: dict[str, Any] | None = None,
+) -> RuntimeAttachment:
+    """Create a safe document attachment reference for a generated artifact."""
+    digest = hashlib.sha256(f"{artifact_name}:{artifact_data}".encode()).hexdigest()
+    return create_runtime_attachment(
+        attachment_id=f"document_artifact_{digest[:12]}",
+        source=source,
+        filename=artifact_name,
+        mime_type=mime_type,
+        kind=AttachmentKind.DOCUMENT,
+        data_ref=f"artifact://sha256/{digest}",
+        size_bytes=_decoded_base64_size(artifact_data),
+        metadata=metadata,
+    )
+
+
+def _get_file_size(path: str) -> int | None:
+    """Return file size when the local path is readable."""
+    try:
+        if os.path.isfile(path) and os.access(path, os.R_OK):
+            return os.path.getsize(path)
+    except OSError:
+        return None
+    return None
+
+
+def _decoded_base64_size(value: str) -> int | None:
+    """Return decoded byte length for base64 content without exposing bytes."""
+    try:
+        return len(base64.b64decode(value, validate=True))
+    except Exception:
+        return None
+
+
+__all__ = [
+    "AttachmentKind",
+    "AttachmentRegistry",
+    "AttachmentSource",
+    "AttachmentEntryPayload",
+    "RuntimeAttachment",
+    "MAX_IMAGE_ATTACHMENT_BYTES",
+    "SUPPORTED_IMAGE_MIME_TYPES",
+    "create_generated_image_attachment",
+    "create_generated_document_attachment",
+    "create_runtime_attachment",
+    "create_input_image_attachment",
+    "create_runtime_image_attachment",
+    "is_supported_image_mime",
+    "is_supported_document_mime",
+    "SUPPORTED_DOCUMENT_MIME_TYPES",
+    "normalize_mime_type",
+    "merge_attachment_registry",
+    "normalize_image_mime_type",
+    "safe_image_data_ref",
+    "sanitize_image_attachment_for_trace",
+]