march-agent 0.1.1__py3-none-any.whl

march_agent/app.py

@@ -0,0 +1,416 @@
+"""Main application class for March AI Agent framework (async)."""
+
+import asyncio
+import logging
+import signal
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from typing import Optional, Dict, Any, List, Set
+
+from .agent import Agent
+from .gateway_client import GatewayClient
+from .conversation_client import ConversationClient
+from .memory_client import MemoryClient
+from .exceptions import RegistrationError, ConfigurationError
+
+logger = logging.getLogger(__name__)
+
+
+class MarchAgentApp:
+    """
+    Main application class for March AI Agent framework.
+
+    Example:
+        from march_agent import MarchAgentApp
+
+        app = MarchAgentApp(
+            gateway_url="agent-gateway:8080",
+            api_key="agent-key-1"
+        )
+
+        # Register agents
+        medical_agent = app.register_me(
+            name="medical-qa-agent",
+            about="Medical Q&A",
+            document="Answers medical questions"
+        )
+
+        @medical_agent.on_message
+        def handle_message(message, sender):
+            # Access conversation history
+            history = message.conversation.get_history(limit=5)
+
+            # Stream response
+            with medical_agent.streamer(message) as s:
+                s.stream("Processing...")
+                s.stream("Done!")
+
+        app.run()
+    """
+
+    def __init__(
+        self,
+        gateway_url: str,
+        api_key: str,
+        heartbeat_interval: int = 60,
+        max_concurrent_tasks: int = 100,
+        error_message_template: str = (
+            "I encountered an error while processing your message. "
+            "Please try again or contact support if the issue persists."
+        ),
+        secure: bool = False,
+    ):
+        """
+        Initialize March Agent App.
+
+        Args:
+            gateway_url: Gateway endpoint (e.g., "agent-gateway:8080")
+            api_key: API key for gateway authentication
+            heartbeat_interval: Heartbeat interval in seconds
+            max_concurrent_tasks: Maximum number of concurrent message handlers (default: 100)
+            error_message_template: Template for error messages sent to users
+            secure: If True, use TLS for gRPC and HTTPS for HTTP requests
+        """
+        self.gateway_url = gateway_url
+        self.api_key = api_key
+        self.heartbeat_interval = heartbeat_interval
+        self.max_concurrent_tasks = max_concurrent_tasks
+        self.error_message_template = error_message_template
+        self.secure = secure
+
+        # Create gateway client
+        self.gateway_client = GatewayClient(gateway_url, api_key, secure=secure)
+
+        # Create conversation client using gateway proxy
+        self.conversation_client = ConversationClient(
+            self.gateway_client.conversation_store_url
+        )
+
+        # Create memory client using gateway proxy
+        self.memory_client = MemoryClient(
+            self.gateway_client.ai_memory_url
+        )
+
+        self._agents: List[Agent] = []
+        self._running = False
+        self._connected = False
+        self._active_tasks: Set[asyncio.Task] = set()
+        self._executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="consume")
+        self._task_start_times: Dict[asyncio.Task, float] = {}
+        self._hung_task_threshold = 300.0  # 5 minutes
+
+        logger.info(f"MarchAgentApp initialized (gateway: {gateway_url}, max_concurrent: {max_concurrent_tasks})")
+
+    def register_me(
+        self,
+        name: str,
+        about: str,
+        document: str,
+        representation_name: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        related_pages: Optional[List[Dict[str, str]]] = None,
+    ) -> Agent:
+        """
+        Register an agent with the backend.
+
+        Example:
+            medical_agent = app.register_me(
+                name="medical-qa",
+                about="Medical question answering",
+                document="Answers medical questions using AI",
+                representation_name="Medical Q&A Bot",
+                related_pages=[
+                    {"name": "Dashboard", "endpoint": "/dashboard"},
+                    {"name": "Reports", "endpoint": "/reports"},
+                ]
+            )
+
+            @medical_agent.on_message
+            def handle_message(message, sender):
+                with medical_agent.streamer(message) as s:
+                    s.stream("Processing...")
+
+        Args:
+            name: Unique agent name (used for routing)
+            about: Short description
+            document: Detailed documentation
+            representation_name: Display name (optional)
+            metadata: Additional metadata (optional)
+            related_pages: List of related pages with 'name' and 'endpoint' keys (optional)
+                          Example: [{"name": "Dashboard", "endpoint": "/dashboard"}]
+
+        Returns:
+            Agent instance
+        """
+        # Register agent with backend via gateway proxy
+        logger.info(f"Registering agent '{name}'...")
+
+        payload = {
+            "name": name,
+            "about": about,
+            "document": document,
+            "representationName": representation_name or name,
+        }
+        if metadata:
+            payload["metadata"] = metadata
+        if related_pages:
+            payload["relatedPages"] = related_pages
+
+        try:
+            response = self.gateway_client.http_post(
+                "ai-inventory",
+                "/api/v1/agents/register",
+                json=payload,
+                timeout=30
+            )
+            if response.status_code == 201:
+                agent_data = response.json()
+                logger.info(f"Agent '{name}' registered successfully")
+            else:
+                raise RegistrationError(f"Registration failed: {response.text}")
+        except Exception as e:
+            raise RegistrationError(f"Failed to register agent '{name}': {e}")
+
+        # Create agent instance
+        agent = Agent(
+            name=name,
+            gateway_client=self.gateway_client,
+            agent_data=agent_data,
+            heartbeat_interval=self.heartbeat_interval,
+            conversation_client=self.conversation_client,
+            memory_client=self.memory_client,
+            error_message_template=self.error_message_template,
+        )
+
+        self._agents.append(agent)
+        logger.info(f"Agent '{name}' ready")
+
+        return agent
+
+    def run(self):
+        """Start all registered agents and block until shutdown."""
+        if not self._agents:
+            raise ConfigurationError(
+                "No agents registered. Use app.register_me() to register agents."
+            )
+
+        logger.info(f"Starting {len(self._agents)} agent(s)...")
+        self._running = True
+
+        # Connect to gateway with all agent names
+        agent_names = [agent.name for agent in self._agents]
+        try:
+            self.gateway_client.connect(agent_names)
+            self._connected = True
+        except Exception as e:
+            raise ConfigurationError(f"Failed to connect to gateway: {e}")
+
+        # Setup signal handlers
+        signal.signal(signal.SIGINT, self._shutdown_handler)
+        signal.signal(signal.SIGTERM, self._shutdown_handler)
+
+        # Initialize all agents with the connected gateway
+        for agent in self._agents:
+            agent._initialize_with_gateway()
+
+        # Start all agents in background threads (they just stay alive)
+        threads = []
+        for agent in self._agents:
+            thread = threading.Thread(target=agent.start_consuming, daemon=True)
+            thread.start()
+            threads.append(thread)
+
+        logger.info("All agents started. Press Ctrl+C to stop.")
+
+        # Run the single consume loop in the main thread
+        try:
+            self._consume_loop()
+        except KeyboardInterrupt:
+            logger.info("Interrupted by user")
+        finally:
+            self._shutdown()
+
+    def _consume_loop(self):
+        """Main consume loop that dispatches messages to agents concurrently."""
+        # Run the async consume loop
+        asyncio.run(self._consume_loop_async())
+
+    async def _consume_loop_async(self):
+        """Async consume loop with concurrent message processing."""
+        # Build topic -> agent mapping
+        topic_to_agent = {
+            f"{agent.name}.inbox": agent for agent in self._agents
+        }
+
+        logger.info(
+            f"Starting consume loop for topics: {list(topic_to_agent.keys())} "
+            f"(max_concurrent: {self.max_concurrent_tasks})"
+        )
+
+        loop = asyncio.get_event_loop()
+
+        try:
+            while self._running:
+                try:
+                    # Run sync consume_one in executor to not block the event loop
+                    msg = await loop.run_in_executor(
+                        self._executor,
+                        lambda: self.gateway_client.consume_one(timeout=0.5)
+                    )
+
+                    if msg:
+                        agent = topic_to_agent.get(msg.topic)
+                        if agent:
+                            # Wait if we've hit max concurrency
+                            while len(self._active_tasks) >= self.max_concurrent_tasks:
+                                if not self._active_tasks:
+                                    break
+
+                                # Wait with timeout to detect hung tasks
+                                done, self._active_tasks = await asyncio.wait(
+                                    self._active_tasks,
+                                    return_when=asyncio.FIRST_COMPLETED,
+                                    timeout=30.0  # 30 second timeout
+                                )
+
+                                # Check for errors in completed tasks
+                                for task in done:
+                                    self._task_start_times.pop(task, None)
+                                    if task.exception():
+                                        logger.error(
+                                            f"Task failed with exception: {task.exception()}"
+                                        )
+
+                                # If timeout with no completions, check for hung tasks
+                                if not done:
+                                    current_time = time.time()
+                                    hung_tasks = [
+                                        task for task in self._active_tasks
+                                        if current_time - self._task_start_times.get(task, current_time) > self._hung_task_threshold
+                                    ]
+
+                                    if hung_tasks:
+                                        logger.error(
+                                            f"Detected {len(hung_tasks)} hung tasks (>{self._hung_task_threshold}s). "
+                                            f"Total active: {len(self._active_tasks)}"
+                                        )
+                                        # Optional: cancel hung tasks
+                                        # for task in hung_tasks:
+                                        #     task.cancel()
+                                        #     self._task_start_times.pop(task, None)
+                                    else:
+                                        logger.warning(
+                                            f"No tasks completed in 30s. Active tasks: {len(self._active_tasks)} "
+                                            f"(may indicate slow handlers or blocking code)"
+                                        )
+
+                            # Create task for concurrent processing
+                            task = asyncio.create_task(
+                                self._handle_message_safe(agent, msg)
+                            )
+                            self._active_tasks.add(task)
+                            self._task_start_times[task] = time.time()
+                            # Auto-remove from set when done
+                            task.add_done_callback(self._active_tasks.discard)
+
+                            logger.debug(
+                                f"Dispatched message to {agent.name}, "
+                                f"active tasks: {len(self._active_tasks)}"
+                            )
+                        else:
+                            logger.warning(f"No agent for topic: {msg.topic}")
+
+                except Exception as e:
+                    if self._running:
+                        logger.error(f"Error in consume loop: {e}")
+                        # Try to reconnect if connection lost
+                        await asyncio.sleep(1.0)
+                        try:
+                            self._reconnect()
+                        except Exception as re:
+                            logger.error(f"Reconnect failed: {re}")
+                            await asyncio.sleep(5.0)
+
+        finally:
+            # Wait for all active tasks to complete on shutdown
+            if self._active_tasks:
+                logger.info(
+                    f"Waiting for {len(self._active_tasks)} active tasks to complete..."
+                )
+                await asyncio.gather(*self._active_tasks, return_exceptions=True)
+                logger.info("All tasks completed")
+
+    async def _handle_message_safe(self, agent: Agent, msg) -> None:
+        """Wrapper that handles errors from message handlers gracefully."""
+        try:
+            await agent._handle_message_async(msg)
+        except Exception as e:
+            logger.error(
+                f"Error handling message for {agent.name}: {e}",
+                exc_info=True
+            )
+
+    def _reconnect(self):
+        """Attempt to reconnect to the gateway."""
+        logger.info("Attempting to reconnect to gateway...")
+        agent_names = [agent.name for agent in self._agents]
+        self.gateway_client.close()
+        self.gateway_client.connect(agent_names)
+        logger.info("Reconnected to gateway successfully")
+
+    def _shutdown_handler(self, signum, frame):
+        """Handle shutdown signals."""
+        logger.info(f"Received signal {signum}, initiating shutdown...")
+        self._running = False
+
+    def _shutdown(self):
+        """Shutdown all agents gracefully."""
+        logger.info("Shutting down all agents...")
+        self._running = False
+
+        # Shutdown the executor
+        try:
+            self._executor.shutdown(wait=False)
+        except Exception as e:
+            logger.error(f"Error shutting down executor: {e}")
+
+        # Close async sessions - need to create a new loop since asyncio.run closed its loop
+        try:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            try:
+                loop.run_until_complete(self._close_async_sessions())
+            finally:
+                loop.close()
+        except Exception as e:
+            logger.error(f"Error closing async sessions: {e}")
+
+        for agent in self._agents:
+            try:
+                agent.shutdown()
+            except Exception as e:
+                logger.error(f"Error shutting down agent {agent.name}: {e}")
+
+        # Close gateway connection
+        if self._connected:
+            self.gateway_client.close()
+
+        logger.info("All agents shut down successfully")
+
+    async def _close_async_sessions(self):
+        """Close all async HTTP sessions."""
+        try:
+            await self.gateway_client.close_async()
+        except Exception as e:
+            logger.error(f"Error closing gateway async session: {e}")
+
+        try:
+            await self.conversation_client.close()
+        except Exception as e:
+            logger.error(f"Error closing conversation client: {e}")
+
+        try:
+            await self.memory_client.close()
+        except Exception as e:
+            logger.error(f"Error closing memory client: {e}")

