jarviscore-framework 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

jarviscore/profiles/customagent.py

@@ -1,87 +1,130 @@
 """
-CustomAgent - User-controlled execution profile.
-
-User provides their own implementation using any framework:
-- LangChain
-- MCP (Model Context Protocol)
-- CrewAI
-- Raw Python
-- Any other tool/framework
+CustomAgent - User-controlled execution profile with P2P message handling.
+
+Unified profile for building agents that:
+- Handle P2P mesh communication (requests, notifications)
+- Execute workflow tasks
+- Integrate with HTTP APIs (FastAPI, Flask, etc.)
+
+Example - Basic P2P Agent:
+    class AnalystAgent(CustomAgent):
+        role = "analyst"
+        capabilities = ["analysis"]
+
+        async def on_peer_request(self, msg):
+            result = await self.analyze(msg.data)
+            return {"status": "success", "result": result}
+
+Example - With FastAPI:
+    from fastapi import FastAPI
+    from jarviscore.integrations.fastapi import JarvisLifespan
+
+    class ProcessorAgent(CustomAgent):
+        role = "processor"
+        capabilities = ["processing"]
+
+        async def on_peer_request(self, msg):
+            return {"result": await self.process(msg.data)}
+
+    app = FastAPI(lifespan=JarvisLifespan(ProcessorAgent(), mode="p2p"))
 """
-from typing import Dict, Any
+from typing import Dict, Any, Optional
+import asyncio
+import logging
+
 from jarviscore.core.profile import Profile
 
+logger = logging.getLogger(__name__)
+
 
 class CustomAgent(Profile):
     """
-    Custom execution profile with full user control.
+    User-controlled agent profile with P2P message handling.
+
+    For P2P messaging, implement these handlers:
+        on_peer_request(msg) - Handle requests, return response
+        on_peer_notify(msg)  - Handle notifications (fire-and-forget)
+        on_error(error, msg) - Handle errors
 
-    User defines:
-    - role: str
-    - capabilities: List[str]
-    - setup(): Initialize custom framework/tools
-    - execute_task(): Custom execution logic
+    For workflow execution:
+        execute_task(task)   - Handle workflow tasks directly
+        (defaults to delegating to on_peer_request)
 
-    Framework provides:
-    - Orchestration (task claiming, dependencies, nudging)
-    - P2P coordination (agent discovery, task routing)
-    - State management (crash recovery, HITL)
-    - Cost tracking (if user provides token counts)
+    Configuration:
+        listen_timeout: Seconds to wait for messages (default: 1.0)
+        auto_respond: Auto-send on_peer_request return value (default: True)
 
-    Example with LangChain:
-        class APIAgent(CustomAgent):
-            role = "api_client"
-            capabilities = ["api_calls"]
+    Example - P2P Agent:
+        class AnalystAgent(CustomAgent):
+            role = "analyst"
+            capabilities = ["analysis"]
+
+            async def on_peer_request(self, msg):
+                result = await self.analyze(msg.data)
+                return {"status": "success", "result": result}
+
+    Example - With LangChain:
+        class LangChainAgent(CustomAgent):
+            role = "assistant"
+            capabilities = ["chat"]
 
             async def setup(self):
                 await super().setup()
                 from langchain.agents import Agent
                 self.lc_agent = Agent(...)
 
-            async def execute_task(self, task):
-                result = await self.lc_agent.run(task["task"])
+            async def on_peer_request(self, msg):
+                result = await self.lc_agent.run(msg.data["query"])
                 return {"status": "success", "output": result}
 
-    Example with MCP:
+    Example - With MCP:
         class MCPAgent(CustomAgent):
             role = "tool_user"
             capabilities = ["mcp_tools"]
-            mcp_server_url = "stdio://./server.py"
 
             async def setup(self):
                 await super().setup()
                 from mcp import Client
-                self.mcp = Client(self.mcp_server_url)
+                self.mcp = Client("stdio://./server.py")
                 await self.mcp.connect()
 
-            async def execute_task(self, task):
-                result = await self.mcp.call_tool("my_tool", task["params"])
+            async def on_peer_request(self, msg):
+                result = await self.mcp.call_tool("my_tool", msg.data)
                 return {"status": "success", "data": result}
 
-    Example with Raw Python:
-        class DataProcessor(CustomAgent):
+    Example - With FastAPI:
+        from fastapi import FastAPI
+        from jarviscore.integrations.fastapi import JarvisLifespan
+
+        class ProcessorAgent(CustomAgent):
             role = "processor"
             capabilities = ["data_processing"]
 
-            async def execute_task(self, task):
-                # Pure Python logic
-                data = task["params"]["data"]
-                processed = [x * 2 for x in data]
-                return {"status": "success", "output": processed}
+            async def on_peer_request(self, msg):
+                if msg.data.get("action") == "process":
+                    return {"result": await self.process(msg.data["payload"])}
+                return {"error": "unknown action"}
+
+        agent = ProcessorAgent()
+        app = FastAPI(lifespan=JarvisLifespan(agent, mode="p2p"))
+
+        @app.post("/process")
+        async def process_endpoint(data: dict, request: Request):
+            # HTTP endpoint - primary interface
+            agent = request.app.state.jarvis_agents["processor"]
+            return await agent.process(data)
     """
 
