context

context

Context injection — retrieve relevant memory and inject into prompts.

Classes

ContextConfig dataclass

ContextConfig(
    enabled: bool = True,
    top_k: int = 5,
    min_score: float = 0.0,
    max_context_tokens: int = 2048,
)

Controls how retrieved context is injected into prompts.

Functions:

format_context

format_context(results: List[RetrievalResult]) -> str

Format retrieval results into a context block.

Each result is prefixed with its source attribution.

Source code in src/diapason/tools/storage/context.py

def format_context(results: List[RetrievalResult]) -> str:
    """Format retrieval results into a context block.

    Each result is prefixed with its source attribution.
    """
    if not results:
        return ""

    lines = []
    for r in results:
        source_tag = f"[Source: {r.source}]" if r.source else ""
        if source_tag:
            lines.append(f"{source_tag} {r.content}")
        else:
            lines.append(r.content)

    return "\n\n".join(lines)

build_context_message

build_context_message(
    results: List[RetrievalResult],
) -> Message

Create a system message with formatted context.

Source code in src/diapason/tools/storage/context.py

def build_context_message(
    results: List[RetrievalResult],
) -> Message:
    """Create a system message with formatted context."""
    context_text = format_context(results)
    content = (
        "The following context was retrieved from the knowledge"
        " base. Use it to inform your response, citing sources"
        " where applicable:\n\n" + context_text
    )
    return Message(role=Role.SYSTEM, content=content)

inject_context

inject_context(
    query: str,
    messages: List[Message],
    backend: MemoryBackend,
    *,
    config: Optional[ContextConfig] = None,
) -> List[Message]

Retrieve relevant context and prepend it to messages.

Returns a new list — the original list is not mutated. If no results pass the score threshold, returns the original messages unchanged.

PARAMETER	DESCRIPTION
query	The user query to search for. TYPE: str
messages	The existing message list. TYPE: List[Message]
backend	The memory backend to search. TYPE: MemoryBackend
config	Context injection settings (uses defaults if None). TYPE: Optional[ContextConfig] DEFAULT: None

Source code in src/diapason/tools/storage/context.py

def inject_context(
    query: str,
    messages: List[Message],
    backend: MemoryBackend,
    *,
    config: Optional[ContextConfig] = None,
) -> List[Message]:
    """Retrieve relevant context and prepend it to *messages*.

    Returns a **new** list — the original list is not mutated.
    If no results pass the score threshold, returns the original
    messages unchanged.

    Parameters
    ----------
    query:
        The user query to search for.
    messages:
        The existing message list.
    backend:
        The memory backend to search.
    config:
        Context injection settings (uses defaults if ``None``).
    """
    cfg = config or ContextConfig()
    if not cfg.enabled:
        return messages

    results = backend.retrieve(query, top_k=cfg.top_k)

    # Filter by minimum score
    results = [r for r in results if r.score >= cfg.min_score]

    if not results:
        return messages

    # Truncate to max_context_tokens
    truncated: List[RetrievalResult] = []
    total_tokens = 0
    for r in results:
        tokens = _count_tokens(r.content)
        if total_tokens + tokens > cfg.max_context_tokens:
            break
        truncated.append(r)
        total_tokens += tokens

    if not truncated:
        return messages

    # Publish event
    bus = get_event_bus()
    bus.publish(
        EventType.MEMORY_RETRIEVE,
        {
            "context_injection": True,
            "query": query,
            "num_results": len(truncated),
            "total_tokens": total_tokens,
        },
    )

    # Build context message and prepend
    ctx_msg = build_context_message(truncated)
    return [ctx_msg] + list(messages)