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

aip_agents/utils/pii/pii_helper.pyi

@@ -0,0 +1,78 @@
+from _typeshed import Incomplete
+from aip_agents.utils.logger import LoggerManager as LoggerManager
+from typing import Any
+
+logger: Incomplete
+
+def normalize_enable_pii(enable_pii: Any) -> bool | None:
+    """Normalize enable_pii value from agent configuration.
+
+    Args:
+        enable_pii: Raw enable_pii value from agent configuration.
+
+    Returns:
+        The normalized enable_pii flag when explicitly set (True/False), otherwise None.
+    """
+def add_pii_mappings(left: dict[str, str] | None, right: dict[str, str] | None) -> dict[str, str]:
+    """Reducer function to merge PII mappings from multiple sources.
+
+    This is a LangGraph reducer function that merges PII mappings from:
+    - Parent agent's initial mapping
+    - Tool outputs with newly discovered PII
+    - Subagent responses with their discovered PII
+
+    Args:
+        left: Existing PII mapping (or None)
+        right: New PII mapping to merge (or None)
+
+    Returns:
+        Merged PII mapping dictionary
+
+    Note:
+        - Right (new) mappings take precedence over left (existing)
+        - Handles None/non-dict cases gracefully
+        - Preserves all unique PII tags
+        - Returns empty dict if both inputs are None/empty
+    """
+def extract_pii_mapping_from_agent_response(result: Any) -> dict[str, str] | None:
+    """Extract PII mapping from subagent response.
+
+    Used by DelegationToolManager to propagate PII mappings from subagents
+    back to parent agents.
+
+    Args:
+        result: The result returned by the delegated agent
+
+    Returns:
+        PII mapping dictionary if found, None otherwise
+
+    Note:
+        - Checks if result is a dict
+        - Extracts 'full_final_state' from result
+        - Extracts 'pii_mapping' from full_final_state
+        - Validates mapping is a non-empty dict
+        - Returns None if any step fails
+    """
+def deanonymize_final_response_content(content: str, is_final_response: bool, metadata: dict[str, Any] | None) -> str:
+    """Deanonymize final response content using PII mapping from metadata.
+
+    Args:
+        content: Final response content that may contain PII tags.
+        is_final_response: Flag indicating whether this message is a final response.
+        metadata: Optional metadata dict (or event payload containing ``metadata``) with
+            ``pii_mapping`` tag-to-value mapping.
+
+    Returns:
+        Content string with PII tags replaced by real values when applicable.
+    """
+def anonymize_final_response_content(content: str, metadata: dict[str, Any] | None) -> str:
+    """Anonymize final response content using PII mapping from metadata.
+
+    Args:
+        content: Final response content that may contain real PII values.
+        metadata: Metadata dict (or event payload containing ``metadata``) with
+            ``pii_mapping`` tag-to-value mapping.
+
+    Returns:
+        Content string with real PII values replaced by their PII tags when mapping is present.
+    """

aip_agents/utils/pii/uuid_deanonymizer_mapping.py

