genxai-framework 0.1.0__py3-none-any.whl

genxai/core/communication/__init__.py

@@ -0,0 +1,6 @@
+"""Communication layer for agent-to-agent messaging."""
+
+from genxai.core.communication.message_bus import MessageBus, Message
+from genxai.core.communication.protocols import CommunicationProtocol
+
+__all__ = ["MessageBus", "Message", "CommunicationProtocol"]

genxai/core/communication/collaboration.py

@@ -0,0 +1,44 @@
+"""Collaboration protocol implementations."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, List
+
+
+@dataclass
+class VotingResult:
+    winner: Any
+    counts: Dict[Any, int]
+
+
+class VotingProtocol:
+    """Simple majority voting protocol."""
+
+    async def run(self, inputs: List[Any], metadata: Dict[str, Any]) -> VotingResult:
+        counts: Dict[Any, int] = {}
+        for value in inputs:
+            counts[value] = counts.get(value, 0) + 1
+        winner = max(counts, key=counts.get)
+        return VotingResult(winner=winner, counts=counts)
+
+
+class NegotiationProtocol:
+    """Simple negotiation protocol that returns consensus if all equal."""
+
+    async def run(self, inputs: List[Any], metadata: Dict[str, Any]) -> Any:
+        if not inputs:
+            return None
+        first = inputs[0]
+        if all(value == first for value in inputs):
+            return first
+        return metadata.get("fallback")
+
+
+class AuctionProtocol:
+    """Simple auction protocol selecting max bid from inputs."""
+
+    async def run(self, inputs: List[Any], metadata: Dict[str, Any]) -> Any:
+        if not inputs:
+            return None
+        return max(inputs)

genxai/core/communication/message_bus.py

@@ -0,0 +1,192 @@
+"""Message bus for agent-to-agent communication."""
+
+from typing import Any, Dict, List, Optional, Callable
+from pydantic import BaseModel, Field, ConfigDict
+from datetime import datetime
+from collections import defaultdict
+import asyncio
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class Message(BaseModel):
+    """Message for agent communication."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    id: str
+    sender: str
+    recipient: Optional[str] = None  # None for broadcast
+    content: Any
+    message_type: str = "default"
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+    timestamp: datetime = Field(default_factory=datetime.now)
+    reply_to: Optional[str] = None
+
+
+
+class MessageBus:
+    """Central message bus for agent communication."""
+
+    def __init__(self) -> None:
+        """Initialize message bus."""
+        self._subscribers: Dict[str, List[Callable]] = defaultdict(list)
+        self._message_history: List[Message] = []
+        self._message_count = 0
+
+    async def send(self, message: Message) -> None:
+        """Send a message to a specific recipient.
+
+        Args:
+            message: Message to send
+        """
+        self._message_count += 1
+        message.id = f"msg_{self._message_count}"
+        self._message_history.append(message)
+
+        logger.info(f"Message sent: {message.sender} -> {message.recipient}")
+
+        # Deliver to recipient's subscribers
+        if message.recipient and message.recipient in self._subscribers:
+            for callback in self._subscribers[message.recipient]:
+                try:
+                    await callback(message)
+                except Exception as e:
+                    logger.error(f"Error delivering message to {message.recipient}: {e}")
+
+    async def broadcast(self, message: Message, group: Optional[str] = None) -> None:
+        """Broadcast a message to all subscribers or a group.
+
+        Args:
+            message: Message to broadcast
+            group: Optional group name to broadcast to
+        """
+        self._message_count += 1
+        message.id = f"msg_{self._message_count}"
+        message.recipient = None  # Broadcast has no specific recipient
+        self._message_history.append(message)
+
+        logger.info(f"Message broadcast from {message.sender} to group: {group or 'all'}")
+
+        # Deliver to all subscribers (or group subscribers)
+        for agent_id, callbacks in self._subscribers.items():
+            if group and not agent_id.startswith(f"{group}_"):
+                continue
+
+            for callback in callbacks:
+                try:
+                    await callback(message)
+                except Exception as e:
+                    logger.error(f"Error broadcasting to {agent_id}: {e}")
+
+    async def request_reply(
+        self, message: Message, timeout: float = 30.0
+    ) -> Optional[Message]:
+        """Send a message and wait for reply.
+
+        Args:
+            message: Message to send
+            timeout: Timeout in seconds
+
+        Returns:
+            Reply message or None if timeout
+        """
+        # Send message
+        await self.send(message)
+
+        # Wait for reply
+        reply_event = asyncio.Event()
+        reply_message: Optional[Message] = None
+
+        async def reply_handler(msg: Message) -> None:
+            nonlocal reply_message
+            if msg.reply_to == message.id:
+                reply_message = msg
+                reply_event.set()
+
+        # Subscribe to replies
+        if message.sender:
+            self.subscribe(message.sender, reply_handler)
+
+        try:
+            await asyncio.wait_for(reply_event.wait(), timeout=timeout)
+            return reply_message
+        except asyncio.TimeoutError:
+            logger.warning(f"Request-reply timeout for message {message.id}")
+            return None
+        finally:
+            if message.sender:
+                self.unsubscribe(message.sender, reply_handler)
+
+    def subscribe(self, agent_id: str, callback: Callable) -> None:
+        """Subscribe an agent to receive messages.
+
+        Args:
+            agent_id: Agent identifier
+            callback: Async callback function to handle messages
+        """
+        self._subscribers[agent_id].append(callback)
+        logger.debug(f"Agent {agent_id} subscribed to message bus")
+
+    def unsubscribe(self, agent_id: str, callback: Optional[Callable] = None) -> None:
+        """Unsubscribe an agent from messages.
+
+        Args:
+            agent_id: Agent identifier
+            callback: Specific callback to remove (None to remove all)
+        """
+        if agent_id in self._subscribers:
+            if callback:
+                self._subscribers[agent_id].remove(callback)
+            else:
+                del self._subscribers[agent_id]
+            logger.debug(f"Agent {agent_id} unsubscribed from message bus")
+
+    def get_history(
+        self, agent_id: Optional[str] = None, limit: Optional[int] = None
+    ) -> List[Message]:
+        """Get message history.
+
+        Args:
+            agent_id: Filter by agent (sender or recipient)
+            limit: Maximum number of messages to return
+
+        Returns:
+            List of messages
+        """
+        messages = self._message_history
+
+        if agent_id:
+            messages = [
+                m
+                for m in messages
+                if m.sender == agent_id or m.recipient == agent_id
+            ]
+
+        if limit:
+            messages = messages[-limit:]
+
+        return messages
+
+    def clear_history(self) -> None:
+        """Clear message history."""
+        self._message_history.clear()
+        logger.info("Message history cleared")
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get message bus statistics.
+
+        Returns:
+            Statistics dictionary
+        """
+        return {
+            "total_messages": self._message_count,
+            "history_size": len(self._message_history),
+            "subscribers": len(self._subscribers),
+            "subscriber_list": list(self._subscribers.keys()),
+        }
+
+    def __repr__(self) -> str:
+        """String representation."""
+        return f"MessageBus(messages={self._message_count}, subscribers={len(self._subscribers)})"

genxai/core/communication/protocols.py

@@ -0,0 +1,35 @@
+"""Communication protocols for agent interaction."""
+
+from enum import Enum
+from typing import Protocol as TypingProtocol, Any, Dict, List
+
+
+class CommunicationProtocol(str, Enum):
+    """Communication protocols for agents."""
+
+    POINT_TO_POINT = "point_to_point"
+    BROADCAST = "broadcast"
+    REQUEST_REPLY = "request_reply"
+    PUB_SUB = "pub_sub"
+    NEGOTIATION = "negotiation"
+    VOTING = "voting"
+    AUCTION = "auction"
+
+
+class MessageHandler(TypingProtocol):
+    """Protocol for message handlers."""
+
+    async def handle_message(self, message: Any) -> None:
+        """Handle incoming message.
+
+        Args:
+            message: Message to handle
+        """
+        ...
+
+
+class CollaborationProtocol(TypingProtocol):
+    """Protocol for collaboration strategies."""
+
+    async def run(self, inputs: List[Any], metadata: Dict[str, Any]) -> Any:
+        ...

genxai/core/execution/__init__.py

@@ -0,0 +1,22 @@
+"""Distributed execution primitives for GenXAI."""
+
+from genxai.core.execution.queue import (
+    QueueBackend,
+    QueueTask,
+    InMemoryQueueBackend,
+    WorkerQueueEngine,
+    RedisQueueBackend,
+    RQQueueBackend,
+)
+from genxai.core.execution.metadata import ExecutionRecord, ExecutionStore
+
+__all__ = [
+    "QueueBackend",
+    "QueueTask",
+    "InMemoryQueueBackend",
+    "WorkerQueueEngine",
+    "RedisQueueBackend",
+    "RQQueueBackend",
+    "ExecutionRecord",
+    "ExecutionStore",
+]

genxai/core/execution/metadata.py

@@ -0,0 +1,181 @@
+"""Execution metadata store for workflow runs."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, Optional
+import json
+import uuid
+
+
+@dataclass
+class ExecutionRecord:
+    """Represents a workflow execution record."""
+
+    run_id: str
+    workflow: str
+    status: str
+    started_at: str
+    completed_at: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    error: Optional[str] = None
+    result: Optional[Dict[str, Any]] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "run_id": self.run_id,
+            "workflow": self.workflow,
+            "status": self.status,
+            "started_at": self.started_at,
+            "completed_at": self.completed_at,
+            "metadata": self.metadata,
+            "error": self.error,
+            "result": self.result,
+        }
+
+
+class ExecutionStore:
+    """Execution store with JSON or SQL persistence support."""
+
+    def __init__(
+        self,
+        persistence_path: Optional[Path] = None,
+        sql_url: Optional[str] = None,
+    ) -> None:
+        self._records: Dict[str, ExecutionRecord] = {}
+        self._persistence_path = persistence_path
+        self._sql_url = sql_url
+        self._engine = None
+        self._table = None
+
+        if sql_url:
+            try:
+                import sqlalchemy as sa
+            except Exception as exc:
+                raise ImportError(
+                    "sqlalchemy is required for SQL persistence. Install with: pip install sqlalchemy"
+                ) from exc
+            self._engine = sa.create_engine(sql_url)
+            metadata = sa.MetaData()
+            self._table = sa.Table(
+                "genxai_executions",
+                metadata,
+                sa.Column("run_id", sa.String, primary_key=True),
+                sa.Column("workflow", sa.String, nullable=False),
+                sa.Column("status", sa.String, nullable=False),
+                sa.Column("started_at", sa.String, nullable=False),
+                sa.Column("completed_at", sa.String),
+                sa.Column("metadata", sa.JSON, nullable=True),
+                sa.Column("error", sa.Text, nullable=True),
+                sa.Column("result", sa.JSON, nullable=True),
+            )
+            metadata.create_all(self._engine)
+
+    def generate_run_id(self) -> str:
+        return str(uuid.uuid4())
+
+    def create(
+        self,
+        run_id: str,
+        workflow: str,
+        status: str,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> ExecutionRecord:
+        if run_id in self._records:
+            return self._records[run_id]
+
+        record = ExecutionRecord(
+            run_id=run_id,
+            workflow=workflow,
+            status=status,
+            started_at=datetime.now().isoformat(),
+            metadata=metadata or {},
+        )
+        self._records[run_id] = record
+        self._persist(record)
+        return record
+
+    def update(
+        self,
+        run_id: str,
+        status: Optional[str] = None,
+        error: Optional[str] = None,
+        result: Optional[Dict[str, Any]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        completed: bool = False,
+    ) -> ExecutionRecord:
+        record = self._records.get(run_id)
+        if record is None:
+            record = self.create(run_id, workflow="unknown", status="unknown")
+        if status is not None:
+            record.status = status
+        if error is not None:
+            record.error = error
+        if result is not None:
+            record.result = result
+        if metadata:
+            record.metadata.update(metadata)
+        if completed:
+            record.completed_at = datetime.now().isoformat()
+        self._persist(record)
+        return record
+
+    def get(self, run_id: str) -> Optional[ExecutionRecord]:
+        record = self._records.get(run_id)
+        if record or not self._engine or not self._table:
+            return record
+
+        import sqlalchemy as sa
+
+        with self._engine.begin() as conn:
+            stmt = sa.select(self._table).where(self._table.c.run_id == run_id)
+            row = conn.execute(stmt).mappings().first()
+            if not row:
+                return None
+            record = ExecutionRecord(
+                run_id=row["run_id"],
+                workflow=row["workflow"],
+                status=row["status"],
+                started_at=row["started_at"],
+                completed_at=row["completed_at"],
+                metadata=row["metadata"] or {},
+                error=row["error"],
+                result=row["result"],
+            )
+            self._records[run_id] = record
+            return record
+
+    def _persist(self, record: ExecutionRecord) -> None:
+        if self._engine is not None and self._table is not None:
+            import sqlalchemy as sa
+
+            payload = record.to_dict()
+            with self._engine.begin() as conn:
+                stmt = sa.select(self._table.c.run_id).where(
+                    self._table.c.run_id == record.run_id
+                )
+                exists = conn.execute(stmt).first()
+                if exists:
+                    conn.execute(
+                        self._table.update()
+                        .where(self._table.c.run_id == record.run_id)
+                        .values(**payload)
+                    )
+                else:
+                    conn.execute(self._table.insert().values(**payload))
+
+        if not self._persistence_path:
+            return
+        self._persistence_path.mkdir(parents=True, exist_ok=True)
+        path = self._persistence_path / f"execution_{record.run_id}.json"
+        path.write_text(json.dumps(record.to_dict(), indent=2, default=str))
+
+    def close(self) -> None:
+        """Dispose of SQL resources if enabled."""
+        if self._engine is not None:
+            self._engine.dispose()
+
+    def __del__(self) -> None:
+        self.close()

genxai/core/execution/queue.py

@@ -0,0 +1,201 @@
+"""Async worker queue engine for distributed execution."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+from typing import Any, Awaitable, Callable, Optional, Protocol
+import uuid
+import logging
+import json
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class QueueTask:
+    """Represents a unit of work for the worker queue."""
+
+    task_id: str
+    payload: dict[str, Any]
+    handler: Optional[Callable[[dict[str, Any]], Awaitable[Any]]]
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+class QueueBackend(Protocol):
+    """Protocol for queue backends."""
+
+    async def put(self, task: QueueTask) -> None:
+        ...
+
+    async def get(self) -> QueueTask:
+        ...
+
+    def qsize(self) -> int:
+        ...
+
+
+class InMemoryQueueBackend:
+    """In-memory asyncio queue backend."""
+
+    def __init__(self) -> None:
+        self._queue: asyncio.Queue[QueueTask] = asyncio.Queue()
+
+    async def put(self, task: QueueTask) -> None:
+        await self._queue.put(task)
+
+    async def get(self) -> QueueTask:
+        return await self._queue.get()
+
+    def qsize(self) -> int:
+        return self._queue.qsize()
+
+
+class WorkerQueueEngine:
+    """Simple async worker engine for processing queued tasks."""
+
+    def __init__(
+        self,
+        backend: Optional[QueueBackend] = None,
+        worker_count: int = 2,
+        max_retries: int = 3,
+        backoff_seconds: float = 0.5,
+        handler_registry: Optional[dict[str, Callable[[dict[str, Any]], Awaitable[Any]]]] = None,
+    ) -> None:
+        self._backend = backend or InMemoryQueueBackend()
+        self._worker_count = worker_count
+        self._max_retries = max_retries
+        self._backoff_seconds = backoff_seconds
+        self._workers: list[asyncio.Task[None]] = []
+        self._running = False
+        self._handler_registry = handler_registry or {}
+
+    def register_handler(
+        self,
+        name: str,
+        handler: Callable[[dict[str, Any]], Awaitable[Any]],
+    ) -> None:
+        """Register a handler by name for distributed queue backends."""
+        self._handler_registry[name] = handler
+
+    async def start(self) -> None:
+        if self._running:
+            return
+        self._running = True
+        for idx in range(self._worker_count):
+            worker = asyncio.create_task(self._worker_loop(idx))
+            self._workers.append(worker)
+
+    async def stop(self) -> None:
+        if not self._running:
+            return
+        self._running = False
+        for worker in self._workers:
+            worker.cancel()
+        await asyncio.gather(*self._workers, return_exceptions=True)
+        self._workers.clear()
+
+    async def enqueue(
+        self,
+        payload: dict[str, Any],
+        handler: Optional[Callable[[dict[str, Any]], Awaitable[Any]]] = None,
+        metadata: Optional[dict[str, Any]] = None,
+        run_id: Optional[str] = None,
+        handler_name: Optional[str] = None,
+    ) -> str:
+        task_id = run_id or str(uuid.uuid4())
+        if handler is None and handler_name:
+            handler = self._handler_registry.get(handler_name)
+        if handler is None:
+            raise ValueError("Handler must be provided or registered via handler_name")
+        task = QueueTask(
+            task_id=task_id,
+            payload=payload,
+            handler=handler,
+            metadata={**(metadata or {}), "handler_name": handler_name},
+        )
+        await self._backend.put(task)
+        return task_id
+
+    async def _worker_loop(self, worker_id: int) -> None:
+        while self._running:
+            try:
+                task = await self._backend.get()
+                await self._execute_with_retry(task)
+                logger.debug(
+                    "Worker %s processed task %s", worker_id, task.task_id
+                )
+            except asyncio.CancelledError:
+                break
+            except Exception as exc:
+                logger.error("Worker %s failed: %s", worker_id, exc)
+
+    async def _execute_with_retry(self, task: QueueTask) -> None:
+        handler = task.handler
+        if handler is None:
+            handler_name = task.metadata.get("handler_name")
+            handler = self._handler_registry.get(handler_name) if handler_name else None
+        if handler is None:
+            raise ValueError(f"No handler registered for task {task.task_id}")
+        attempts = 0
+        while True:
+            try:
+                await handler(task.payload)
+                return
+            except Exception as exc:
+                attempts += 1
+                if attempts > self._max_retries:
+                    raise exc
+                await asyncio.sleep(self._backoff_seconds * attempts)
+
+
+class RedisQueueBackend:
+    """Redis-backed queue backend for distributed execution.
+
+    Stores serialized QueueTask payloads in a Redis list and uses BLPOP to
+    retrieve work items.
+    """
+
+    def __init__(self, url: str, queue_name: str = "genxai:queue") -> None:
+        try:
+            import redis.asyncio as redis  # type: ignore
+        except Exception as exc:
+            raise ImportError(
+                "redis package is required for RedisQueueBackend. Install with: pip install redis"
+            ) from exc
+
+        self._redis = redis.from_url(url)
+        self._queue_name = queue_name
+
+    async def put(self, task: QueueTask) -> None:
+        payload = {
+            "task_id": task.task_id,
+            "payload": task.payload,
+            "metadata": task.metadata,
+        }
+        await self._redis.rpush(self._queue_name, json.dumps(payload))
+
+    async def get(self) -> QueueTask:
+        _, raw = await self._redis.blpop(self._queue_name)
+        data = json.loads(raw)
+        return QueueTask(
+            task_id=data["task_id"],
+            payload=data["payload"],
+            handler=None,
+            metadata=data.get("metadata", {}),
+        )
+
+    def qsize(self) -> int:
+        return int(self._redis.llen(self._queue_name))
+
+
+class RQQueueBackend:
+    """Placeholder backend for Redis/RQ integration.
+
+    This is a stub that documents the interface needed for an RQ backend.
+    """
+
+    def __init__(self) -> None:
+        raise NotImplementedError(
+            "RQQueueBackend is a stub. Implement with Redis + rq when ready."
+        )