Newer
Older
navi-1 / navi / workers / base.py
"""Base classes for post-response background workers."""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from navi.core.events import AgentEvent

if TYPE_CHECKING:
    from navi.core.session import Session, SessionStore
    from navi.llm.base import LLMBackend


@dataclass
class WorkerContext:
    """Runtime data passed to every worker after a response completes."""

    session_id: str
    context_tokens: int | None        # tokens used in last turn (from Ollama)
    max_context_tokens: int           # ollama_num_ctx
    llm: LLMBackend
    model: str
    temperature: float
    session_store: SessionStore


@dataclass
class WorkerResult:
    """Returned by a worker. `events` will be yielded to the WebSocket client."""

    events: list[AgentEvent] = field(default_factory=list)


class Worker(ABC):
    """
    A post-response background task.

    Workers run sequentially after `StreamEnd` is yielded.
    Each worker may modify `session` (e.g. compress context) and return
    events to forward to the client. Failures are isolated — one broken
    worker does not block others.
    """

    @abstractmethod
    async def run(self, session: Session, ctx: WorkerContext) -> WorkerResult:
        """Execute the worker. May mutate session and save via ctx.session_store."""