@@ -0,0 +1,195 @@
+"""UUID-based deanonymizer mapping utilities.
+
+This module provides a UUID-based implementation of the deanonymizer mapping,
+where anonymized entities are suffixed with UUIDs instead of sequential numbers.
+
+Authors:
+    Fachriza Adhiatma (fachriza.d.adhiatma@gdplabs.id)
+"""
+
+import logging
+import re
+import uuid
+from collections import defaultdict
+from typing import Any
+
+try:
+    from gllm_privacy.pii_detector.utils.deanonymizer_mapping import (
+        DeanonymizerMapping,
+        MappingDataType,
+    )
+except ImportError:  # pragma: no cover
+    MappingDataType = dict[str, dict[str, str]]
+
+    class DeanonymizerMapping:  # type: ignore[no-redef]
+        """Fallback deanonymizer mapping when optional dependency is missing.
+
+        This class exists only to keep the module importable when `gllm-privacy`
+        is not installed.
+        """
+
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            """Initialize the mapping.
+
+            Raises:
+                ImportError: Always raised because `gllm-privacy` is required.
+            """
+            raise ImportError("optional dependency 'gllm-privacy' is required for UUIDDeanonymizerMapping")
+
+
+from aip_agents.utils.constants import DEFAULT_PII_TAG_NAMESPACE
+
+logger = logging.getLogger(__name__)
+
+
+def format_operator_with_uuid(operator_name: str, uuid_suffix: str) -> str:
+    """Format the operator name with a UUID suffix.
+
+    Args:
+        operator_name: The operator name.
+        uuid_suffix: The UUID suffix to append.
+
+    Returns:
+        The formatted operator name with UUID suffix.
+    """
+    clean_operator_name = re.sub(r"[<>]", "", operator_name)
+    clean_operator_name = re.sub(r"_[a-f0-9-]+$", "", clean_operator_name)
+
+    if operator_name.startswith("<") and operator_name.endswith(">"):
+        return f"<{clean_operator_name}_{uuid_suffix}>"
+    else:
+        return f"{clean_operator_name}_{uuid_suffix}"
+
+
+class UUIDDeanonymizerMapping(DeanonymizerMapping):
+    """Class to store the deanonymizer mapping with UUID suffixes.
+
+    This class extends DeanonymizerMapping to use UUID suffixes instead of
+    sequential numbers for differentiating multiple entities of the same type.
+
+    Attributes:
+        mapping: The deanonymizer mapping.
+        skip_format_duplicates: Whether to skip formatting duplicated operators.
+        uuid_length: The length of the UUID suffix to use (default: 8).
+    """
+
+    def __init__(
+        self,
+        mapping: MappingDataType | None = None,
+        skip_format_duplicates: bool = False,
+        uuid_length: int = 8,
+    ) -> None:
+        """Initialize UUIDDeanonymizerMapping.
+
+        Args:
+            mapping: The deanonymizer mapping. If None, creates an empty defaultdict(dict).
+            skip_format_duplicates: Whether to skip formatting duplicated operators.
+            uuid_length: The length of the UUID suffix to use (default: 8).
+        """
+        if mapping is None:
+            mapping = defaultdict(dict)
+
+        try:
+            super().__init__(mapping=mapping, skip_format_duplicates=skip_format_duplicates)
+        except TypeError as exc:
+            logger.warning(
+                "DeanonymizerMapping init skipped due to incompatible signature; " "using local attributes only: %s",
+                exc,
+            )
+
+        self.mapping = mapping
+        self.skip_format_duplicates = skip_format_duplicates
+        self.uuid_length = uuid_length
+
+    def update(self, new_mapping: MappingDataType, use_uuid_suffix: bool | None = None) -> None:
+        """Update the deanonymizer mapping with new values using UUID suffixes.
+
+        Duplicated values will not be added. If there are multiple entities of the same type,
+        the mapping will include a UUID suffix to differentiate them. For example, if there are
+        two names in the input text, the mapping will include NAME_<uuid1> and NAME_<uuid2>.
+
+        Args:
+            new_mapping: The new mapping to be added to the existing deanonymizer mapping.
+            use_uuid_suffix: Whether to apply UUID suffixes to keys.
+                If True, keys will always be formatted with UUID suffixes.
+                If False, keys will be used as-is without UUID formatting.
+                If None, behavior falls back to the instance configuration via
+                skip_format_duplicates (preserving existing behavior).
+
+        Returns:
+            None
+        """
+        seen_values: set[str] = set()
+        format_with_uuid = self._should_format_with_uuid(use_uuid_suffix)
+
+        for entity_type, values in new_mapping.items():
+            self._update_entity_type(entity_type, values, seen_values, format_with_uuid)
+
+    def _update_entity_type(
+        self,
+        entity_type: str,
+        values: dict[str, str],
+        seen_values: set[str],
+        format_with_uuid: bool,
+    ) -> None:
+        """Update mapping entries for a single entity type.
+
+        Args:
+            entity_type: The entity category being updated.
+            values: Mapping of anonymized keys to original values for the entity.
+            seen_values: Set tracking values already processed in this update call.
+            format_with_uuid: Whether UUID formatting should be applied to keys.
+
+        Returns:
+            None
+        """
+        for key, value in values.items():
+            if not self._should_store_value(entity_type, value, seen_values):
+                continue
+
+            new_key = self._format_key(key, value, format_with_uuid)
+            self.mapping[entity_type][new_key] = value
+            seen_values.add(value)
+
+    def _should_format_with_uuid(self, use_uuid_suffix: bool | None) -> bool:
+        """Determine whether UUID suffix formatting is required.
+
+        Args:
+            use_uuid_suffix: Override controlling UUID formatting behavior.
+
+        Returns:
+            True if UUID suffixes should be applied, otherwise False.
+        """
+        if use_uuid_suffix is None:
+            return not self.skip_format_duplicates
+        return use_uuid_suffix
+
+    def _should_store_value(self, entity_type: str, value: str, seen_values: set[str]) -> bool:
+        """Check if the provided value should be stored in the mapping.
+
+        Args:
+            entity_type: Entity category being updated.
+            value: Original value corresponding to the anonymized key.
+            seen_values: Values already processed during this update call.
+
+        Returns:
+            True if the value is new for the entity type and current update scope.
+        """
+        return value not in seen_values and value not in self.mapping[entity_type].values()
+
+    def _format_key(self, key: str, value: str, format_with_uuid: bool) -> str:
+        """Return the appropriate key representation based on format flag.
+
+        Args:
+            key: Original anonymized key produced by the detector.
+            value: Original value used to derive deterministic UUID suffixes.
+            format_with_uuid: Whether UUID suffix formatting should be applied.
+
+        Returns:
+            The formatted key respecting the desired UUID formatting behavior.
+        """
+        if not format_with_uuid:
+            return key
+
+        uuid_suffix = uuid.uuid5(DEFAULT_PII_TAG_NAMESPACE, value).hex[: self.uuid_length]
+        return format_operator_with_uuid(key, uuid_suffix)

