march-agent 0.1.1__py3-none-any.whl

march_agent/heartbeat.py

@@ -0,0 +1,84 @@
+"""Heartbeat manager for keeping agent status active."""
+
+import logging
+import time
+import threading
+from typing import Optional, TYPE_CHECKING
+
+from .exceptions import HeartbeatError
+
+if TYPE_CHECKING:
+    from .gateway_client import GatewayClient
+
+logger = logging.getLogger(__name__)
+
+
+class HeartbeatManager:
+    """Manages periodic heartbeat signals to the API via the gateway proxy."""
+
+    def __init__(
+        self,
+        gateway_client: "GatewayClient",
+        agent_name: str,
+        interval: int = 60
+    ):
+        self.gateway_client = gateway_client
+        self.agent_name = agent_name
+        self.interval = interval
+        self._running = False
+        self._thread: Optional[threading.Thread] = None
+
+    def start(self):
+        """Start the heartbeat thread."""
+        if self._running:
+            raise HeartbeatError("Heartbeat already running")
+
+        self._running = True
+        self._thread = threading.Thread(target=self._heartbeat_loop, daemon=True)
+        self._thread.start()
+        logger.info(f"Heartbeat started for agent '{self.agent_name}'")
+
+    def stop(self):
+        """Stop the heartbeat thread."""
+        if not self._running:
+            return
+
+        self._running = False
+        if self._thread:
+            self._thread.join(timeout=5.0)
+        logger.info("Heartbeat stopped")
+
+    def _heartbeat_loop(self):
+        """Background loop that sends heartbeat periodically."""
+        while self._running:
+            try:
+                self._send_heartbeat()
+            except Exception as e:
+                logger.error(f"Heartbeat send failed: {e}", exc_info=True)
+
+            # Sleep in small intervals to allow quick shutdown
+            for _ in range(self.interval):
+                if not self._running:
+                    break
+                time.sleep(1)
+
+    def _send_heartbeat(self):
+        """Send a single heartbeat to the API via gateway proxy."""
+        payload = {"name": self.agent_name}
+
+        try:
+            response = self.gateway_client.http_post(
+                "ai-inventory",
+                "/api/v1/health/heartbeat",
+                json=payload,
+                timeout=5.0
+            )
+            if response.status_code == 404:
+                logger.warning(f"Agent '{self.agent_name}' not found. Re-registration may be needed.")
+            elif response.status_code not in (200, 201):
+                logger.warning(f"Heartbeat returned status {response.status_code}: {response.text}")
+            else:
+                logger.debug(f"Heartbeat sent successfully for agent '{self.agent_name}'")
+        except Exception as e:
+            logger.error(f"Heartbeat request failed: {e}")
+            raise HeartbeatError(f"Failed to send heartbeat: {e}")

march_agent/memory.py

@@ -0,0 +1,73 @@
+"""Memory context for accessing AI memory from message handler."""
+
+from typing import List, Optional
+
+from .memory_client import MemoryClient, MemorySearchResult, UserSummary
+
+
+class Memory:
+    """User-facing memory context, attached to Message.
+
+    Provides easy access to memory functions scoped to the current user/conversation.
+    """
+
+    def __init__(self, user_id: str, conversation_id: str, client: MemoryClient):
+        self.user_id = user_id
+        self.conversation_id = conversation_id
+        self._client = client
+
+    async def query_about_user(
+        self,
+        query: str,
+        limit: int = 10,
+        threshold: int = 70,
+        context_messages: int = 0,
+    ) -> List[MemorySearchResult]:
+        """Query user's long-term memory with semantic search.
+
+        Args:
+            query: Semantic search query
+            limit: Maximum results (1-100)
+            threshold: Similarity threshold 0-100
+            context_messages: Surrounding messages to include (0-20)
+
+        Returns:
+            List of MemorySearchResult
+        """
+        return await self._client.search(
+            query=query,
+            user_id=self.user_id,
+            limit=limit,
+            min_similarity=threshold,
+            context_messages=context_messages,
+        )
+
+    async def query_about_conversation(
+        self,
+        query: str,
+        limit: int = 10,
+        threshold: int = 70,
+        context_messages: int = 0,
+    ) -> List[MemorySearchResult]:
+        """Query conversation's long-term memory with semantic search.
+
+        Args:
+            query: Semantic search query
+            limit: Maximum results (1-100)
+            threshold: Similarity threshold 0-100
+            context_messages: Surrounding messages to include (0-20)
+
+        Returns:
+            List of MemorySearchResult
+        """
+        return await self._client.search(
+            query=query,
+            conversation_id=self.conversation_id,
+            limit=limit,
+            min_similarity=threshold,
+            context_messages=context_messages,
+        )
+
+    async def get_user_summary(self) -> Optional[UserSummary]:
+        """Get user's summary (aggregated across all conversations)."""
+        return await self._client.get_user_summary(self.user_id)