-    def __init__(self, agent_id=None):
-        super().__init__(agent_id)
+    # Configuration - can be overridden in subclasses
+    listen_timeout: float = 1.0  # Seconds to wait for messages
+    auto_respond: bool = True    # Automatically send response for requests
 
-        # User can add any custom attributes
-        # e.g., mcp_server_url, langchain_config, etc.
+    def __init__(self, agent_id: Optional[str] = None):
+        super().__init__(agent_id)
 
     async def setup(self):
         """
-        User implements this to initialize custom framework/tools.
-
-        DAY 1: Base implementation (user overrides)
-        DAY 5+: Full examples with LangChain, MCP, etc.
+        Initialize agent resources. Override to add custom setup.
 
         Example:
             async def setup(self):
@@ -91,47 +134,232 @@ class CustomAgent(Profile):
                 self.agent = Agent(...)
         """
         await super().setup()
-
         self._logger.info(f"CustomAgent setup: {self.agent_id}")
-        self._logger.info(
-            f"  Note: Override setup() to initialize your custom framework"
+
+    # ─────────────────────────────────────────────────────────────────
+    # P2P Message Handling
+    # ─────────────────────────────────────────────────────────────────
+
+    async def run(self):
+        """
+        Listener loop - receives and dispatches P2P messages.
+
+        Runs automatically in P2P mode. Dispatches messages to:
+        - on_peer_request() for request-response messages
+        - on_peer_notify() for fire-and-forget notifications
+
+        You typically don't need to override this. Just implement the handlers.
+        """
+        self._logger.info(f"[{self.role}] Listener loop started")
+
+        while not self.shutdown_requested:
+            try:
+                # Wait for incoming message with timeout
+                # Timeout allows periodic shutdown_requested checks
+                msg = await self.peers.receive(timeout=self.listen_timeout)
+
+                if msg is None:
+                    # Timeout - no message, continue loop to check shutdown
+                    continue
+
+                # Dispatch to appropriate handler
+                await self._dispatch_message(msg)
+
+            except asyncio.CancelledError:
+                self._logger.debug(f"[{self.role}] Listener loop cancelled")
+                raise
+            except Exception as e:
+                self._logger.error(f"[{self.role}] Listener loop error: {e}")
+                await self.on_error(e, None)
+
+        self._logger.info(f"[{self.role}] Listener loop stopped")
+
+    async def _dispatch_message(self, msg):
+        """
+        Dispatch message to appropriate handler based on message type.
+
+        Handles:
+        - REQUEST messages: calls on_peer_request, sends response if auto_respond=True
+        - NOTIFY messages: calls on_peer_notify
+        - RESPONSE messages: ignored (handled by _deliver_message resolving futures)
+        """
+        from jarviscore.p2p.messages import MessageType
+
+        try:
+            # Skip RESPONSE messages - they should be handled by pending request futures
+            if msg.type == MessageType.RESPONSE:
+                self._logger.debug(
+                    f"[{self.role}] Ignoring orphaned RESPONSE from {msg.sender} (no pending request)"
+                )
+                return
+            
+            # Check if this is a request (expects response)
+            is_request = (
+                msg.type == MessageType.REQUEST or
+                getattr(msg, 'is_request', False)
+            )
+
+            if is_request:
+                # Request-response: call handler, optionally send response
+                response = await self.on_peer_request(msg)
+
+                if self.auto_respond and response is not None:
+                    await self.peers.respond(msg, response)
+                    self._logger.debug(
+                        f"[{self.role}] Sent response to {msg.sender}"
+                    )
+            else:
+                # Notification: fire-and-forget
+                await self.on_peer_notify(msg)
+
+        except Exception as e:
+            self._logger.error(
+                f"[{self.role}] Error handling message from {msg.sender}: {e}"
+            )
+            await self.on_error(e, msg)
+
+    # ─────────────────────────────────────────────────────────────────
+    # Message Handlers - Override in your agent
+    # ─────────────────────────────────────────────────────────────────
+
+    async def on_peer_request(self, msg) -> Any:
+        """
+        Handle incoming peer request.
+
+        Override to process request-response messages from other agents.
+        The return value is automatically sent as response (if auto_respond=True).
+
+        Args:
+            msg: IncomingMessage with:
+                - msg.sender: Sender agent ID or role
+                - msg.data: Request payload (dict)
+                - msg.correlation_id: For response matching (handled automatically)
+
+        Returns:
+            Response data (dict) to send back to the requester.
+            Return None to skip sending a response.
+
+        Example:
+            async def on_peer_request(self, msg):
+                action = msg.data.get("action")
+
+                if action == "analyze":
+                    result = await self.analyze(msg.data["payload"])
+                    return {"status": "success", "result": result}
+
+                elif action == "status":
+                    return {"status": "ok", "queue_size": self.queue_size}
+
+                return {"status": "error", "message": f"Unknown action: {action}"}
+        """
+        return None
+
+    async def on_peer_notify(self, msg) -> None:
+        """
+        Handle incoming peer notification.
+
+        Override to process fire-and-forget messages from other agents.
+        No response is expected or sent.
+
+        Args:
+            msg: IncomingMessage with:
+                - msg.sender: Sender agent ID or role
+                - msg.data: Notification payload (dict)
+
+        Example:
+            async def on_peer_notify(self, msg):
+                event = msg.data.get("event")
+
+                if event == "task_complete":
+                    await self.update_dashboard(msg.data)
+                    self._logger.info(f"Task completed by {msg.sender}")
+
+                elif event == "peer_joined":
+                    self._logger.info(f"New peer in mesh: {msg.data.get('role')}")
+        """
+        self._logger.debug(
+            f"[{self.role}] Received notify from {msg.sender}: "
+            f"{list(msg.data.keys()) if isinstance(msg.data, dict) else 'data'}"
         )
 
+    async def on_error(self, error: Exception, msg=None) -> None:
+        """
+        Handle errors during message processing.
+
+        Override to customize error handling (logging, alerting, metrics, etc.)
+        Default implementation logs the error and continues processing.
+
+        Args:
+            error: The exception that occurred
+            msg: The message being processed when error occurred (may be None)
+
+        Example:
+            async def on_error(self, error, msg):
+                # Log with context
+                self._logger.error(
+                    f"Error processing message: {error}",
+                    extra={"sender": msg.sender if msg else None}
+                )
+
+                # Send to error tracking service
+                await self.error_tracker.capture(error, context={"msg": msg})
+
+                # Optionally notify the sender of failure
+                if msg and msg.correlation_id:
+                    await self.peers.respond(msg, {
+                        "status": "error",
+                        "error": str(error)
+                    })
+        """
+        if msg:
+            self._logger.error(
+                f"[{self.role}] Error processing message from {msg.sender}: {error}"
+            )
+        else:
+            self._logger.error(f"[{self.role}] Error in listener loop: {error}")
+
+    # ─────────────────────────────────────────────────────────────────
+    # Workflow Compatibility
+    # ─────────────────────────────────────────────────────────────────
+
     async def execute_task(self, task: Dict[str, Any]) -> Dict[str, Any]:
         """
-        User implements this with custom execution logic.
+        Execute a task (for workflow/distributed modes).
 
-        DAY 1: Raises NotImplementedError (user must override)
-        DAY 5+: Full examples provided
+        Default: Delegates to on_peer_request via synthetic message.
+        Override for custom workflow logic.
 
         Args:
-            task: Task specification
+            task: Task specification dict
 
         Returns:
-            Result dictionary with at least:
-            - status: "success" or "failure"
-            - output: Task result
-            - error (optional): Error message if failed
-            - tokens_used (optional): For cost tracking
-            - cost_usd (optional): For cost tracking
+            Result dict with status and output
 
         Raises:
-            NotImplementedError: User must override this method
-
-        Example:
-            async def execute_task(self, task):
-                result = await self.my_framework.run(task)
-                return {
-                    "status": "success",
-                    "output": result,
-                    "tokens_used": 1000,  # Optional
-                    "cost_usd": 0.002     # Optional
-                }
+            NotImplementedError: If on_peer_request returns None and
+                                 execute_task is not overridden
         """
+        from jarviscore.p2p.messages import IncomingMessage, MessageType
+
+        # Create a synthetic message to pass to the handler
+        synthetic_msg = IncomingMessage(
+            sender="workflow",
+            sender_node="local",
+            type=MessageType.REQUEST,
+            data=task,
+            correlation_id=None,
+            timestamp=0
+        )
+
+        result = await self.on_peer_request(synthetic_msg)
+
+        if result is not None:
+            return {"status": "success", "output": result}
+
         raise NotImplementedError(
-            f"{self.__class__.__name__} must implement execute_task()\n\n"
+            f"{self.__class__.__name__} must implement on_peer_request() or execute_task()\n\n"
             f"Example:\n"
-            f"    async def execute_task(self, task):\n"
-            f"        result = await self.my_framework.run(task['task'])\n"
-            f"        return {{'status': 'success', 'output': result}}\n"
+            f"    async def on_peer_request(self, msg):\n"
+            f"        result = await self.process(msg.data)\n"
+            f"        return {{'status': 'success', 'result': result}}\n"
         )

jarviscore/testing/__init__.py

@@ -0,0 +1,35 @@
+"""
+Testing utilities for JarvisCore.
+
+Provides mock implementations for unit testing agents without
+requiring real P2P infrastructure or network connections.
+
+Example:
+    from jarviscore.testing import MockMesh, MockPeerClient
+
+    # Using MockMesh for full integration testing
+    mesh = MockMesh()
+    mesh.add(MyAgent)
+    await mesh.start()
+
+    agent = mesh.get_agent("my_role")
+    agent.peers.set_mock_response("analyst", {"result": "test"})
+
+    # Test and verify
+    await agent.process_request(...)
+    agent.peers.assert_requested("analyst")
+
+    # Using MockPeerClient for unit testing
+    agent = MyAgent()
+    agent.peers = MockPeerClient(mock_peers=[
+        {"role": "analyst", "capabilities": ["analysis"]}
+    ])
+"""
+
+from .mocks import MockMesh, MockPeerClient, MockPeerInfo
+
+__all__ = [
+    'MockMesh',
+    'MockPeerClient',
+    'MockPeerInfo',
+]