aip_agents/utils/pii/uuid_deanonymizer_mapping.pyi

@@ -0,0 +1,73 @@
+from _typeshed import Incomplete
+from aip_agents.utils.constants import DEFAULT_PII_TAG_NAMESPACE as DEFAULT_PII_TAG_NAMESPACE
+from gllm_privacy.pii_detector.utils.deanonymizer_mapping import DeanonymizerMapping, MappingDataType
+from typing import Any
+
+MappingDataType = dict[str, dict[str, str]]
+
+class DeanonymizerMapping:
+    """Fallback deanonymizer mapping when optional dependency is missing.
+
+        This class exists only to keep the module importable when `gllm-privacy`
+        is not installed.
+        """
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        """Initialize the mapping.
+
+            Raises:
+                ImportError: Always raised because `gllm-privacy` is required.
+            """
+
+logger: Incomplete
+
+def format_operator_with_uuid(operator_name: str, uuid_suffix: str) -> str:
+    """Format the operator name with a UUID suffix.
+
+    Args:
+        operator_name: The operator name.
+        uuid_suffix: The UUID suffix to append.
+
+    Returns:
+        The formatted operator name with UUID suffix.
+    """
+
+class UUIDDeanonymizerMapping(DeanonymizerMapping):
+    """Class to store the deanonymizer mapping with UUID suffixes.
+
+    This class extends DeanonymizerMapping to use UUID suffixes instead of
+    sequential numbers for differentiating multiple entities of the same type.
+
+    Attributes:
+        mapping: The deanonymizer mapping.
+        skip_format_duplicates: Whether to skip formatting duplicated operators.
+        uuid_length: The length of the UUID suffix to use (default: 8).
+    """
+    mapping: Incomplete
+    skip_format_duplicates: Incomplete
+    uuid_length: Incomplete
+    def __init__(self, mapping: MappingDataType | None = None, skip_format_duplicates: bool = False, uuid_length: int = 8) -> None:
+        """Initialize UUIDDeanonymizerMapping.
+
+        Args:
+            mapping: The deanonymizer mapping. If None, creates an empty defaultdict(dict).
+            skip_format_duplicates: Whether to skip formatting duplicated operators.
+            uuid_length: The length of the UUID suffix to use (default: 8).
+        """
+    def update(self, new_mapping: MappingDataType, use_uuid_suffix: bool | None = None) -> None:
+        """Update the deanonymizer mapping with new values using UUID suffixes.
+
+        Duplicated values will not be added. If there are multiple entities of the same type,
+        the mapping will include a UUID suffix to differentiate them. For example, if there are
+        two names in the input text, the mapping will include NAME_<uuid1> and NAME_<uuid2>.
+
+        Args:
+            new_mapping: The new mapping to be added to the existing deanonymizer mapping.
+            use_uuid_suffix: Whether to apply UUID suffixes to keys.
+                If True, keys will always be formatted with UUID suffixes.
+                If False, keys will be used as-is without UUID formatting.
+                If None, behavior falls back to the instance configuration via
+                skip_format_duplicates (preserving existing behavior).
+
+        Returns:
+            None
+        """