march_agent/artifact.py

@@ -0,0 +1,58 @@
+"""Artifact class for attaching files/URLs to messages."""
+
+from dataclasses import dataclass
+from typing import Optional, Dict, Any
+from enum import Enum
+
+
+class ArtifactType(str, Enum):
+    """Artifact type enum for categorizing attachments."""
+
+    DOCUMENT = "document"  # PDF, DOC, etc.
+    IMAGE = "image"  # PNG, JPG, GIF, etc.
+    IFRAME = "iframe"  # Embeddable content (maps, charts)
+    VIDEO = "video"  # Video files or embeds
+    AUDIO = "audio"  # Audio files
+    CODE = "code"  # Code snippets with syntax highlighting
+    LINK = "link"  # External links with preview
+    FILE = "file"  # Generic file download
+
+
+@dataclass
+class Artifact:
+    """Represents an artifact (URL attachment) for a message.
+
+    Artifacts are URLs to files, images, iframes, or other resources
+    that agents can attach to their responses.
+
+    Example:
+        artifact = Artifact(
+            url="https://example.com/report.pdf",
+            type=ArtifactType.DOCUMENT,
+            title="Monthly Report",
+            metadata={"size": 1024000, "mimeType": "application/pdf"}
+        )
+    """
+
+    url: str
+    type: ArtifactType
+    title: Optional[str] = None
+    description: Optional[str] = None
+    metadata: Optional[Dict[str, Any]] = None
+    position: Optional[int] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dict for JSON serialization."""
+        d: Dict[str, Any] = {
+            "url": self.url,
+            "type": self.type.value,
+        }
+        if self.title:
+            d["title"] = self.title
+        if self.description:
+            d["description"] = self.description
+        if self.metadata:
+            d["metadata"] = self.metadata
+        if self.position is not None:
+            d["position"] = self.position
+        return d

march_agent/checkpoint_client.py

@@ -0,0 +1,169 @@
+"""HTTP client for checkpoint storage API."""
+
+import logging
+from typing import List, Dict, Any, Optional
+import aiohttp
+
+from .exceptions import APIException
+
+logger = logging.getLogger(__name__)
+
+
+class CheckpointClient:
+    """Async HTTP client for checkpoint-store API.
+
+    This client communicates with the conversation-store's checkpoint endpoints
+    to store and retrieve LangGraph-compatible checkpoints.
+    """
+
+    def __init__(self, base_url: str):
+        """Initialize checkpoint client.
+
+        Args:
+            base_url: Base URL for the checkpoint API (e.g., http://gateway/s/conversation-store)
+        """
+        self.base_url = base_url.rstrip("/")
+        self._session: Optional[aiohttp.ClientSession] = None
+
+    async def _get_session(self) -> aiohttp.ClientSession:
+        """Get or create aiohttp session."""
+        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):
+        """Close the aiohttp session."""
+        if self._session and not self._session.closed:
+            await self._session.close()
+
+    async def put(
+        self,
+        config: Dict[str, Any],
+        checkpoint: Dict[str, Any],
+        metadata: Dict[str, Any],
+        new_versions: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """Store a checkpoint.
+
+        Args:
+            config: RunnableConfig with configurable containing thread_id, checkpoint_ns, checkpoint_id
+            checkpoint: Checkpoint data (channel_values, channel_versions, etc.)
+            metadata: Checkpoint metadata (source, step, parents, writes)
+            new_versions: New channel versions (optional)
+
+        Returns:
+            Config of the stored checkpoint
+        """
+        url = f"{self.base_url}/checkpoints/"
+        payload = {
+            "config": config,
+            "checkpoint": checkpoint,
+            "metadata": metadata,
+            "new_versions": new_versions or {},
+        }
+
+        session = await self._get_session()
+        try:
+            async with session.put(url, json=payload) as response:
+                if response.status >= 400:
+                    error_text = await response.text()
+                    raise APIException(f"Failed to store checkpoint: {response.status} - {error_text}")
+                return await response.json()
+        except aiohttp.ClientError as e:
+            raise APIException(f"Failed to store checkpoint: {e}")
+
+    async def get_tuple(
+        self,
+        thread_id: str,
+        checkpoint_ns: str = "",
+        checkpoint_id: Optional[str] = None,
+    ) -> Optional[Dict[str, Any]]:
+        """Get a checkpoint tuple.
+
+        Args:
+            thread_id: Thread identifier
+            checkpoint_ns: Checkpoint namespace (default "")
+            checkpoint_id: Specific checkpoint ID (latest if not provided)
+
+        Returns:
+            CheckpointTuple dict or None if not found
+        """
+        url = f"{self.base_url}/checkpoints/{thread_id}"
+        params = {"checkpoint_ns": checkpoint_ns}
+        if checkpoint_id:
+            params["checkpoint_id"] = checkpoint_id
+
+        session = await self._get_session()
+        try:
+            async with session.get(url, params=params) as response:
+                if response.status == 404:
+                    return None
+                if response.status >= 400:
+                    error_text = await response.text()
+                    raise APIException(f"Failed to get checkpoint: {response.status} - {error_text}")
+                result = await response.json()
+                # API returns null for not found
+                return result if result else None
+        except aiohttp.ClientError as e:
+            raise APIException(f"Failed to get checkpoint: {e}")
+
+    async def list(
+        self,
+        thread_id: Optional[str] = None,
+        checkpoint_ns: Optional[str] = None,
+        before: Optional[str] = None,
+        limit: Optional[int] = None,
+    ) -> List[Dict[str, Any]]:
+        """List checkpoints.
+
+        Args:
+            thread_id: Filter by thread ID
+            checkpoint_ns: Filter by namespace
+            before: Return checkpoints before this checkpoint_id
+            limit: Maximum number of checkpoints to return
+
+        Returns:
+            List of CheckpointTuple dicts
+        """
+        url = f"{self.base_url}/checkpoints/"
+        params = {}
+        if thread_id:
+            params["thread_id"] = thread_id
+        if checkpoint_ns is not None:
+            params["checkpoint_ns"] = checkpoint_ns
+        if before:
+            params["before"] = before
+        if limit:
+            params["limit"] = limit
+
+        session = await self._get_session()
+        try:
+            async with session.get(url, params=params) as response:
+                if response.status >= 400:
+                    error_text = await response.text()
+                    raise APIException(f"Failed to list checkpoints: {response.status} - {error_text}")
+                return await response.json()
+        except aiohttp.ClientError as e:
+            raise APIException(f"Failed to list checkpoints: {e}")
+
+    async def delete_thread(self, thread_id: str) -> Dict[str, Any]:
+        """Delete all checkpoints for a thread.
+
+        Args:
+            thread_id: Thread identifier
+
+        Returns:
+            Dict with thread_id and deleted count
+        """
+        url = f"{self.base_url}/checkpoints/{thread_id}"
+
+        session = await self._get_session()
+        try:
+            async with session.delete(url) as response:
+                if response.status >= 400:
+                    error_text = await response.text()
+                    raise APIException(f"Failed to delete checkpoints: {response.status} - {error_text}")
+                return await response.json()
+        except aiohttp.ClientError as e:
+            raise APIException(f"Failed to delete checkpoints: {e}")

march_agent/checkpointer.py

@@ -0,0 +1,16 @@
+"""Checkpoint client module.
+
+This module provides the low-level HTTP client for checkpoint operations.
+For the LangGraph-compatible checkpointer, see march_agent.extensions.langgraph.
+
+Usage:
+    from march_agent.extensions.langgraph import HTTPCheckpointSaver
+
+    app = MarchAgentApp(gateway_url="agent-gateway:8080", api_key="key")
+    checkpointer = HTTPCheckpointSaver(app=app)
+"""
+
+# Re-export CheckpointClient for backwards compatibility
+from .checkpoint_client import CheckpointClient
+
+__all__ = ["CheckpointClient"]