agentrun-sdk 0.1.2__py3-none-any.whl

agentrun_sdk/types/guardrails.py

@@ -0,0 +1,254 @@
+"""Guardrail-related type definitions for the SDK.
+
+These types are modeled after the Bedrock API.
+
+- Bedrock docs: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Types_Amazon_Bedrock_Runtime.html
+"""
+
+from typing import Dict, List, Literal, Optional
+
+from typing_extensions import TypedDict
+
+
+class GuardrailConfig(TypedDict, total=False):
+    """Configuration for content filtering guardrails.
+
+    Attributes:
+        guardrailIdentifier: Unique identifier for the guardrail.
+        guardrailVersion: Version of the guardrail to apply.
+        streamProcessingMode: Processing mode.
+        trace: The trace behavior for the guardrail.
+    """
+
+    guardrailIdentifier: str
+    guardrailVersion: str
+    streamProcessingMode: Optional[Literal["sync", "async"]]
+    trace: Literal["enabled", "disabled"]
+
+
+class Topic(TypedDict):
+    """Information about a topic guardrail.
+
+    Attributes:
+        action: The action the guardrail should take when it intervenes on a topic.
+        name: The name for the guardrail.
+        type: The type behavior that the guardrail should perform when the model detects the topic.
+    """
+
+    action: Literal["BLOCKED"]
+    name: str
+    type: Literal["DENY"]
+
+
+class TopicPolicy(TypedDict):
+    """A behavior assessment of a topic policy.
+
+    Attributes:
+        topics: The topics in the assessment.
+    """
+
+    topics: List[Topic]
+
+
+class ContentFilter(TypedDict):
+    """The content filter for a guardrail.
+
+    Attributes:
+        action: Action to take when content is detected.
+        confidence: Confidence level of the detection.
+        type: The type of content to filter.
+    """
+
+    action: Literal["BLOCKED"]
+    confidence: Literal["NONE", "LOW", "MEDIUM", "HIGH"]
+    type: Literal["INSULTS", "HATE", "SEXUAL", "VIOLENCE", "MISCONDUCT", "PROMPT_ATTACK"]
+
+
+class ContentPolicy(TypedDict):
+    """An assessment of a content policy for a guardrail.
+
+    Attributes:
+        filters: List of content filters to apply.
+    """
+
+    filters: List[ContentFilter]
+
+
+class CustomWord(TypedDict):
+    """Definition of a custom word to be filtered.
+
+    Attributes:
+        action: Action to take when the word is detected.
+        match: The word or phrase to match.
+    """
+
+    action: Literal["BLOCKED"]
+    match: str
+
+
+class ManagedWord(TypedDict):
+    """Definition of a managed word to be filtered.
+
+    Attributes:
+        action: Action to take when the word is detected.
+        match: The word or phrase to match.
+        type: Type of the word.
+    """
+
+    action: Literal["BLOCKED"]
+    match: str
+    type: Literal["PROFANITY"]
+
+
+class WordPolicy(TypedDict):
+    """The word policy assessment.
+
+    Attributes:
+        customWords: List of custom words to filter.
+        managedWordLists: List of managed word lists to filter.
+    """
+
+    customWords: List[CustomWord]
+    managedWordLists: List[ManagedWord]
+
+
+class PIIEntity(TypedDict):
+    """Definition of a Personally Identifiable Information (PII) entity to be filtered.
+
+    Attributes:
+        action: Action to take when PII is detected.
+        match: The specific PII instance to match.
+        type: The type of PII to detect.
+    """
+
+    action: Literal["ANONYMIZED", "BLOCKED"]
+    match: str
+    type: Literal[
+        "ADDRESS",
+        "AGE",
+        "AWS_ACCESS_KEY",
+        "AWS_SECRET_KEY",
+        "CA_HEALTH_NUMBER",
+        "CA_SOCIAL_INSURANCE_NUMBER",
+        "CREDIT_DEBIT_CARD_CVV",
+        "CREDIT_DEBIT_CARD_EXPIRY",
+        "CREDIT_DEBIT_CARD_NUMBER",
+        "DRIVER_ID",
+        "EMAIL",
+        "INTERNATIONAL_BANK_ACCOUNT_NUMBER",
+        "IP_ADDRESS",
+        "LICENSE_PLATE",
+        "MAC_ADDRESS",
+        "NAME",
+        "PASSWORD",
+        "PHONE",
+        "PIN",
+        "SWIFT_CODE",
+        "UK_NATIONAL_HEALTH_SERVICE_NUMBER",
+        "UK_NATIONAL_INSURANCE_NUMBER",
+        "UK_UNIQUE_TAXPAYER_REFERENCE_NUMBER",
+        "URL",
+        "USERNAME",
+        "US_BANK_ACCOUNT_NUMBER",
+        "US_BANK_ROUTING_NUMBER",
+        "US_INDIVIDUAL_TAX_IDENTIFICATION_NUMBER",
+        "US_PASSPORT_NUMBER",
+        "US_SOCIAL_SECURITY_NUMBER",
+        "VEHICLE_IDENTIFICATION_NUMBER",
+    ]
+
+
+class Regex(TypedDict):
+    """Definition of a custom regex pattern for filtering sensitive information.
+
+    Attributes:
+        action: Action to take when the pattern is matched.
+        match: The regex filter match.
+        name: Name of the regex pattern for identification.
+        regex: The regex query.
+    """
+
+    action: Literal["ANONYMIZED", "BLOCKED"]
+    match: str
+    name: str
+    regex: str
+
+
+class SensitiveInformationPolicy(TypedDict):
+    """Policy defining sensitive information filtering rules.
+
+    Attributes:
+        piiEntities: List of Personally Identifiable Information (PII) entities to detect and handle.
+        regexes: The regex queries in the assessment.
+    """
+
+    piiEntities: List[PIIEntity]
+    regexes: List[Regex]
+
+
+class ContextualGroundingFilter(TypedDict):
+    """Filter for ensuring responses are grounded in provided context.
+
+    Attributes:
+        action: Action to take when the threshold is not met.
+        score: The score generated by contextual grounding filter (range [0, 1]).
+        threshold: Threshold used by contextual grounding filter to determine whether the content is grounded or not.
+        type: The contextual grounding filter type.
+    """
+
+    action: Literal["BLOCKED", "NONE"]
+    score: float
+    threshold: float
+    type: Literal["GROUNDING", "RELEVANCE"]
+
+
+class ContextualGroundingPolicy(TypedDict):
+    """The policy assessment details for the guardrails contextual grounding filter.
+
+    Attributes:
+        filters: The filter details for the guardrails contextual grounding filter.
+    """
+
+    filters: List[ContextualGroundingFilter]
+
+
+class GuardrailAssessment(TypedDict):
+    """A behavior assessment of the guardrail policies used in a call to the Converse API.
+
+    Attributes:
+        contentPolicy: The content policy.
+        contextualGroundingPolicy: The contextual grounding policy used for the guardrail assessment.
+        sensitiveInformationPolicy: The sensitive information policy.
+        topicPolicy: The topic policy.
+        wordPolicy: The word policy.
+    """
+
+    contentPolicy: ContentPolicy
+    contextualGroundingPolicy: ContextualGroundingPolicy
+    sensitiveInformationPolicy: SensitiveInformationPolicy
+    topicPolicy: TopicPolicy
+    wordPolicy: WordPolicy
+
+
+class GuardrailTrace(TypedDict):
+    """Trace information from guardrail processing.
+
+    Attributes:
+        inputAssessment: Assessment of input content against guardrail policies, keyed by input identifier.
+        modelOutput: The original output from the model before guardrail processing.
+        outputAssessments: Assessments of output content against guardrail policies, keyed by output identifier.
+    """
+
+    inputAssessment: Dict[str, GuardrailAssessment]
+    modelOutput: List[str]
+    outputAssessments: Dict[str, List[GuardrailAssessment]]
+
+
+class Trace(TypedDict):
+    """A Top level guardrail trace object.
+
+    Attributes:
+        guardrail: Trace information from guardrail processing.
+    """
+
+    guardrail: GuardrailTrace