aip_agents/utils/reference_helper.py

@@ -0,0 +1,273 @@
+"""Reference helper utilities for handling and validating reference data.
+
+Authors:
+    Putu Ravindra Wiguna (putu.r.wiguna@gdplabs.id)
+"""
+
+from typing import Any
+
+from gllm_core.schema import Chunk
+from langgraph.types import Command
+
+from aip_agents.utils.logger import get_logger
+from aip_agents.utils.metadata_helper import MetadataFieldKeys
+
+logger = get_logger(__name__)
+
+
+# Tool attribute constants for reference extraction
+SAVE_OUTPUT_HISTORY_ATTR = "save_output_history"
+FORMAT_AGENT_REFERENCE = "_format_agent_reference"
+
+
+def extract_references_from_tool(tool: Any, tool_output: Any) -> list[Chunk]:
+    """Extract reference data from tools that support it.
+
+    Extraction priority:
+    1. Direct tool references via _format_agent_reference (preferred)
+    2. Command.update references (fallback for delegation tools)
+
+    Args:
+        tool: The tool instance to extract references from
+        tool_output: The output from the tool execution
+    Returns:
+        List of deduplicated Chunk objects containing reference data
+
+    Note:
+        - Never raises exceptions; logs warnings for issues
+        - Direct tool references take precedence over Command references
+    """
+    if getattr(tool, SAVE_OUTPUT_HISTORY_ATTR, False):
+        logger.debug(
+            "Ignoring save_output_history for %s during reference extraction.",
+            tool.__class__.__name__,
+        )
+
+    # Prefer direct tool references
+    direct_references = _extract_direct_tool_references(tool, tool_output)
+    if direct_references:
+        return validate_references(direct_references)
+
+    if isinstance(tool_output, Command):
+        command_references = extract_references_from_command_update(tool_output)
+        if command_references:
+            return validate_references(command_references)
+
+    return []
+
+
+def _extract_direct_tool_references(tool: Any, tool_output: Any) -> list[Any]:
+    """Extract references directly from tool hooks when available.
+
+    Args:
+        tool: The tool instance to extract references from
+        tool_output: The output from the tool execution
+
+    Returns:
+        List of reference data (not yet validated/deduplicated)
+    """
+    if not hasattr(tool, FORMAT_AGENT_REFERENCE):
+        return []
+
+    try:
+        reference_data = tool._format_agent_reference(tool_output)
+    except (AttributeError, ValueError, TypeError) as e:
+        logger.warning(f"Failed to extract references from {tool.__class__.__name__}: {e}")
+        return []
+    except Exception as e:
+        logger.error(
+            f"Unexpected error extracting references from {tool.__class__.__name__}: {e}",
+            exc_info=True,
+        )
+        return []
+
+    return _normalize_reference_data(reference_data, tool)
+
+
+def _normalize_reference_data(reference_data: Any, tool: Any | None = None) -> list[Any]:
+    """Normalize reference payloads into a list for downstream processing.
+
+    Args:
+        reference_data: Arbitrary payload returned by a tool or command.
+        tool: Optional tool instance used for contextual logging when payload type is invalid.
+
+    Returns:
+        A list of reference objects ready for validation. Returns an empty list when
+        the payload is ``None`` or cannot be interpreted as a reference collection.
+
+    Notes:
+        - Accepts `Chunk` instances and wraps them in a list.
+        - Pass-through for list inputs to preserve existing semantics.
+        - Emits a warning when encountering unsupported payload types, including
+          the originating tool class when available.
+    """
+    if reference_data is None:
+        return []
+    if isinstance(reference_data, Chunk):
+        return [reference_data]
+    if isinstance(reference_data, list):
+        return reference_data
+
+    if tool is not None:
+        logger.warning(
+            "Invalid reference data format from %s: %s",
+            tool.__class__.__name__,
+            reference_data,
+        )
+    else:
+        logger.warning("Invalid reference data format: %s", reference_data)
+
+    return []
+
+
+def extract_references_from_command_update(command: Command) -> list[Any]:
+    """Extract references from a Command object's update dictionary.
+
+    Args:
+        command: A Command object potentially containing references in its update dict.
+
+    Returns:
+        List of reference data (not yet validated/deduplicated)
+
+    Note:
+        - Never raises exceptions; logs warnings for malformed data
+        - Skips non-Chunk items with warning log
+    """
+    update = getattr(command, "update", None)
+    if not isinstance(update, dict):
+        return []
+
+    raw_refs = update.get(MetadataFieldKeys.REFERENCES)
+    if not isinstance(raw_refs, list):
+        return []
+
+    valid_refs = []
+    for ref in raw_refs:
+        if isinstance(ref, Chunk):
+            valid_refs.append(ref)
+        else:
+            logger.warning(
+                "Ignoring non-Chunk reference payload from delegation command: %s (type: %s)",
+                ref,
+                type(ref).__name__,
+            )
+
+    return valid_refs
+
+
+def validate_references(references: list[Any]) -> list[Chunk]:
+    """Validate and deduplicate reference data.
+
+    Args:
+        references: List of reference data (expected to be Chunk objects).
+
+    Returns:
+        List of validated, deduplicated Chunk objects by content.
+    """
+    if not references:
+        return []
+
+    validated = []
+    seen_content = set()
+
+    for ref in references:
+        # Skip non-Chunk objects
+        if not isinstance(ref, Chunk):
+            logger.debug(f"Skipping non-Chunk reference: {type(ref)}")
+            continue
+
+        # Skip invalid content/metadata
+        if not isinstance(ref.content, str | bytes) or not isinstance(ref.metadata, dict):
+            logger.debug(f"Skipping reference with invalid content/metadata: {type(ref.content)}, {type(ref.metadata)}")
+            continue
+
+        # Deduplicate by content
+        if ref.content not in seen_content:
+            seen_content.add(ref.content)
+            validated.append(ref)
+
+    return validated
+
+
+def serialize_references_for_metadata(references: list[Any]) -> list[dict[str, Any]]:
+    """Serialize references for inclusion in A2A metadata.
+
+    Args:
+        references: List of reference objects (typically Chunk objects).
+
+    Returns:
+        List of serialized reference dictionaries.
+    """
+    serialized_refs = []
+    for ref in references:
+        if hasattr(ref, "model_dump"):
+            # For Pydantic models like Chunk
+            serialized_refs.append(ref.model_dump(mode="python"))
+        elif isinstance(ref, dict):
+            # Plain dicts pass through unchanged
+            serialized_refs.append(ref)
+        elif hasattr(ref, "__dict__"):
+            # For regular objects with attributes
+            serialized_refs.append(ref.__dict__)
+        else:
+            # Fallback for other types
+            serialized_refs.append(str(ref))
+    return serialized_refs
+
+
+def add_references_chunks(left: list[Chunk], right: list[Chunk]) -> list[Chunk]:
+    """Reducer function to accumulate reference data from multiple tool calls.
+
+    This is a LangGraph reducer function that should be forgiving and handle
+    edge cases gracefully. Non-Chunk items are filtered out.
+
+    Args:
+        left: Existing list of reference chunks (or None/non-list)
+        right: New list of reference chunks to add (or None/non-list)
+
+    Returns:
+        Combined list of valid Chunk objects
+    """
+    # Handle None/non-list cases by converting to empty lists
+    left_list = left if isinstance(left, list) else []
+    right_list = right if isinstance(right, list) else []
+
+    # Filter to only include valid Chunk objects
+    valid_left = [item for item in left_list if isinstance(item, Chunk)]
+    valid_right = [item for item in right_list if isinstance(item, Chunk)]
+
+    # Handle empty cases efficiently
+    if not valid_left:
+        return valid_right
+    if not valid_right:
+        return valid_left
+
+    return valid_left + valid_right
+
+
+def extract_references_from_agent_response(result: Any) -> list[dict[str, Any]] | None:
+    """Extract references from agent response for delegation tools.
+
+    Args:
+        result: The result returned by the delegated agent.
+
+    Returns:
+        List of reference chunks if found, None otherwise.
+    """
+    if not isinstance(result, dict):
+        return None
+
+    full_state = result.get("full_final_state", {})
+    if not isinstance(full_state, dict):
+        return None
+
+    references = full_state.get("references")
+    if not isinstance(references, list) or not references:
+        return None
+
+    try:
+        validated_refs = validate_references(references)
+        return validated_refs if validated_refs else None
+    except Exception as e:
+        logger.warning(f"ExtractRefFromAgentResponse: Error validating references: {e}")
+        return None

