PRD-107 — Context Interface Abstraction | Design Doc's | Automatos AI Docs

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import StrEnum
from typing import Any, Optional


class ContextModeType(StrEnum):
    """Port-level mode names. Decoupled from ContextMode enum in modes.py."""
    CHATBOT = "chatbot"
    TASK_EXECUTION = "task_execution"
    HEARTBEAT_ORCHESTRATOR = "heartbeat_orchestrator"
    HEARTBEAT_AGENT = "heartbeat_agent"
    RECIPE = "recipe"
    ROUTER = "router"
    ORCHESTRATOR_STAGE = "orchestrator_stage"
    NL2SQL = "nl2sql"
    COORDINATOR = "coordinator"        # PRD-102
    VERIFIER = "verifier"              # PRD-103


@dataclass(frozen=True)
class AgentContext:
    """Immutable domain object — the coordinator's view of context.

    Any backend must produce this shape. Replaces ContextResult as the
    public contract. ContextResult becomes adapter-internal.
    """
    system_prompt: str
    messages: list[dict[str, Any]] = field(default_factory=list)
    tools: list[dict[str, Any]] = field(default_factory=list)
    tool_choice: str = "auto"

    # Metadata for observability and telemetry (PRD-106)
    mode: str = ""
    sections_included: list[str] = field(default_factory=list)
    sections_trimmed: list[str] = field(default_factory=list)
    token_estimate: int = 0
    token_budget: int = 0
    preparation_time_ms: float = 0.0
    memory_context: Optional[str] = None
    user_name: Optional[str] = None
    metadata: dict[str, Any] = field(default_factory=dict)


class ContextProvider(ABC):
    """Driven port: consumers call this to build LLM-ready context.

    Phase 2: DefaultContextProvider wraps ContextService
    Phase 3: NeuralFieldContextProvider queries neural field
    """

    @abstractmethod
    async def build_context(
        self,
        mode: str,
        agent: Any,
        workspace_id: str,
        messages: Optional[list[dict]] = None,
        task_description: Optional[str] = None,
        **kwargs: Any,
    ) -> AgentContext:
        """Build complete LLM-ready context for the given mode.

        Args:
            mode: Context mode name (ContextModeType value or string).
            agent: Agent record (SQLAlchemy model or dict).
            workspace_id: UUID string for workspace scoping.
            messages: Conversation history for chatbot modes.
            task_description: Task text for task execution modes.
            **kwargs: Backend-specific overrides (recipe_step, etc.).

        Returns:
            Immutable AgentContext ready for LLM call.
        """
        ...

class SharedContextPort(ABC):
    """Cross-agent context sharing within a mission.

    Phase 2: RedisSharedContext wraps SharedContextManager
    Phase 3: NeuralFieldSharedContext injects/queries neural field
    """

    @abstractmethod
    async def inject(
        self,
        context_id: str,
        key: str,
        value: Any,
        agent_id: int,
        strength: float = 1.0,
    ) -> None:
        """Inject a finding/result into shared context.

        Phase 2: Redis HSET with 2h TTL.
        Phase 3: Neural field injection with strength parameter.
        """
        ...

    @abstractmethod
    async def query(
        self,
        context_id: str,
        query: str,
        agent_id: int,
        top_k: int = 10,
    ) -> list[dict[str, Any]]:
        """Query shared context relevant to the query string.

        Phase 2: Key-value lookup from Redis/in-memory dict.
        Phase 3: Resonance-based retrieval from neural field.

        Returns list of {key, value, score?} dicts, ranked by relevance.
        """
        ...

    @abstractmethod
    async def create_context(
        self,
        team_agent_ids: list[int],
        initial_data: Optional[dict] = None,
    ) -> str:
        """Create a scoped shared context for a team of agents.

        Returns: context_id (opaque string).
        """
        ...

    @abstractmethod
    async def destroy_context(self, context_id: str) -> None:
        """Destroy a shared context and release resources."""
        ...

# orchestrator/modules/context/adapters/default.py

from core.ports.context import ContextProvider, AgentContext
from modules.context.service import ContextService
from modules.context.modes import ContextMode


class DefaultContextProvider(ContextProvider):
    """Phase 2 adapter: delegates to existing ContextService.

    Maps AgentContext (domain contract) ← ContextResult (internal).
    All existing behavior, sections, budget management, tool loading
    is preserved. This is a thin translation layer.
    """

    def __init__(self, db_session=None):
        self._service = ContextService(db_session)

    async def build_context(
        self,
        mode: str,
        agent,
        workspace_id: str,
        messages=None,
        task_description=None,
        **kwargs,
    ) -> AgentContext:
        # Map string mode → ContextMode enum (backwards compat)
        try:
            context_mode = ContextMode(mode)
        except ValueError:
            # New modes (coordinator, verifier) not yet in enum
            # Fall back to TASK_EXECUTION as default
            context_mode = ContextMode.TASK_EXECUTION

        result = await self._service.build_context(
            mode=context_mode,
            agent=agent,
            workspace_id=workspace_id,
            messages=messages,
            task_description=task_description,
            **kwargs,
        )

        # Map ContextResult → AgentContext
        return AgentContext(
            system_prompt=result.system_prompt,
            messages=result.messages,
            tools=result.tools,
            tool_choice=result.tool_choice,
            mode=result.mode,
            sections_included=result.sections_included,
            sections_trimmed=result.sections_trimmed,
            token_estimate=result.token_estimate,
            token_budget=result.token_budget,
            preparation_time_ms=result.preparation_time_ms,
            memory_context=result.memory_context,
            user_name=result.user_name,
        )