agentrun_sdk/types/media.py

@@ -0,0 +1,89 @@
+"""Media-related type definitions for the SDK.
+
+These types are modeled after the Bedrock API.
+
+- Bedrock docs: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Types_Amazon_Bedrock_Runtime.html
+"""
+
+from typing import Literal
+
+from typing_extensions import TypedDict
+
+DocumentFormat = Literal["pdf", "csv", "doc", "docx", "xls", "xlsx", "html", "txt", "md"]
+"""Supported document formats."""
+
+
+class DocumentSource(TypedDict):
+    """Contains the content of a document.
+
+    Attributes:
+        bytes: The binary content of the document.
+    """
+
+    bytes: bytes
+
+
+class DocumentContent(TypedDict):
+    """A document to include in a message.
+
+    Attributes:
+        format: The format of the document (e.g., "pdf", "txt").
+        name: The name of the document.
+        source: The source containing the document's binary content.
+    """
+
+    format: Literal["pdf", "csv", "doc", "docx", "xls", "xlsx", "html", "txt", "md"]
+    name: str
+    source: DocumentSource
+
+
+ImageFormat = Literal["png", "jpeg", "gif", "webp"]
+"""Supported image formats."""
+
+
+class ImageSource(TypedDict):
+    """Contains the content of an image.
+
+    Attributes:
+        bytes: The binary content of the image.
+    """
+
+    bytes: bytes
+
+
+class ImageContent(TypedDict):
+    """An image to include in a message.
+
+    Attributes:
+        format: The format of the image (e.g., "png", "jpeg").
+        source: The source containing the image's binary content.
+    """
+
+    format: ImageFormat
+    source: ImageSource
+
+
+VideoFormat = Literal["flv", "mkv", "mov", "mpeg", "mpg", "mp4", "three_gp", "webm", "wmv"]
+"""Supported video formats."""
+
+
+class VideoSource(TypedDict):
+    """Contains the content of a video.
+
+    Attributes:
+        bytes: The binary content of the video.
+    """
+
+    bytes: bytes
+
+
+class VideoContent(TypedDict):
+    """A video to include in a message.
+
+    Attributes:
+        format: The format of the video (e.g., "mp4", "avi").
+        source: The source containing the video's binary content.
+    """
+
+    format: VideoFormat
+    source: VideoSource

