Milestone 3.5.1 — Context Assembly Engine Skeleton (Python gRPC Server)

// packages/proto/ibex/context/v1/context.proto
syntax = "proto3";
package ibex.context.v1;

option go_package = "github.com/Rick1330/ibex-harness/packages/proto/ibex/context/v1;contextv1";
option python_package = "ibex_proto.context.v1";

import "google/protobuf/timestamp.proto";

// AssembleContext is called by the proxy on every LLM request.
// It returns the enriched messages array with memories and directive injected.
service ContextAssemblyService {
    rpc AssembleContext(AssembleContextRequest) returns (AssembleContextResponse);
}

message AssembleContextRequest {
    string org_id      = 1; // (required) UUID string
    string agent_id    = 2; // (required) UUID string
    string session_id  = 3; // (optional) UUID string; empty if new session
    string model       = 4; // (required) e.g. "gpt-4o", "gpt-4o-mini"
    string request_id  = 5; // (required) for tracing/logging

    // The original messages from the client (before injection).
    // The service returns an enriched version of this array.
    repeated Message messages = 6;
}

message AssembleContextResponse {
    // Enriched messages ready to send to the LLM provider.
    // The service may add, reorder, or modify messages.
    // The client's original messages are always preserved in order;
    // only injected content (directive, memories) is prepended.
    repeated Message messages = 1;

    // Metadata about what was injected (for tracing and dashboard display).
    InjectionMetadata metadata = 2;
}

message Message {
    string role    = 1; // "system", "user", "assistant"
    string content = 2;
}

message InjectionMetadata {
    bool   directive_injected    = 1;
    int32  memories_injected     = 2;  // count of memories included
    int32  memories_available    = 3;  // count of memories found (before token truncation)
    int32  total_tokens_injected = 4;  // approximate token count of injected content
    int32  context_window_used   = 5;  // tokens used as % of context window
    bool   was_truncated         = 6;  // true if memories were cut to fit token budget
    string fallback_reason       = 7;  // non-empty if assembly fell back to directive-only
    repeated string memory_ids   = 8;  // UUIDs of injected memories (for hot cache score update)
}

services/context/
  pyproject.toml
  Dockerfile
  .env.example
  src/
    context/
      __init__.py
      app.py             # gRPC server startup
      settings.py        # pydantic-settings
      grpc_server.py     # ContextAssemblyServicer implementation
      services/
        budget.py          # Token budget calculator (3.5.2)
        retrieval.py       # Parallel retrieval orchestrator (3.5.3)
        scorer.py          # Composite scoring (3.5.4)
        packer.py          # Greedy knapsack (3.5.5)
        formatter.py       # Context formatting and injection (3.5.6)
      clients/
        memory_client.py   # HTTP client to memory service
        embedder_client.py # HTTP client to embedding service
        redis_client.py    # Direct Redis access for hot cache + directive
      proto/
        # generated by make proto-gen
  tests/

# src/context/grpc_server.py
from __future__ import annotations
 
import asyncio
import time
from uuid import UUID
 
import grpclib.server
from ibex_proto.context.v1 import (
    ContextAssemblyBase,
    AssembleContextRequest,
    AssembleContextResponse,
    InjectionMetadata,
    Message,
)
 
from context.services.budget import TokenBudgetCalculator
from context.services.retrieval import ContextRetriever
from context.services.scorer import MemoryScorer
from context.services.packer import GreedyKnapsackPacker
from context.services.formatter import ContextFormatter
from context import logger
 
RETRIEVAL_DEADLINE_SECONDS = 0.040  # 40ms hard deadline on parallel retrieval
ASSEMBLY_DEADLINE_SECONDS  = 0.048  # 48ms total — gives 2ms margin before proxy times out
 