# orchestrator/modules/context/adapters/shared_redis.py

from core.ports.context import SharedContextPort
from modules.agents.communication.inter_agent import SharedContextManager


class RedisSharedContext(SharedContextPort):
    """Phase 2 adapter: Redis + in-memory shared context.

    Wraps SharedContextManager. 2h TTL via Redis EXPIRE.
    query() returns all entries (no semantic ranking in Phase 2).
    """

    def __init__(self, db_session=None):
        self._manager = SharedContextManager(db_session)

    async def inject(self, context_id, key, value, agent_id, strength=1.0):
        # strength is ignored in Phase 2 — Redis doesn't support weighted entries
        await self._manager.update_shared_context(
            context_id=context_id,
            agent=agent_id,
            updates={key: value},
            merge_strategy="append",
        )

    async def query(self, context_id, query, agent_id, top_k=10):
        ctx = self._manager.contexts.get(context_id)
        if not ctx:
            # Try Redis fallback
            if self._manager._redis:
                import json as _json
                redis_key = f"swarm:{context_id}:shared_context"
                raw = self._manager._redis.hget(redis_key, "context_data")
                if raw:
                    data = _json.loads(raw)
                    return [{"key": k, "value": v, "score": 1.0} for k, v in data.items()][:top_k]
            return []

        # Phase 2: return all entries, no semantic ranking
        return [
            {"key": k, "value": v, "score": 1.0}
            for k, v in ctx.context_data.items()
        ][:top_k]

    async def create_context(self, team_agent_ids, initial_data=None):
        ctx = await self._manager.create_shared_context(
            team=team_agent_ids,
            initial_context=initial_data or {},
        )
        return ctx.id

    async def destroy_context(self, context_id):
        # Remove from in-memory dict
        self._manager.contexts.pop(context_id, None)
        # Remove from Redis
        if self._manager._redis:
            redis_key = f"swarm:{context_id}:shared_context"
            self._manager._redis.delete(redis_key)

# orchestrator/modules/context/adapters/memory.py

from core.ports.context import ContextProvider, SharedContextPort, AgentContext


class InMemoryContextProvider(ContextProvider):
    """Test fake: returns configurable context without DB or LLM calls."""

    def __init__(self, system_prompt: str = "You are a test agent.", tools: list = None):
        self._system_prompt = system_prompt
        self._tools = tools or []

    async def build_context(self, mode, agent, workspace_id,
                            messages=None, task_description=None, **kwargs):
        return AgentContext(
            system_prompt=self._system_prompt,
            tools=self._tools,
            messages=messages or [],
            mode=mode,
            token_estimate=100,
            token_budget=4096,
        )


class InMemorySharedContext(SharedContextPort):
    """Test fake: in-memory dict-based shared context."""

    def __init__(self):
        self._contexts: dict[str, dict] = {}

    async def inject(self, context_id, key, value, agent_id, strength=1.0):
        if context_id not in self._contexts:
            raise ValueError(f"Context {context_id} not found")
        self._contexts[context_id][key] = value

    async def query(self, context_id, query, agent_id, top_k=10):
        data = self._contexts.get(context_id, {})
        return [{"key": k, "value": v, "score": 1.0} for k, v in data.items()][:top_k]

    async def create_context(self, team_agent_ids, initial_data=None):
        import uuid
        ctx_id = str(uuid.uuid4())
        self._contexts[ctx_id] = initial_data or {}
        return ctx_id

    async def destroy_context(self, context_id):
        self._contexts.pop(context_id, None)

# orchestrator/core/factories/context.py

from core.ports.context import ContextProvider, SharedContextPort
from config import config


def create_context_provider(
    db_session=None,
    backend: str = None,
) -> ContextProvider:
    """Create the context provider based on config.

    Args:
        db_session: SQLAlchemy session for DB-backed adapters.
        backend: Override CONTEXT_BACKEND config. One of:
            "default" — Phase 2 message-passing (ContextService)
            "neural_field" — Phase 3 neural field (PRD-108+)

    Returns:
        ContextProvider instance.
    """
    backend = backend or getattr(config, 'CONTEXT_BACKEND', 'default')

    if backend == "neural_field":
        from modules.context.adapters.neural_field import NeuralFieldContextProvider
        return NeuralFieldContextProvider()
    else:
        from modules.context.adapters.default import DefaultContextProvider
        return DefaultContextProvider(db_session)