march_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}")

march_agent/message.py

@@ -0,0 +1,80 @@
+"""Message data structure."""
+
+import json
+from dataclasses import dataclass
+from typing import Dict, Any, Optional, TYPE_CHECKING
+
+from .conversation import Conversation
+from .conversation_client import ConversationClient
+
+if TYPE_CHECKING:
+    from .memory import Memory
+    from .memory_client import MemoryClient
+
+
+@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
+
+    @classmethod
+    def from_kafka_message(
+        cls,
+        body: Dict[str, Any],
+        headers: Dict[str, str],
+        conversation_client: Optional[ConversationClient] = None,
+        memory_client: Optional["MemoryClient"] = None,
+        agent_name: Optional[str] = None,
+    ):
+        """Create Message from Kafka message."""
+        from .memory import Memory
+
+        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
+
+        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,
+        )

march_agent/streamer.py

@@ -0,0 +1,220 @@
+"""Streamer class for streaming responses with context manager support."""
+
+import json
+import logging
+from typing import Optional, Dict, Any, List, TYPE_CHECKING
+
+from .artifact import Artifact, ArtifactType
+
+if TYPE_CHECKING:
+    from .gateway_client import GatewayClient
+    from .message import Message
+    from .conversation_client import ConversationClient
+
+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
+        self._artifacts: List[Artifact] = []
+        self._artifacts_sent: bool = False
+
+    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 add_artifact(
+        self,
+        url: str,
+        type: str,
+        title: Optional[str] = None,
+        description: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> "Streamer":
+        """Add an artifact to the message (fluent API).
+
+        Args:
+            url: URL to the artifact (file or iframe src)
+            type: Type of artifact ('document', 'image', 'iframe', 'video', 'audio', 'code', 'link', 'file')
+            title: Optional display title
+            description: Optional description or alt text
+            metadata: Optional metadata dict (size, mimeType, dimensions, etc.)
+
+        Returns:
+            self for method chaining
+
+        Example:
+            streamer.add_artifact(
+                url="https://example.com/report.pdf",
+                type="document",
+                title="Monthly Report",
+                metadata={"size": 1024000, "mimeType": "application/pdf"}
+            )
+        """
+        artifact = Artifact(
+            url=url,
+            type=ArtifactType(type),
+            title=title,
+            description=description,
+            metadata=metadata,
+            position=len(self._artifacts),
+        )
+        self._artifacts.append(artifact)
+        return self
+
+    def set_artifacts(self, artifacts: List[Dict[str, Any]]) -> "Streamer":
+        """Set all artifacts at once (replaces any existing).
+
+        Args:
+            artifacts: List of artifact dicts with url, type, title, description, metadata
+
+        Returns:
+            self for method chaining
+        """
+        self._artifacts = [
+            Artifact(
+                url=a["url"],
+                type=ArtifactType(a["type"]),
+                title=a.get("title"),
+                description=a.get("description"),
+                metadata=a.get("metadata"),
+                position=i,
+            )
+            for i, a in enumerate(artifacts)
+        ]
+        return 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 self._finished:
+            return  # Idempotent
+
+        self._send("", done=True, persist=True)
+        self._finished = True
+
+        # Use constructor's awaiting if not explicitly provided
+        should_await = awaiting if awaiting is not None else self._awaiting
+
+        # 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)
+
+        # Send artifacts with first chunk (only once)
+        if self._artifacts and not self._artifacts_sent:
+            headers["artifacts"] = json.dumps([a.to_dict() for a in self._artifacts])
+            self._artifacts_sent = True
+
+        # 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:
+            await self.finish(awaiting=self._awaiting)
+        return False  # Don't suppress exceptions