class ContextAssemblyServicer(ContextAssemblyBase):
    """
    gRPC servicer for context assembly.
 
    Graceful degradation contract:
    - If ANY part of assembly exceeds ASSEMBLY_DEADLINE_SECONDS:
      return directive-only context with metadata.fallback_reason set.
    - If directive retrieval fails: return original messages unchanged.
    - If memory retrieval partially fails: use whatever was retrieved in time.
 
    The proxy MUST enforce its own timeout (45ms) independently.
    Never trust the context service to stay within budget.
    """
 
    def __init__(
        self,
        retriever: ContextRetriever,
        budget_calc: TokenBudgetCalculator,
        scorer: MemoryScorer,
        packer: GreedyKnapsackPacker,
        formatter: ContextFormatter,
    ) -> None:
        self._retriever  = retriever
        self._budget     = budget_calc
        self._scorer     = scorer
        self._packer     = packer
        self._formatter  = formatter
 
    async def assemble_context(
        self,
        stream: grpclib.server.Stream[AssembleContextRequest, AssembleContextResponse],
    ) -> None:
        request: AssembleContextRequest = await stream.recv_message()
        start = time.monotonic()
 
        try:
            response = await asyncio.wait_for(
                self._assemble(request),
                timeout=ASSEMBLY_DEADLINE_SECONDS,
            )
        except asyncio.TimeoutError:
            logger.WarnCtx(
                f"context assembly timeout for agent={request.agent_id} "
                f"request_id={request.request_id}"
            )
            response = self._directive_only_fallback(
                request, reason="assembly_timeout"
            )
        except Exception as exc:
            logger.ErrorCtx(f"context assembly error: {exc}", request_id=request.request_id)
            response = self._directive_only_fallback(
                request, reason=f"assembly_error:{type(exc).__name__}"
            )
 
        elapsed_ms = int((time.monotonic() - start) * 1000)
        # Emit latency metric regardless of success/fallback
        metrics.context_assembly_duration.observe(elapsed_ms / 1000)
 
        await stream.send_message(response)
 
    async def _assemble(self, req: AssembleContextRequest) -> AssembleContextResponse:
        org_id   = UUID(req.org_id)
        agent_id = UUID(req.agent_id)
 
        # Step 1: Calculate token budget for this model
        budget = self._budget.calculate(model=req.model, messages=req.messages)
 
        # Step 2: Parallel retrieval (40ms deadline)
        retrieval = await asyncio.wait_for(
            self._retriever.retrieve(
                org_id=org_id, agent_id=agent_id,
                session_id=req.session_id or None,
                messages=req.messages,
                budget=budget,
            ),
            timeout=RETRIEVAL_DEADLINE_SECONDS,
        )
 
        # Step 3: Score and rank candidate memories
        scored = self._scorer.score(
            memories=retrieval.memories,
            query_embedding=retrieval.query_embedding,
        )
 
        # Step 4: Greedy knapsack — pack memories into token budget
        packed = self._packer.pack(
            scored_memories=scored,
            token_budget=budget.memory_tokens,
        )
 
        # Step 5: Format and inject into messages array
        enriched_messages, meta = self._formatter.inject(
            original_messages=req.messages,
            directive=retrieval.directive,
            packed_memories=packed,
            budget=budget,
        )
 
        return AssembleContextResponse(
            messages=enriched_messages,
            metadata=meta,
        )
 
    @staticmethod
    def _directive_only_fallback(
        req: AssembleContextRequest, reason: str
    ) -> AssembleContextResponse:
        """Return the original messages unchanged, or with directive prepended if available."""
        return AssembleContextResponse(
            messages=list(req.messages),
            metadata=InjectionMetadata(
                directive_injected=False,
                memories_injected=0,
                fallback_reason=reason,
            ),
        )

// services/proxy/internal/context/client.go
// Package context provides a gRPC client for the context assembly service.
package context
 
import (
    "context"
    "fmt"
    "time"
 
    "google.golang.org/grpc"
    contextv1 "github.com/Rick1330/ibex-harness/packages/proto/ibex/context/v1"
)
 
const assemblyTimeout = 45 * time.Millisecond // 5ms margin before the service's 48ms limit
 
// Client wraps the context assembly gRPC client with timeout enforcement.
type Client struct {
    inner  contextv1.ContextAssemblyServiceClient
    tracer trace.Tracer
    log    *logger.Logger
}
 
// AssembleContext enriches the messages array with the agent's memories and directive.
// On timeout or error, returns the original messages with a fallback flag set.
// This function NEVER fails the LLM request — context assembly is a quality enhancement, not a gate.
func (c *Client) AssembleContext(
    ctx context.Context,
    req *contextv1.AssembleContextRequest,
) (*contextv1.AssembleContextResponse, error) {
    ctx, cancel := context.WithTimeout(ctx, assemblyTimeout)
    defer cancel()
 
    ctx, span := c.tracer.Start(ctx, "ContextAssembly.AssembleContext")
    defer span.End()
 
    resp, err := c.inner.AssembleContext(ctx, req)
    if err != nil {
        // Do NOT propagate: return original messages, log warning
        c.log.WarnCtx(ctx, "context assembly failed; proceeding without memory injection",
            "error", err, "agent_id", req.AgentId)
        span.SetStatus(codes.Error, err.Error())
        return &contextv1.AssembleContextResponse{
            Messages: req.Messages,
            Metadata: &contextv1.InjectionMetadata{
                FallbackReason: fmt.Sprintf("grpc_error:%T", err),
            },
        }, nil
    }
    return resp, nil
}