lorebase-guide 0.2.3__py3-none-any.whl

lorebase_guide/__init__.py

@@ -0,0 +1,10 @@
+"""Lorebase Guide: optional exploration assistant and query embeddings over Keeper.
+
+Importing the package registers the `guide` settings section with `lorebase_platform`.
+"""
+
+from lorebase_guide.settings import GuideSettings, guide_settings
+
+__version__ = "0.2.3"
+
+__all__ = ["GuideSettings", "guide_settings"]

lorebase_guide/__main__.py

@@ -0,0 +1,93 @@
+"""Run Guide: `lorebase-guide` or `python -m lorebase_guide`.
+
+Host, port and everything else come from the `guide` and `inference` settings.
+
+lorebase-guide                 serve until stopped
+lorebase-guide modes clear     forget the structured output modes remembered per chat endpoint
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import sys
+from collections.abc import Sequence
+from pathlib import Path
+
+import uvicorn
+
+from lorebase_guide.app import create_app
+from lorebase_guide.modes import FileModeMemory, mode_memory_path
+from lorebase_guide.settings import guide_settings
+from lorebase_inference import inference_settings
+from lorebase_inference.chat import use_mode_memory
+from lorebase_platform.logging import configure_logging
+from lorebase_platform.settings import ConfigError, load_settings
+from lorebase_platform.telemetry import setup_telemetry
+
+SERVICE_NAME = "guide"
+
+
+def _parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="lorebase-guide", description=__doc__.split("\n")[0])
+    commands = parser.add_subparsers(dest="command")
+    commands.add_parser("run", help="Serve until stopped (default).")
+    modes = commands.add_parser("modes", help="Structured output modes of the chat endpoints.")
+    modes.add_argument("action", choices=["clear"], help="Forget every remembered mode.")
+    return parser
+
+
+def _clear(path: Path) -> int:
+    try:
+        return FileModeMemory(path).clear()
+    except OSError as exc:
+        sys.exit(f"cannot clear {path}: {exc}")
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    """Load settings, configure logging and telemetry, and serve until stopped (or run the
+    command)."""
+    args = _parser().parse_args(argv)
+    try:
+        settings = load_settings()
+        guide = guide_settings(settings)
+        inference = inference_settings(settings)
+    except ConfigError as exc:
+        sys.exit(str(exc))
+    if args.command == "modes":
+        path = mode_memory_path(guide.data_dir, guide.mode_memory_file)
+        cleared = 0 if path is None else _clear(path)
+        sys.stdout.write(json.dumps({"cleared": cleared}) + "\n")
+        return
+    configure_logging(settings.logging, service=SERVICE_NAME)
+    # httpx logs every request URL at INFO; keep it quiet so record keys stay out of the logs.
+    for name in ("httpx", "httpcore"):
+        logging.getLogger(name).setLevel(logging.WARNING)
+    telemetry = setup_telemetry(settings.telemetry)
+    modes = mode_memory_path(guide.data_dir, guide.mode_memory_file)
+    if modes is not None:
+        use_mode_memory(
+            FileModeMemory(modes),
+            secret=(
+                None
+                if inference.mode_memory_key is None
+                else inference.mode_memory_key.get_secret_value()
+            ),
+            max_age_s=inference.mode_memory_ttl_s,
+        )
+    try:
+        uvicorn.run(
+            create_app(guide, inference),
+            host=guide.host,
+            port=guide.port,
+            log_config=None,
+            access_log=False,
+            proxy_headers=False,
+        )
+    finally:
+        telemetry.shutdown()
+
+
+if __name__ == "__main__":
+    main()

lorebase_guide/agent.py

@@ -0,0 +1,337 @@
+"""The exploration loop: model turn, read-only tool call, repeat; bounded on every axis.
+
+Budgets (from `guide.agent`, lowered per request, never raised):
+
+- `max_steps` tool calls; when they are used up the model gets one last turn to answer;
+- `max_tokens_per_call` completion tokens per model call;
+- `max_total_tokens` input plus output tokens over the whole exploration; before each call the
+  input is estimated from the text length, and the call is only made when that estimate fits in
+  what remains, with a completion cap of at most what is left after it;
+- `timeout_s` wall-clock time from the start, covering model calls and Keeper reads alike.
+
+The loop yields `Event`s that the HTTP layer streams as server-sent events. Logs and metrics
+carry counts, tool names and outcomes only: never the question, data or answers.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import math
+import secrets
+import time
+from collections.abc import AsyncIterator, Awaitable
+from dataclasses import dataclass, field
+from typing import Any
+
+from lorebase_client import AsyncLorebaseClient
+from lorebase_client.errors import LorebaseError
+from lorebase_guide import prompts
+from lorebase_guide.embed import Embedder
+from lorebase_guide.keeper import tool_scope
+from lorebase_guide.settings import AgentSettings
+from lorebase_guide.telemetry import STEPS, TOKENS, tracer
+from lorebase_guide.tools import TOOLS, ToolContext, ToolResult, run_tool
+from lorebase_guide.trace import TraceSession
+from lorebase_inference import (
+    ChatModel,
+    ChatOptions,
+    InferenceError,
+    InvalidOutputError,
+    Message,
+)
+
+logger = logging.getLogger("lorebase.guide.agent")
+
+_CHARS_PER_TOKEN = 4  # input estimate before a call; same ratio as the inference package
+_MIN_CHARS_PER_TOKEN = 1.0  # floor of the ratio measured on a call
+
+
+@dataclass(frozen=True)
+class Event:
+    """One server-sent event: `name` and a JSON object."""
+
+    name: str
+    data: dict[str, Any]
+
+
+@dataclass(frozen=True)
+class Budget:
+    max_steps: int
+    max_total_tokens: int
+    max_tokens_per_call: int
+    timeout_s: float
+
+    @classmethod
+    def of(cls, settings: AgentSettings, *, max_steps: int | None, budget: int | None) -> Budget:
+        return cls(
+            max_steps=min(max_steps or settings.max_steps, settings.max_steps),
+            max_total_tokens=min(budget or settings.max_total_tokens, settings.max_total_tokens),
+            max_tokens_per_call=settings.max_tokens_per_call,
+            timeout_s=settings.timeout_s,
+        )
+
+
+@dataclass
+class _Usage:
+    input_tokens: int = 0
+    output_tokens: int = 0
+
+    @property
+    def total(self) -> int:
+        return self.input_tokens + self.output_tokens
+
+
+@dataclass
+class Exploration:
+    """Outcome counters, read by the HTTP layer for logs and metrics."""
+
+    steps: int = 0
+    outcome: str = "running"
+    usage: _Usage = field(default_factory=_Usage)
+    # Characters per input token: 4 until the endpoint reports its own count of a call, then
+    # what that call measured (JSON tool results ran at about 2.2 on Claude models).
+    chars_per_token: float = _CHARS_PER_TOKEN
+
+
+class _StopError(Exception):
+    def __init__(self, code: str, message: str) -> None:
+        super().__init__(message)
+        self.code = code
+        self.message = message
+
+
+class Explorer:
+    """Runs explorations with one chat model; the Keeper client is per exploration."""
+
+    def __init__(self, chat: ChatModel, settings: AgentSettings, embedder: Embedder | None) -> None:
+        self._chat = chat
+        self._settings = settings
+        self._embedder = embedder
+
+    async def run(
+        self,
+        *,
+        client: AsyncLorebaseClient,
+        trace: TraceSession,
+        domain: str,
+        branch: str | None,
+        question: str,
+        budget: Budget,
+        state: Exploration | None = None,
+    ) -> AsyncIterator[Event]:
+        """Events of one exploration, ending with `answer` (or `error`) and `done`."""
+        state = state or Exploration()
+        started = time.monotonic()
+        deadline = asyncio.get_running_loop().time() + budget.timeout_s
+        marker = secrets.token_hex(12)
+        ctx = ToolContext(
+            client=client,
+            domain=domain,
+            branch=branch,
+            settings=self._settings,
+            embedder=self._embedder,
+        )
+        messages = [
+            Message("system", prompts.system_prompt(marker)),
+            Message("user", prompts.question_prompt(domain, branch, question.replace(marker, ""))),
+        ]
+        yield Event(
+            "session",
+            {
+                "session_id": trace.session_id,
+                "domain": domain,
+                "branch": branch,
+                "model": getattr(self._chat, "model", None),
+                "max_steps": budget.max_steps,
+            },
+        )
+        final: dict[str, Any] | None = None
+        calls = 0
+        try:
+            while final is None:
+                last_turn = state.steps >= budget.max_steps or not self._fits_another_step(
+                    messages, budget, state, calls
+                )
+                if last_turn:
+                    messages.append(Message("user", prompts.FINAL_CALL))
+                turn = await self._turn(messages, budget, state, deadline, final=last_turn)
+                calls += 1
+                thought = str(turn.get("thought") or "").strip()
+                if thought:
+                    yield Event("thought", {"step": state.steps, "text": thought})
+                    await _bounded(trace.note(client, thought), deadline)
+                action = str(turn.get("action") or "")
+                if action == prompts.ANSWER or last_turn:
+                    final = self._final(turn if action == prompts.ANSWER else {}, ctx)
+                    break
+                state.steps += 1
+                raw_args = turn.get("args")
+                args: dict[str, Any] = raw_args if isinstance(raw_args, dict) else {}
+                yield Event("tool_call", {"step": state.steps, "tool": action, "args": args})
+                result = await self._tool(ctx, trace, action, args, deadline)
+                yield Event(
+                    "tool_result",
+                    {
+                        "step": state.steps,
+                        "tool": action,
+                        "ok": result.ok,
+                        "summary": result.summary,
+                        "refs": result.refs[: self._settings.max_items],
+                    },
+                )
+                messages.append(Message("assistant", json.dumps(turn, ensure_ascii=False)))
+                messages.append(
+                    Message(
+                        "user",
+                        prompts.data_block(
+                            marker,
+                            action,
+                            result.ok,
+                            result.data,
+                            self._settings.max_result_chars,
+                        ),
+                    )
+                )
+        except _StopError as stop:
+            state.outcome = stop.code
+            yield Event("error", {"code": stop.code, "message": stop.message})
+            final = {"text": prompts.NOT_FOUND, "citations": [], "found": False}
+        else:
+            state.outcome = "answered" if final["found"] else "not_found"
+        yield Event("answer", final)
+        yield Event(
+            "done",
+            {
+                "session_id": trace.session_id,
+                "steps": state.steps,
+                "outcome": state.outcome,
+                "tokens": {
+                    "input": state.usage.input_tokens,
+                    "output": state.usage.output_tokens,
+                },
+                "elapsed_ms": round((time.monotonic() - started) * 1000),
+            },
+        )
+
+    def _fits_another_step(
+        self, messages: list[Message], budget: Budget, state: Exploration, calls: int
+    ) -> bool:
+        """Whether the token budget still covers a tool step and the answer after it.
+
+        A tool step costs this call (estimated input, plus the output of an average call so
+        far) and leaves a larger context for the next one: this input plus at most
+        `max_result_chars` of tool result. The answer after it may use up to
+        `max_tokens_per_call`. When that does not fit, this turn is the last one, so the
+        exploration answers from what it read instead of stopping on the budget with nothing.
+        """
+        estimate = math.ceil(_chars(messages) / state.chars_per_token)
+        output = (
+            math.ceil(state.usage.output_tokens / calls)
+            if calls
+            else budget.max_tokens_per_call // 4
+        )
+        result = math.ceil(self._settings.max_result_chars / state.chars_per_token)
+        needed = (estimate + output) + (estimate + output + result) + budget.max_tokens_per_call
+        return state.usage.total + needed <= budget.max_total_tokens
+
+    async def _turn(
+        self,
+        messages: list[Message],
+        budget: Budget,
+        state: Exploration,
+        deadline: float,
+        *,
+        final: bool,
+    ) -> dict[str, Any]:
+        schema = prompts.turn_schema()
+        chars = _chars(messages)
+        remaining = (
+            budget.max_total_tokens - state.usage.total - math.ceil(chars / state.chars_per_token)
+        )
+        if remaining <= 0:
+            raise _StopError("budget", "the token budget of this exploration is used up")
+        options = ChatOptions(
+            task=self._settings.task,
+            temperature=self._settings.temperature,
+            max_tokens=min(budget.max_tokens_per_call, remaining),
+        )
+        try:
+            with tracer.start_as_current_span("guide.model"):
+                result = await _bounded(
+                    self._chat.complete_json(messages, schema, options), deadline
+                )
+        except InferenceError as exc:
+            # The calls of an unusable answer were made and billed: count them. The problems
+            # are schema paths and keywords (for example `/: json`), never answer text.
+            detail: dict[str, Any] = {"error": type(exc).__name__}
+            if isinstance(exc, InvalidOutputError):
+                state.usage.input_tokens += exc.usage.input_tokens
+                state.usage.output_tokens += exc.usage.output_tokens
+                detail.update(attempts=exc.attempts, problems=list(exc.problems[:5]))
+            logger.warning("model call failed", extra=detail)
+            raise _StopError("model", "the model could not produce a valid step") from None
+        state.usage.input_tokens += result.usage.input_tokens
+        state.usage.output_tokens += result.usage.output_tokens
+        if result.usage.input_tokens > 0:
+            state.chars_per_token = min(
+                max(chars / result.usage.input_tokens, _MIN_CHARS_PER_TOKEN), _CHARS_PER_TOKEN
+            )
+        TOKENS.add(result.usage.input_tokens, {"direction": "input"})
+        TOKENS.add(result.usage.output_tokens, {"direction": "output"})
+        # The answer of the final turn is kept even when the endpoint counted more input than
+        # estimated: its tokens are spent either way, and no further call follows.
+        if state.usage.total > budget.max_total_tokens and not final:
+            raise _StopError("budget", "the token budget of this exploration is used up")
+        data = result.data
+        return data if isinstance(data, dict) else {}
+
+    async def _tool(
+        self,
+        ctx: ToolContext,
+        trace: TraceSession,
+        action: str,
+        args: dict[str, Any],
+        deadline: float,
+    ) -> ToolResult:
+        with (
+            tracer.start_as_current_span("guide.tool") as span,
+            tool_scope(trace.tool_headers(action)),
+        ):
+            span.set_attribute("guide.tool", action)
+            try:
+                result = await _bounded(run_tool(ctx, action, args), deadline)
+            except LorebaseError as exc:
+                result = ToolResult(False, f"Keeper error {exc.status}", {"error": "keeper"})
+        STEPS.add(1, {"tool": action if action in TOOLS else "unknown", "ok": result.ok})
+        return result
+
+    def _final(self, turn: dict[str, Any], ctx: ToolContext) -> dict[str, Any]:
+        """The answer, citing only Records a tool returned; without such a citation, or when
+        the model says so, the answer is that the graph does not contain it."""
+        raw = turn.get("citations")
+        cited: list[Any] = raw if isinstance(raw, list) else []
+        citations: list[dict[str, str]] = []
+        for ref in dict.fromkeys(str(c) for c in cited):
+            if ref in ctx.seen:
+                kind, _, key = ref.partition("/")
+                citations.append({"ref": ref, "kind": kind, "key": key})
+        text = str(turn.get("answer") or "").strip()[: self._settings.max_answer_chars]
+        found = bool(turn.get("found", True)) and bool(citations) and bool(text)
+        if not found:
+            return {"text": prompts.NOT_FOUND, "citations": [], "found": False}
+        return {"text": text, "citations": citations, "found": True}
+
+
+def _chars(messages: list[Message]) -> int:
+    """Characters of a model call with these messages and the turn schema."""
+    return len(json.dumps(prompts.turn_schema())) + sum(len(m.content) for m in messages)
+
+
+async def _bounded[T](awaitable: Awaitable[T], deadline: float) -> T:
+    try:
+        async with asyncio.timeout_at(deadline):
+            return await awaitable
+    except TimeoutError:
+        raise _StopError("timeout", "the exploration took longer than allowed") from None