basion-agent 0.4.0__py3-none-any.whl

basion_agent/memory_client.py

@@ -0,0 +1,155 @@
+"""HTTP client for AI Memory service."""
+
+import logging
+from dataclasses import dataclass, field
+from typing import Dict, Any, Optional, List
+from urllib.parse import urlencode
+import aiohttp
+
+from .exceptions import APIException
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MemoryMessage:
+    """A message from ai-memory storage.
+
+    Exactly matches the ai-memory API MessageStored schema.
+    """
+
+    id: str
+    role: str  # "user", "assistant", "system"
+    content: str
+    tenant_id: Optional[str] = None  # Agent scope/namespace
+    user_id: Optional[str] = None
+    conversation_id: Optional[str] = None
+    metadata: Optional[Dict[str, Any]] = None
+    timestamp: Optional[str] = None
+    sequence_number: Optional[int] = None
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "MemoryMessage":
+        """Create from ai-memory API response."""
+        return cls(
+            id=data.get("id", ""),
+            role=data.get("role", ""),
+            content=data.get("content", ""),
+            tenant_id=data.get("tenant_id"),
+            user_id=data.get("user_id"),
+            conversation_id=data.get("conversation_id"),
+            metadata=data.get("metadata"),
+            timestamp=data.get("timestamp"),
+            sequence_number=data.get("sequence_number"),
+        )
+
+
+@dataclass
+class MemorySearchResult:
+    """A memory search result with similarity score."""
+
+    message: MemoryMessage
+    score: float
+    context: List[MemoryMessage] = field(default_factory=list)
+
+
+@dataclass
+class UserSummary:
+    """User conversation summary."""
+
+    text: str
+    last_updated: str
+    message_count: int
+    version: int
+
+
+class MemoryClient:
+    """Low-level async HTTP client for AI Memory service."""
+
+    def __init__(self, base_url: str):
+        self.base_url = base_url.rstrip("/")
+        self._session: Optional[aiohttp.ClientSession] = None
+
+    async def _get_session(self) -> aiohttp.ClientSession:
+        if self._session is None or self._session.closed:
+            timeout = aiohttp.ClientTimeout(total=30.0)
+            self._session = aiohttp.ClientSession(timeout=timeout)
+        return self._session
+
+    async def close(self):
+        if self._session and not self._session.closed:
+            await self._session.close()
+
+    async def search(
+        self,
+        query: str,
+        user_id: Optional[str] = None,
+        conversation_id: Optional[str] = None,
+        tenant_id: Optional[str] = None,
+        limit: int = 10,
+        min_similarity: int = 70,
+        context_messages: int = 0,
+    ) -> List[MemorySearchResult]:
+        """Search long-term memory with semantic search."""
+        params = {
+            "q": query,
+            "limit": limit,
+            "min_similarity": min_similarity,
+            "context_messages": context_messages,
+        }
+        if user_id:
+            params["user_id"] = user_id
+        if conversation_id:
+            params["conversation_id"] = conversation_id
+        if tenant_id:
+            params["tenant_id"] = tenant_id
+
+        url = f"{self.base_url}/conversation/search?{urlencode(params)}"
+        session = await self._get_session()
+        try:
+            async with session.get(url) as response:
+                if response.status >= 400:
+                    error_text = await response.text()
+                    raise APIException(
+                        f"Memory search failed: {response.status} - {error_text}"
+                    )
+                data = await response.json()
+
+            results = []
+            for item in data.get("results", []):
+                msg = MemoryMessage.from_dict(item["message"])
+                context = [MemoryMessage.from_dict(c) for c in item.get("context", [])]
+                results.append(
+                    MemorySearchResult(message=msg, score=item["score"], context=context)
+                )
+            return results
+        except aiohttp.ClientError as e:
+            raise APIException(f"Memory search failed: {e}")
+
+    async def get_user_summary(self, user_id: str) -> Optional[UserSummary]:
+        """Get user's conversation summary."""
+        url = f"{self.base_url}/conversation/user/{user_id}/summary"
+        session = await self._get_session()
+        try:
+            async with session.get(url) as response:
+                if response.status == 404:
+                    return None
+                if response.status >= 400:
+                    error_text = await response.text()
+                    raise APIException(
+                        f"Summary fetch failed: {response.status} - {error_text}"
+                    )
+                data = await response.json()
+
+            if not data.get("has_summary") or not data.get("summary"):
+                return None
+
+            s = data["summary"]
+            return UserSummary(
+                text=s["text"],
+                last_updated=s["last_updated"],
+                message_count=s["message_count"],
+                version=s.get("version", 1),
+            )
+        except aiohttp.ClientError as e:
+            raise APIException(f"Summary fetch failed: {e}")