agentrun_sdk/types/session.py

@@ -0,0 +1,152 @@
+"""Data models for session management."""
+
+import base64
+import inspect
+from dataclasses import asdict, dataclass, field
+from datetime import datetime, timezone
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+from .content import Message
+
+if TYPE_CHECKING:
+    from ..agent.agent import Agent
+
+
+class SessionType(str, Enum):
+    """Enumeration of session types.
+
+    As sessions are expanded to support new usecases like multi-agent patterns,
+    new types will be added here.
+    """
+
+    AGENT = "AGENT"
+
+
+def encode_bytes_values(obj: Any) -> Any:
+    """Recursively encode any bytes values in an object to base64.
+
+    Handles dictionaries, lists, and nested structures.
+    """
+    if isinstance(obj, bytes):
+        return {"__bytes_encoded__": True, "data": base64.b64encode(obj).decode()}
+    elif isinstance(obj, dict):
+        return {k: encode_bytes_values(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [encode_bytes_values(item) for item in obj]
+    else:
+        return obj
+
+
+def decode_bytes_values(obj: Any) -> Any:
+    """Recursively decode any base64-encoded bytes values in an object.
+
+    Handles dictionaries, lists, and nested structures.
+    """
+    if isinstance(obj, dict):
+        if obj.get("__bytes_encoded__") is True and "data" in obj:
+            return base64.b64decode(obj["data"])
+        return {k: decode_bytes_values(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [decode_bytes_values(item) for item in obj]
+    else:
+        return obj
+
+
+@dataclass
+class SessionMessage:
+    """Message within a SessionAgent.
+
+    Attributes:
+        message: Message content
+        message_id: Index of the message in the conversation history
+        redact_message: If the original message is redacted, this is the new content to use
+        created_at: ISO format timestamp for when this message was created
+        updated_at: ISO format timestamp for when this message was last updated
+    """
+
+    message: Message
+    message_id: int
+    redact_message: Optional[Message] = None
+    created_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+    updated_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+
+    @classmethod
+    def from_message(cls, message: Message, index: int) -> "SessionMessage":
+        """Convert from a Message, base64 encoding bytes values."""
+        return cls(
+            message=message,
+            message_id=index,
+            created_at=datetime.now(timezone.utc).isoformat(),
+            updated_at=datetime.now(timezone.utc).isoformat(),
+        )
+
+    def to_message(self) -> Message:
+        """Convert SessionMessage back to a Message, decoding any bytes values.
+
+        If the message was redacted, return the redact content instead.
+        """
+        if self.redact_message is not None:
+            return self.redact_message
+        else:
+            return self.message
+
+    @classmethod
+    def from_dict(cls, env: dict[str, Any]) -> "SessionMessage":
+        """Initialize a SessionMessage from a dictionary, ignoring keys that are not class parameters."""
+        extracted_relevant_parameters = {k: v for k, v in env.items() if k in inspect.signature(cls).parameters}
+        return cls(**decode_bytes_values(extracted_relevant_parameters))
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert the SessionMessage to a dictionary representation."""
+        return encode_bytes_values(asdict(self))  # type: ignore
+
+
+@dataclass
+class SessionAgent:
+    """Agent that belongs to a Session."""
+
+    agent_id: str
+    state: Dict[str, Any]
+    conversation_manager_state: Dict[str, Any]
+    created_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+    updated_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+
+    @classmethod
+    def from_agent(cls, agent: "Agent") -> "SessionAgent":
+        """Convert an Agent to a SessionAgent."""
+        if agent.agent_id is None:
+            raise ValueError("agent_id needs to be defined.")
+        return cls(
+            agent_id=agent.agent_id,
+            conversation_manager_state=agent.conversation_manager.get_state(),
+            state=agent.state.get(),
+        )
+
+    @classmethod
+    def from_dict(cls, env: dict[str, Any]) -> "SessionAgent":
+        """Initialize a SessionAgent from a dictionary, ignoring keys that are not class parameters."""
+        return cls(**{k: v for k, v in env.items() if k in inspect.signature(cls).parameters})
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert the SessionAgent to a dictionary representation."""
+        return asdict(self)
+
+
+@dataclass
+class Session:
+    """Session data model."""
+
+    session_id: str
+    session_type: SessionType
+    created_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+    updated_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+
+    @classmethod
+    def from_dict(cls, env: dict[str, Any]) -> "Session":
+        """Initialize a Session from a dictionary, ignoring keys that are not class parameters."""
+        return cls(**{k: v for k, v in env.items() if k in inspect.signature(cls).parameters})
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert the Session to a dictionary representation."""
+        return asdict(self)

agentrun_sdk/types/streaming.py

@@ -0,0 +1,201 @@
+"""Streaming-related type definitions for the SDK.
+
+These types are modeled after the Bedrock API.
+
+- Bedrock docs: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Types_Amazon_Bedrock_Runtime.html
+"""
+
+from typing import Optional, Union
+
+from typing_extensions import TypedDict
+
+from .content import ContentBlockStart, Role
+from .event_loop import Metrics, StopReason, Usage
+from .guardrails import Trace
+
+
+class MessageStartEvent(TypedDict):
+    """Event signaling the start of a message in a streaming response.
+
+    Attributes:
+        role: The role of the message sender (e.g., "assistant", "user").
+    """
+
+    role: Role
+
+
+class ContentBlockStartEvent(TypedDict, total=False):
+    """Event signaling the start of a content block in a streaming response.
+
+    Attributes:
+        contentBlockIndex: Index of the content block within the message.
+            This is optional to accommodate different model providers.
+        start: Information about the content block being started.
+    """
+
+    contentBlockIndex: Optional[int]
+    start: ContentBlockStart
+
+
+class ContentBlockDeltaText(TypedDict):
+    """Text content delta in a streaming response.
+
+    Attributes:
+        text: The text fragment being streamed.
+    """
+
+    text: str
+
+
+class ContentBlockDeltaToolUse(TypedDict):
+    """Tool use input delta in a streaming response.
+
+    Attributes:
+        input: The tool input fragment being streamed.
+    """
+
+    input: str
+
+
+class ReasoningContentBlockDelta(TypedDict, total=False):
+    """Delta for reasoning content block in a streaming response.
+
+    Attributes:
+        redactedContent: The content in the reasoning that was encrypted by the model provider for safety reasons.
+        signature: A token that verifies that the reasoning text was generated by the model.
+        text: The reasoning that the model used to return the output.
+    """
+
+    redactedContent: Optional[bytes]
+    signature: Optional[str]
+    text: Optional[str]
+
+
+class ContentBlockDelta(TypedDict, total=False):
+    """A block of content in a streaming response.
+
+    Attributes:
+        reasoningContent: Contains content regarding the reasoning that is carried out by the model.
+        text: Text fragment being streamed.
+        toolUse: Tool use input fragment being streamed.
+    """
+
+    reasoningContent: ReasoningContentBlockDelta
+    text: str
+    toolUse: ContentBlockDeltaToolUse
+
+
+class ContentBlockDeltaEvent(TypedDict, total=False):
+    """Event containing a delta update for a content block in a streaming response.
+
+    Attributes:
+        contentBlockIndex: Index of the content block within the message.
+            This is optional to accommodate different model providers.
+        delta: The incremental content update for the content block.
+    """
+
+    contentBlockIndex: Optional[int]
+    delta: ContentBlockDelta
+
+
+class ContentBlockStopEvent(TypedDict, total=False):
+    """Event signaling the end of a content block in a streaming response.
+
+    Attributes:
+        contentBlockIndex: Index of the content block within the message.
+            This is optional to accommodate different model providers.
+    """
+
+    contentBlockIndex: Optional[int]
+
+
+class MessageStopEvent(TypedDict, total=False):
+    """Event signaling the end of a message in a streaming response.
+
+    Attributes:
+        additionalModelResponseFields: Additional fields to include in model response.
+        stopReason: The reason why the model stopped generating content.
+    """
+
+    additionalModelResponseFields: Optional[Union[dict, list, int, float, str, bool, None]]
+    stopReason: StopReason
+
+
+class MetadataEvent(TypedDict, total=False):
+    """Event containing metadata about the streaming response.
+
+    Attributes:
+        metrics: Performance metrics related to the model invocation.
+        trace: Trace information for debugging and monitoring.
+        usage: Resource usage information for the model invocation.
+    """
+
+    metrics: Metrics
+    trace: Optional[Trace]
+    usage: Usage
+
+
+class ExceptionEvent(TypedDict):
+    """Base event for exceptions in a streaming response.
+
+    Attributes:
+        message: The error message describing what went wrong.
+    """
+
+    message: str
+
+
+class ModelStreamErrorEvent(ExceptionEvent):
+    """Event for model streaming errors.
+
+    Attributes:
+        originalMessage: The original error message from the model provider.
+        originalStatusCode: The HTTP status code returned by the model provider.
+    """
+
+    originalMessage: str
+    originalStatusCode: int
+
+
+class RedactContentEvent(TypedDict, total=False):
+    """Event for redacting content.
+
+    Attributes:
+        redactUserContentMessage: The string to overwrite the users input with.
+        redactAssistantContentMessage: The string to overwrite the assistants output with.
+
+    """
+
+    redactUserContentMessage: Optional[str]
+    redactAssistantContentMessage: Optional[str]
+
+
+class StreamEvent(TypedDict, total=False):
+    """The messages output stream.
+
+    Attributes:
+        contentBlockDelta: Delta content for a content block.
+        contentBlockStart: Start of a content block.
+        contentBlockStop: End of a content block.
+        internalServerException: Internal server error information.
+        messageStart: Start of a message.
+        messageStop: End of a message.
+        metadata: Metadata about the streaming response.
+        modelStreamErrorException: Model streaming error information.
+        serviceUnavailableException: Service unavailable error information.
+        throttlingException: Throttling error information.
+        validationException: Validation error information.
+    """
+
+    contentBlockDelta: ContentBlockDeltaEvent
+    contentBlockStart: ContentBlockStartEvent
+    contentBlockStop: ContentBlockStopEvent
+    internalServerException: ExceptionEvent
+    messageStart: MessageStartEvent
+    messageStop: MessageStopEvent
+    metadata: MetadataEvent
+    redactContent: RedactContentEvent
+    modelStreamErrorException: ModelStreamErrorEvent
+    serviceUnavailableException: ExceptionEvent
+    throttlingException: ExceptionEvent
+    validationException: ExceptionEvent