def create_shared_context(
    db_session=None,
    backend: str = None,
) -> SharedContextPort:
    """Create the shared context port based on config."""
    backend = backend or getattr(config, 'CONTEXT_BACKEND', 'default')

    if backend == "neural_field":
        from modules.context.adapters.neural_field import NeuralFieldSharedContext
        return NeuralFieldSharedContext()
    else:
        from modules.context.adapters.shared_redis import RedisSharedContext
        return RedisSharedContext(db_session)

# In CoordinatorService
mission_backend = mission_config.get("context_backend", None)  # None = use default
context_provider = create_context_provider(db, backend=mission_backend)
shared_context = create_shared_context(db, backend=mission_backend)

# orchestrator/modules/context/sections/mission_context.py

from modules.context.sections import BaseSection, SectionContext


class MissionContextSection(BaseSection):
    """Injects shared mission context into agent's system prompt.

    Activated when kwargs contains 'mission_context_id' and 'shared_context_port'.
    Priority 3: after identity/task, before skills/plugins.
    """
    name = "mission_context"
    priority = 3
    max_tokens = 2000

    async def render(self, ctx: SectionContext) -> str:
        mission_context_id = ctx.kwargs.get("mission_context_id")
        shared_port = ctx.kwargs.get("shared_context_port")

        if not mission_context_id or not shared_port:
            return ""

        entries = await shared_port.query(
            context_id=mission_context_id,
            query=ctx.task_description or "",
            agent_id=ctx.agent.id if hasattr(ctx.agent, 'id') else 0,
            top_k=5,
        )

        if not entries:
            return ""

        lines = ["## Mission Shared Context", ""]
        for entry in entries:
            lines.append(f"**{entry['key']}:** {entry['value']}")
            lines.append("")

        return "\n".join(lines)

# Before:
class SmartChatOrchestrator:
    def __init__(self, db_session, ...):
        self._db_session = db_session

    async def _build_context(self):
        context = await ContextService(self._db_session).build_context(
            mode=ContextMode.CHATBOT, ...
        )

# After:
class SmartChatOrchestrator:
    def __init__(self, db_session, context_provider: ContextProvider = None, ...):
        self._db_session = db_session
        self._context = context_provider or create_context_provider(db_session)

    async def _build_context(self):
        context = await self._context.build_context(
            mode="chatbot", ...
        )

# orchestrator/modules/context/adapters/neural_field.py

from core.ports.context import ContextProvider, SharedContextPort, AgentContext


class NeuralFieldContextProvider(ContextProvider):
    """Phase 3 adapter: builds context from neural field resonance.

    PRD-108 implements the prototype. This stub documents the contract.
    """

    async def build_context(self, mode, agent, workspace_id,
                            messages=None, task_description=None, **kwargs):
        raise NotImplementedError(
            "NeuralFieldContextProvider requires PRD-108 implementation. "
            "Set CONTEXT_BACKEND=default to use Phase 2 message-passing."
        )


class NeuralFieldSharedContext(SharedContextPort):
    """Phase 3 adapter: shared context via neural field injection.

    PRD-108 implements the prototype. This stub documents the contract.
    """

    async def inject(self, context_id, key, value, agent_id, strength=1.0):
        raise NotImplementedError("Requires PRD-108 implementation.")

    async def query(self, context_id, query, agent_id, top_k=10):
        raise NotImplementedError("Requires PRD-108 implementation.")

    async def create_context(self, team_agent_ids, initial_data=None):
        raise NotImplementedError("Requires PRD-108 implementation.")

    async def destroy_context(self, context_id):
        raise NotImplementedError("Requires PRD-108 implementation.")

orchestrator/
├── core/
│   ├── ports/
│   │   ├── __init__.py              # Public exports
│   │   └── context.py               # ContextProvider, SharedContextPort, AgentContext, ContextModeType
│   └── factories/
│       └── context.py               # create_context_provider(), create_shared_context()
└── modules/
    └── context/
        ├── adapters/
        │   ├── __init__.py           # Adapter exports
        │   ├── default.py            # DefaultContextProvider (wraps ContextService)
        │   ├── shared_redis.py       # RedisSharedContext (wraps SharedContextManager)
        │   ├── memory.py             # InMemoryContextProvider, InMemorySharedContext (test fakes)
        │   └── neural_field.py       # NeuralFieldContextProvider stub (PRD-108+)
        └── sections/
            └── mission_context.py    # MissionContextSection

Coordinator (PRD-102)
    │
    │ imports
    ▼
core/ports/context.py  ←── THE BOUNDARY ──→  modules/context/adapters/
    │                                              │
    │ defines:                                     │ implements:
    │  - ContextProvider ABC                       │  - DefaultContextProvider
    │  - SharedContextPort ABC                     │  - RedisSharedContext
    │  - AgentContext dataclass                    │  - NeuralFieldContextProvider
    │  - ContextModeType StrEnum                   │
    │                                              │ wraps:
    │                                              │  - ContextService (unchanged)
    │                                              │  - SharedContextManager (unchanged)
    ▼
core/factories/context.py
    │
    │ selects adapter at startup
    │ based on config.CONTEXT_BACKEND
    ▼
CoordinatorService(context=adapter)