basion_agent/message.py

@@ -0,0 +1,333 @@
+"""Message data structure."""
+
+import asyncio
+import json
+from dataclasses import dataclass, field
+from io import BytesIO
+from typing import Dict, Any, Optional, List, TYPE_CHECKING
+
+from .conversation import Conversation
+from .conversation_client import ConversationClient
+
+if TYPE_CHECKING:
+    from .memory import Memory
+    from .memory_client import MemoryClient
+    from .attachment_client import AttachmentClient, AttachmentInfo
+
+
+@dataclass
+class Message:
+    """Represents an incoming message."""
+
+    content: str
+    conversation_id: str
+    user_id: str
+    headers: Dict[str, str]
+    raw_body: Dict[str, Any]
+    conversation: Optional[Conversation] = None
+    memory: Optional["Memory"] = None
+    metadata: Optional[Dict[str, Any]] = None
+    schema: Optional[Dict[str, Any]] = None
+    # Keep singular for backward compatibility
+    attachment: Optional["AttachmentInfo"] = None
+    # Array of all attachments
+    attachments: List["AttachmentInfo"] = field(default_factory=list)
+    _attachment_client: Optional["AttachmentClient"] = field(default=None, repr=False)
+
+    def __post_init__(self):
+        """Ensure backward compatibility between attachment and attachments."""
+        # If attachments is empty but attachment exists, populate attachments
+        if not self.attachments and self.attachment:
+            self.attachments = [self.attachment]
+        # If attachment is None but attachments exist, set attachment to first
+        if self.attachment is None and self.attachments:
+            self.attachment = self.attachments[0]
+
+    @classmethod
+    def from_kafka_message(
+        cls,
+        body: Dict[str, Any],
+        headers: Dict[str, str],
+        conversation_client: Optional[ConversationClient] = None,
+        memory_client: Optional["MemoryClient"] = None,
+        attachment_client: Optional["AttachmentClient"] = None,
+        agent_name: Optional[str] = None,
+    ):
+        """Create Message from Kafka message."""
+        from .memory import Memory
+        from .attachment_client import AttachmentInfo
+
+        conversation_id = headers.get("conversationId")
+        user_id = headers.get("userId", "anonymous")
+
+        # Parse metadata from header
+        metadata = None
+        metadata_header = headers.get("messageMetadata")
+        if metadata_header:
+            try:
+                metadata = json.loads(metadata_header)
+            except json.JSONDecodeError:
+                pass
+
+        # Parse schema from header
+        schema = None
+        schema_header = headers.get("messageSchema")
+        if schema_header:
+            try:
+                schema = json.loads(schema_header)
+            except json.JSONDecodeError:
+                pass
+
+        # Parse attachments from header or body (prefer plural, fallback to singular)
+        attachments: List[AttachmentInfo] = []
+
+        # Try parsing plural 'attachments' header first (new format)
+        attachments_header = headers.get("attachments")
+        if attachments_header:
+            try:
+                attachments_data = json.loads(attachments_header)
+                attachments = [AttachmentInfo.from_dict(a) for a in attachments_data]
+            except json.JSONDecodeError:
+                pass
+
+        # Fallback to body attachments (plural)
+        if not attachments and body.get("attachments"):
+            try:
+                attachments = [
+                    AttachmentInfo.from_dict(a) for a in body["attachments"]
+                ]
+            except (TypeError, KeyError):
+                pass
+
+        # BACKWARD COMPATIBILITY: Fallback to singular 'attachment' header
+        if not attachments:
+            attachment_header = headers.get("attachment")
+            if attachment_header:
+                try:
+                    attachment_data = json.loads(attachment_header)
+                    attachments = [AttachmentInfo.from_dict(attachment_data)]
+                except json.JSONDecodeError:
+                    pass
+
+        # Fallback to singular body attachment
+        if not attachments and body.get("attachment"):
+            try:
+                attachments = [AttachmentInfo.from_dict(body["attachment"])]
+            except (TypeError, KeyError):
+                pass
+
+        conversation = None
+        if conversation_id and conversation_client:
+            conversation = Conversation(conversation_id, conversation_client, agent_name)
+
+        memory = None
+        if memory_client and user_id and conversation_id:
+            memory = Memory(user_id, conversation_id, memory_client)
+
+        return cls(
+            content=body.get("content", ""),
+            conversation_id=conversation_id,
+            user_id=user_id,
+            headers=headers,
+            raw_body=body,
+            conversation=conversation,
+            memory=memory,
+            metadata=metadata,
+            schema=schema,
+            attachment=attachments[0] if attachments else None,
+            attachments=attachments,
+            _attachment_client=attachment_client,
+        )
+
+    def has_attachment(self) -> bool:
+        """Check if message has at least one attachment."""
+        return len(self.attachments) > 0
+
+    def has_attachments(self) -> bool:
+        """Check if message has any attachments. Alias for has_attachment()."""
+        return len(self.attachments) > 0
+
+    def get_attachment_count(self) -> int:
+        """Get the number of attachments."""
+        return len(self.attachments)
+
+    def get_attachments(self) -> List["AttachmentInfo"]:
+        """Get all attachments."""
+        from .attachment_client import AttachmentInfo
+
+        return list(self.attachments)
+
+    async def get_attachment_bytes(self) -> bytes:
+        """
+        Download first attachment into memory as bytes.
+        For backward compatibility with single-attachment code.
+
+        Returns:
+            File content as bytes
+
+        Raises:
+            ValueError: If no attachment is available
+        """
+        if not self.attachments:
+            raise ValueError("No attachment available")
+        if not self._attachment_client:
+            raise ValueError("AttachmentClient not available")
+        return await self._attachment_client.download(self.attachments[0].url)
+
+    async def get_attachment_base64(self) -> str:
+        """
+        Get first attachment as base64 string (for LLM vision APIs).
+        For backward compatibility with single-attachment code.
+
+        Returns:
+            Base64 encoded string of the file content
+
+        Raises:
+            ValueError: If no attachment is available
+        """
+        if not self.attachments:
+            raise ValueError("No attachment available")
+        if not self._attachment_client:
+            raise ValueError("AttachmentClient not available")
+        return await self._attachment_client.download_as_base64(self.attachments[0].url)
+
+    async def get_attachment_buffer(self) -> BytesIO:
+        """
+        Get first attachment as BytesIO buffer (for document processing libraries).
+        For backward compatibility with single-attachment code.
+
+        Returns:
+            BytesIO buffer containing the file content
+
+        Raises:
+            ValueError: If no attachment is available
+        """
+        if not self.attachments:
+            raise ValueError("No attachment available")
+        if not self._attachment_client:
+            raise ValueError("AttachmentClient not available")
+        return await self._attachment_client.download_to_buffer(self.attachments[0].url)
+
+    async def get_attachment_bytes_at(self, index: int) -> bytes:
+        """
+        Download a specific attachment by index as bytes.
+
+        Args:
+            index: Zero-based index of the attachment
+
+        Returns:
+            File content as bytes
+
+        Raises:
+            ValueError: If index is out of bounds or no attachment client
+        """
+        if index < 0 or index >= len(self.attachments):
+            raise ValueError(
+                f"Invalid attachment index: {index}. "
+                f"Valid range: 0-{len(self.attachments) - 1}"
+            )
+        if not self._attachment_client:
+            raise ValueError("AttachmentClient not available")
+        return await self._attachment_client.download(self.attachments[index].url)
+
+    async def get_attachment_base64_at(self, index: int) -> str:
+        """
+        Get a specific attachment by index as base64 string.
+
+        Args:
+            index: Zero-based index of the attachment
+
+        Returns:
+            Base64 encoded string of the file content
+
+        Raises:
+            ValueError: If index is out of bounds or no attachment client
+        """
+        if index < 0 or index >= len(self.attachments):
+            raise ValueError(
+                f"Invalid attachment index: {index}. "
+                f"Valid range: 0-{len(self.attachments) - 1}"
+            )
+        if not self._attachment_client:
+            raise ValueError("AttachmentClient not available")
+        return await self._attachment_client.download_as_base64(
+            self.attachments[index].url
+        )
+
+    async def get_attachment_buffer_at(self, index: int) -> BytesIO:
+        """
+        Get a specific attachment by index as BytesIO buffer.
+
+        Args:
+            index: Zero-based index of the attachment
+
+        Returns:
+            BytesIO buffer containing the file content
+
+        Raises:
+            ValueError: If index is out of bounds or no attachment client
+        """
+        if index < 0 or index >= len(self.attachments):
+            raise ValueError(
+                f"Invalid attachment index: {index}. "
+                f"Valid range: 0-{len(self.attachments) - 1}"
+            )
+        if not self._attachment_client:
+            raise ValueError("AttachmentClient not available")
+        return await self._attachment_client.download_to_buffer(
+            self.attachments[index].url
+        )
+
+    async def get_all_attachment_bytes(self) -> List[bytes]:
+        """
+        Download all attachments as bytes.
+
+        Returns:
+            List of file contents as bytes in the same order as attachments
+
+        Raises:
+            ValueError: If no attachment client available
+        """
+        if not self._attachment_client:
+            raise ValueError("AttachmentClient not available")
+        return await asyncio.gather(
+            *[self._attachment_client.download(a.url) for a in self.attachments]
+        )
+
+    async def get_all_attachment_base64(self) -> List[str]:
+        """
+        Get all attachments as base64 strings.
+
+        Returns:
+            List of base64 encoded strings in the same order as attachments
+
+        Raises:
+            ValueError: If no attachment client available
+        """
+        if not self._attachment_client:
+            raise ValueError("AttachmentClient not available")
+        return await asyncio.gather(
+            *[
+                self._attachment_client.download_as_base64(a.url)
+                for a in self.attachments
+            ]
+        )
+
+    async def get_all_attachment_buffers(self) -> List[BytesIO]:
+        """
+        Get all attachments as BytesIO buffers.
+
+        Returns:
+            List of BytesIO buffers in the same order as attachments
+
+        Raises:
+            ValueError: If no attachment client available
+        """
+        if not self._attachment_client:
+            raise ValueError("AttachmentClient not available")
+        return await asyncio.gather(
+            *[
+                self._attachment_client.download_to_buffer(a.url)
+                for a in self.attachments
+            ]
+        )