aip_agents/utils/reference_helper.pyi

@@ -0,0 +1,81 @@
+from _typeshed import Incomplete
+from aip_agents.utils.logger import get_logger as get_logger
+from aip_agents.utils.metadata_helper import MetadataFieldKeys as MetadataFieldKeys
+from gllm_core.schema import Chunk
+from langgraph.types import Command
+from typing import Any
+
+logger: Incomplete
+SAVE_OUTPUT_HISTORY_ATTR: str
+FORMAT_AGENT_REFERENCE: str
+
+def extract_references_from_tool(tool: Any, tool_output: Any) -> list[Chunk]:
+    """Extract reference data from tools that support it.
+
+    Extraction priority:
+    1. Direct tool references via _format_agent_reference (preferred)
+    2. Command.update references (fallback for delegation tools)
+
+    Args:
+        tool: The tool instance to extract references from
+        tool_output: The output from the tool execution
+    Returns:
+        List of deduplicated Chunk objects containing reference data
+
+    Note:
+        - Never raises exceptions; logs warnings for issues
+        - Direct tool references take precedence over Command references
+    """
+def extract_references_from_command_update(command: Command) -> list[Any]:
+    """Extract references from a Command object's update dictionary.
+
+    Args:
+        command: A Command object potentially containing references in its update dict.
+
+    Returns:
+        List of reference data (not yet validated/deduplicated)
+
+    Note:
+        - Never raises exceptions; logs warnings for malformed data
+        - Skips non-Chunk items with warning log
+    """
+def validate_references(references: list[Any]) -> list[Chunk]:
+    """Validate and deduplicate reference data.
+
+    Args:
+        references: List of reference data (expected to be Chunk objects).
+
+    Returns:
+        List of validated, deduplicated Chunk objects by content.
+    """
+def serialize_references_for_metadata(references: list[Any]) -> list[dict[str, Any]]:
+    """Serialize references for inclusion in A2A metadata.
+
+    Args:
+        references: List of reference objects (typically Chunk objects).
+
+    Returns:
+        List of serialized reference dictionaries.
+    """
+def add_references_chunks(left: list[Chunk], right: list[Chunk]) -> list[Chunk]:
+    """Reducer function to accumulate reference data from multiple tool calls.
+
+    This is a LangGraph reducer function that should be forgiving and handle
+    edge cases gracefully. Non-Chunk items are filtered out.
+
+    Args:
+        left: Existing list of reference chunks (or None/non-list)
+        right: New list of reference chunks to add (or None/non-list)
+
+    Returns:
+        Combined list of valid Chunk objects
+    """
+def extract_references_from_agent_response(result: Any) -> list[dict[str, Any]] | None:
+    """Extract references from agent response for delegation tools.
+
+    Args:
+        result: The result returned by the delegated agent.
+
+    Returns:
+        List of reference chunks if found, None otherwise.
+    """