basion_agent/streamer.py

@@ -0,0 +1,184 @@
+"""Streamer class for streaming responses with context manager support."""
+
+import json
+import logging
+from typing import Optional, Dict, Any, TypeVar, TYPE_CHECKING
+
+from .structural.base import StructuralStreamer
+
+if TYPE_CHECKING:
+    from .gateway_client import GatewayClient
+    from .message import Message
+    from .conversation_client import ConversationClient
+
+T = TypeVar('T', bound=StructuralStreamer)
+
+logger = logging.getLogger(__name__)
+
+
+class Streamer:
+    """Handles streaming responses back to the conversation via the gateway (async)."""
+
+    def __init__(
+        self,
+        agent_name: str,
+        original_message: "Message",
+        gateway_client: "GatewayClient",
+        conversation_client: Optional["ConversationClient"] = None,
+        awaiting: bool = False,
+        send_to: str = "user",
+    ):
+        self._agent_name = agent_name
+        self._original_message = original_message
+        self._gateway_client = gateway_client
+        self._conversation_client = conversation_client
+        self._awaiting = awaiting  # Used by context manager
+        self._send_to = send_to
+        self._finished = False
+        self._response_schema: Optional[Dict[str, Any]] = None
+        self._message_metadata: Optional[Dict[str, Any]] = None
+
+    def set_response_schema(self, schema: Dict[str, Any]) -> "Streamer":
+        """Set response schema (fluent API)."""
+        self._response_schema = schema
+        return self
+
+    def set_message_metadata(self, metadata: Dict[str, Any]) -> "Streamer":
+        """Set message metadata (fluent API)."""
+        self._message_metadata = metadata
+        return self
+
+    def stream_by(self, structural: T) -> T:
+        """Bind a structural streamer to this streamer for event sending.
+
+        Returns the structural object itself with streaming capability enabled.
+
+        Args:
+            structural: StructuralStreamer instance (Artifact, Surface, etc.)
+
+        Returns:
+            The same structural object, now bound to this streamer
+
+        Example:
+            artifact = Artifact()
+            s.stream_by(artifact).generating("Creating...")
+            s.stream_by(artifact).done(url="...", type="image")
+        """
+        return structural._bind_streamer(self)
+
+    def stream(
+        self, content: str, persist: bool = True, event_type: Optional[str] = None
+    ) -> None:
+        """Stream a content chunk (not done).
+
+        Args:
+            content: The content to stream
+            persist: If True, this content will be saved to DB. If False, only streamed.
+            event_type: Optional event type for the chunk (e.g., 'thinking', 'tool_call').
+        """
+        if self._finished:
+            raise RuntimeError("Streamer already finished")
+        self._send(content, done=False, persist=persist, event_type=event_type)
+
+    def write(self, content: str, persist: bool = True) -> None:
+        """Alias for stream() - write a content chunk.
+
+        Args:
+            content: The content to write
+            persist: If True, this content will be saved to DB. If False, only streamed.
+        """
+        self.stream(content, persist=persist)
+
+    async def finish(self, awaiting: Optional[bool] = None) -> None:
+        """Finish streaming with empty done=True chunk.
+
+        Args:
+            awaiting: If True, sets awaiting_route to this agent's name.
+                     If None, uses the awaiting value from constructor.
+                     If response_schema was set and awaiting is not explicitly False,
+                     awaiting is automatically set to True.
+        """
+        if self._finished:
+            return  # Idempotent
+
+        self._send("", done=True, persist=True)
+        self._finished = True
+
+        # Determine if we should set awaiting_route
+        # Priority: explicit awaiting arg > constructor awaiting > auto-await for schema
+        if awaiting is not None:
+            should_await = awaiting
+        elif self._awaiting:
+            should_await = True
+        elif self._response_schema is not None:
+            # Auto-await when response_schema is set (so next message comes back to this agent)
+            should_await = True
+        else:
+            should_await = False
+
+        # Store pending_response_schema if schema was set
+        if self._response_schema and self._conversation_client:
+            await self._set_pending_response_schema()
+
+        if should_await and self._conversation_client:
+            await self._set_awaiting_route()
+
+    def _send(
+        self,
+        content: str,
+        done: bool,
+        persist: bool = True,
+        event_type: Optional[str] = None,
+    ) -> None:
+        """Send message to router via gateway (sync - uses gRPC)."""
+        message_body = {"content": content, "done": done, "persist": persist}
+        if event_type:
+            message_body["eventType"] = event_type
+        headers = {
+            "conversationId": self._original_message.conversation_id,
+            "userId": self._original_message.user_id,
+            "from_": self._agent_name,
+            "to_": self._send_to,
+            "nextRoute": self._send_to,
+        }
+        if self._response_schema:
+            headers["responseSchema"] = json.dumps(self._response_schema)
+        if self._message_metadata:
+            headers["messageMetadata"] = json.dumps(self._message_metadata)
+
+        # Produce via gateway gRPC (synchronous operation)
+        self._gateway_client.produce(
+            topic="router.inbox",
+            key=self._original_message.conversation_id,
+            headers=headers,
+            body=message_body,
+        )
+
+    async def _set_pending_response_schema(self) -> None:
+        """Store response schema on conversation for form validation."""
+        try:
+            await self._conversation_client.update_conversation(
+                self._original_message.conversation_id,
+                {"pending_response_schema": self._response_schema},
+            )
+        except Exception as e:
+            logger.error(f"Failed to set pending_response_schema: {e}")
+
+    async def _set_awaiting_route(self) -> None:
+        """Set awaiting_route to this agent's name."""
+        try:
+            await self._conversation_client.update_conversation(
+                self._original_message.conversation_id,
+                {"awaiting_route": self._agent_name},
+            )
+        except Exception as e:
+            logger.error(f"Failed to set awaiting_route: {e}")
+
+    async def __aenter__(self) -> "Streamer":
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        if not self._finished:
+            # Pass None to let finish() determine awaiting based on schema or constructor value
+            await self.finish(awaiting=None)
+        return False  # Don't suppress exceptions

basion_agent/structural/__init__.py

@@ -0,0 +1,6 @@
+from .artifact import Artifact
+from .surface import Surface
+from .text_block import TextBlock
+from .stepper import Stepper
+
+__all__ = ["Artifact", "Surface", "TextBlock", "Stepper"]