haeaeae/Stirling-PDF
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge origin/main into AI-Form-Fill

Browse Source 
Resolution favours main for all conflicts (new orchestrator shape with
artifacts/file_names/conversation_history/resume_with, WorkflowOutcome
enum, RAG, ledger agent, deleted engine/config/.env.example).

Re-wires our three form-fill agents (FormAnalyserAgent, FormFillerAgent,
DocumentExtractorAgent) into the new app layout:
- contracts/__init__.py: adds form-fill exports to __all__
- api/dependencies.py: adds get_form_analyser_agent / get_form_filler_agent / get_document_extractor_agent
- api/app.py: instantiates the three agents in lifespan and registers form_fill_router
- api/routes/__init__.py: exports form_fill_router
- tests/test_stirling_api.py: imports and registers form-fill stubs
- tests/test_stirling_contracts.py: adds back KnowledgeUpdateResponse discriminator test

Form fill remains accessible via its own endpoints (/api/v1/form/ai/analyse,
/fill-batch, /extract). Not wired as an orchestrator delegate — following
main's pattern where orchestrator only routes pdf_edit/pdf_question/user_spec
/math_auditor.

Follow-ups still needed:
- Thread conversation_history into form-fill agent prompts
- Align form-fill response outcomes with WorkflowOutcome enum
- Decide whether engine should return ToolOperationStep plans (main's
  new pattern per #6116) or keep returning fill values directly

All 127 engine tests pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Reece
2026-04-22 13:13:12 +01:00
co-authored by Claude Opus 4.7
parent 80cf2f9a73 fdd8bc23ca
commit b32a3cf271
1770 changed files with 145755 additions and 64948 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/.env
+44
View File
@@ -0,0 +1,44 @@
###############################################################################
# Environment variables used within the AI Engine.
# Values can be overridden in the uncommitted sibling `.env.local` file.
# Note: This file is committed to Git, so should not contain any private keys.
###############################################################################
# Configure the model strings passed to pydantic-ai. Provider credentials are handled by
# pydantic-ai and should be set using the provider's native environment variables, for example
# ANTHROPIC_API_KEY or OPENAI_API_KEY.
STIRLING_SMART_MODEL=anthropic:claude-haiku-4-5
STIRLING_FAST_MODEL=anthropic:claude-haiku-4-5
# Default output token limits applied by the engine for each model tier.
STIRLING_SMART_MODEL_MAX_TOKENS=8192
STIRLING_FAST_MODEL_MAX_TOKENS=2048
# RAG Configuration — retrieval-augmented generation is always on.
# Embedding provider credentials are handled natively (e.g. VOYAGE_API_KEY for VoyageAI).
STIRLING_RAG_EMBEDDING_MODEL=voyageai:voyage-4
# Vector store backend: "sqlite" (embedded) or "pgvector" (external Postgres).
STIRLING_RAG_BACKEND=sqlite
# Path to the sqlite-vec database file (used when backend=sqlite).
STIRLING_RAG_STORE_PATH=data/rag.db
# Postgres DSN for pgvector (used when backend=pgvector). Leave empty when backend=sqlite.
# Example: postgresql://user:password@host:5432/dbname
STIRLING_RAG_PGVECTOR_DSN=
STIRLING_RAG_CHUNK_SIZE=512
STIRLING_RAG_CHUNK_OVERLAP=64
STIRLING_RAG_TOP_K=5
# PostHog analytics. Set STIRLING_POSTHOG_ENABLED=true and provide an API key to enable.
STIRLING_POSTHOG_ENABLED=false
STIRLING_POSTHOG_API_KEY=phc_VOdeYnlevc2T63m3myFGjeBlRcIusRgmhfx6XL5a1iz
STIRLING_POSTHOG_HOST=https://eu.i.posthog.com
# Log level for the stirling logger hierarchy (DEBUG, INFO, WARNING, ERROR)
STIRLING_LOG_LEVEL=INFO
# Path to log file. Rolls daily, keeps 1 backup. Leave empty for console only.
STIRLING_LOG_FILE=
engine/.gitignore
-1
View File
@@ -19,7 +19,6 @@ yarn-error.log*
.vite/
# Environment
.env
.env.local
# LaTeX outputs
engine/AGENTS.md
+15 -1
View File
@@ -4,9 +4,23 @@ This file is for AI agents working in `engine/`.
The engine is a Python reasoning service for Stirling. It plans and interprets work, but it does not own durable state, and it does not execute Stirling PDF operations directly. Keep the service narrow: typed contracts in, typed contracts out, with AI only where it adds reasoning value.
## Commands
All engine commands can be run from the repository root using Task:
- `task engine:check` — run all checks (typecheck + lint + format-check + test)
- `task engine:fix` — auto-fix lint + formatting
- `task engine:install` — install Python dependencies via uv
- `task engine:dev` — start FastAPI with hot reload (localhost:5001)
- `task engine:test` — run pytest
- `task engine:lint` — run ruff linting
- `task engine:typecheck` — run pyright
- `task engine:format` — format code with ruff
- `task engine:tool-models` — generate tool_models.py from Java OpenAPI spec
## Code Style
- Keep `make check` passing.
- Keep `task engine:check` passing.
- Use modern Python when it improves clarity.
- Prefer explicit names to cleverness.
- Avoid nested functions and nested classes unless the language construct requires them.
engine/Dockerfile
+13 -2
View File
@@ -1,9 +1,20 @@
# syntax=docker/dockerfile:1.5
FROM ghcr.io/astral-sh/uv:python3.13-bookworm-slim
ARG TASK_VERSION=3.49.1
RUN apt-get update \
&& apt-get install -y --no-install-recommends curl ca-certificates \
&& ARCH=$(dpkg --print-architecture) \
&& curl -fsSL "https://github.com/go-task/task/releases/download/v${TASK_VERSION}/task_${TASK_VERSION}_linux_${ARCH}.deb" -o /tmp/task.deb \
&& dpkg -i /tmp/task.deb \
&& rm /tmp/task.deb \
&& rm -rf /var/lib/apt/lists/*
WORKDIR /app
COPY pyproject.toml uv.lock ./
COPY pyproject.toml uv.lock Taskfile.yml .env ./
COPY .taskfiles/ ./.taskfiles/
COPY scripts/ ./scripts/
RUN --mount=type=cache,target=/root/.cache/uv \
uv sync --frozen --no-dev
@@ -14,4 +25,4 @@ ENV PYTHONUNBUFFERED=1
EXPOSE 5001
CMD ["uv", "run", "uvicorn", "stirling.api.app:app", "--host", "0.0.0.0", "--port", "5001"]
CMD ["task", "engine:run"]
engine/Makefile
-82
View File
@@ -1,82 +0,0 @@
.PHONY: help install prep check fix lint lint-fix lint-fix-unsafe format format-check typecheck test run run-dev clean tool-models docker-build docker-run
REPO_ROOT_DIR = ..
ROOT_DIR = .
VENV_DIR = $(ROOT_DIR)/.venv
DEPS_STAMP := $(VENV_DIR)/.deps-installed
TOOL_MODELS := $(CURDIR)/src/stirling/models/tool_models.py
FRONTEND_DIR := $(REPO_ROOT_DIR)/frontend
FRONTEND_TSX := $(FRONTEND_DIR)/node_modules/.bin/tsx
$(DEPS_STAMP): $(ROOT_DIR)/uv.lock $(ROOT_DIR)/pyproject.toml
uv python install 3.13.8
uv sync
touch $(DEPS_STAMP)
help:
@echo "Engine commands:"
@echo " make install - Install production dependencies"
@echo " make prep - Set up .env file from .env.example"
@echo " make lint - Run linting checks"
@echo " make format - Format code"
@echo " make format-check - Show whether code is correctly formatted"
@echo " make typecheck - Run type checking"
@echo " make test - Run tests"
@echo " make run - Run the FastAPI backend with uvicorn"
@echo " make run-dev - Run the FastAPI backend with reload"
@echo " make tool-models - Generate src/stirling/models/tool_models.py from frontend TypeScript tool defs"
@echo " make clean - Clean up generated files"
@echo " make docker-build - Build Docker image"
@echo " make docker-run - Run Docker container"
install: $(DEPS_STAMP)
prep: install
uv run scripts/setup_env.py
check: typecheck lint format-check test
fix: lint-fix format
lint: install
uv run ruff check .
lint-fix: install
uv run ruff check . --fix
lint-fix-unsafe: install
uv run ruff check . --fix --unsafe-fixes
format: install
uv run ruff format .
format-check: install
uv run ruff format . --diff
typecheck: install
uv run pyright . --warnings
test: prep
uv run pytest tests
run: prep
cd src && PYTHONUNBUFFERED=1 uv run uvicorn stirling.api.app:app --host 0.0.0.0 --port 5001
run-dev: prep
cd src && PYTHONUNBUFFERED=1 uv run uvicorn stirling.api.app:app --host 0.0.0.0 --port 5001 --reload
frontend-deps:
if [ ! -x "$(FRONTEND_TSX)" ]; then cd $(FRONTEND_DIR) && npm install; fi
tool-models: install frontend-deps
uv run python scripts/generate_tool_models.py --output $(TOOL_MODELS)
$(MAKE) fix
clean:
rm -rf $(VENV_DIR) data logs output
docker-build:
docker build -t stirling-pdf-engine .
docker-run:
docker run -p 5001:5001 stirling-pdf-engine
engine/pyproject.toml
+13 -3
View File
@@ -5,17 +5,26 @@ description = "AI Document Engine"
requires-python = ">=3.13"
dependencies = [
"fastapi>=0.116.0",
"pgvector>=0.3.6",
"psycopg[binary]>=3.2",
"pydantic>=2.0.0",
"pydantic-ai>=1.67.0",
"pydantic-ai-slim[voyageai]>=1.67.0",
"pydantic-settings>=2.0.0",
"python-dotenv>=1.2.1",
"sqlite-vec>=0.1.6",
"uvicorn>=0.35.0",
"opentelemetry-sdk>=1.39.0",
"posthog>=3.0.0",
]
[dependency-groups]
dev = [
"anyio>=4.0.0",
"datamodel-code-generator[ruff]>=0.26.0",
"pytest>=8.0.0",
"pyright>=1.1.408",
"referencing>=0.35.0",
"ruff>=0.14.10",
]
@@ -42,9 +51,9 @@ select = [
"W",
"RUF100",
"UP",
]
ignore = [
"E501", # Temporarily disable line length limit until codebase conformat
"PYI", # flake8-pyi: flags deprecated typing constructs
"FA", # flake8-future-annotations: flags missing future annotations imports
"BLE", # flake8-blind-except: flags bare `except Exception`
]
[tool.pyright]
@@ -55,6 +64,7 @@ reportUnnecessaryCast = "warning"
reportUnnecessaryTypeIgnoreComment = "warning"
reportUnusedImport = "warning"
reportUnknownParameterType = "warning"
reportDeprecated = "warning"
[tool.pytest.ini_options]
testpaths = ["tests"]
engine/scripts/generate_tool_models.py
+184 -471
View File
@@ -1,509 +1,222 @@
#!/usr/bin/env python3
"""Generate Python tool models from the Java backend's OpenAPI spec (SwaggerDoc.json).
Uses datamodel-code-generator to convert OpenAPI request schemas to Pydantic models.
Run via:
task engine:tool-models
"""
from __future__ import annotations
import argparse
import json
import keyword
import re
import subprocess
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import Any
TOOL_MODELS_HEADER = """# AUTO-GENERATED FILE. DO NOT EDIT.
# Generated by scripts/generate_tool_models.py from frontend TypeScript sources.
# ruff: noqa: N815
"""
from datamodel_code_generator import InputFileType, PythonVersion, generate
from datamodel_code_generator.enums import DataModelType
from datamodel_code_generator.format import Formatter
from referencing import Registry, Resource
from referencing.jsonschema import DRAFT202012
# Fields inherited from PDFFile base class — not tool parameters.
BASE_CLASS_FIELDS = frozenset({"fileInput", "fileId"})
OPERATION_TYPE_RE = re.compile(r"operationType\s*:\s*['\"]([A-Za-z0-9_]+)['\"]")
DEFAULT_REF_RE = re.compile(r"defaultParameters\s*:\s*([A-Za-z0-9_]+)")
DEFAULT_SHORTHAND_RE = re.compile(r"\bdefaultParameters\b")
IMPORT_RE = re.compile(r"import\s*\{([^}]+)\}\s*from\s*['\"]([^'\"]+)['\"]")
VAR_OBJ_RE_TEMPLATE = r"(?:export\s+)?const\s+{name}\b[^=]*=\s*\{{"
_ENGINE_ROOT = Path(__file__).resolve().parents[1]
_FILE_HEADER = (
"# AUTO-GENERATED FILE. DO NOT EDIT.\n"
"# Generated by scripts/generate_tool_models.py from Java OpenAPI spec (SwaggerDoc.json).\n"
"# ruff: noqa: E501"
)
@dataclass
class ToolModelSpec:
tool_id: str
params: dict[str, Any]
param_types: dict[str, Any]
class ToolSpec:
path: str
enum_name: str
class_name: str
class ParseError(Exception):
pass
@dataclass
class DiscoveryResult:
tools: list[ToolSpec]
combined_schema: dict[str, Any]
def _find_matching(text: str, start: int, open_char: str, close_char: str) -> int:
depth = 0
i = start
in_str: str | None = None
while i < len(text):
ch = text[i]
if in_str:
if ch == "\\":
i += 2
continue
if ch == in_str:
in_str = None
i += 1
continue
if ch in {"'", '"'}:
in_str = ch
elif ch == open_char:
depth += 1
elif ch == close_char:
depth -= 1
if depth == 0:
return i
i += 1
raise ParseError(f"Unmatched {open_char}{close_char} block")
class ToolDiscovery:
"""Discovers tool endpoints from an OpenAPI spec and builds a combined JSON Schema."""
def _extract_block(text: str, pattern: str) -> str | None:
match = re.search(pattern, text)
if not match:
return None
brace_start = text.find("{", match.end() - 1)
if brace_start == -1:
return None
brace_end = _find_matching(text, brace_start, "{", "}")
return text[brace_start : brace_end + 1]
def _split_top_level_items(obj_body: str) -> list[str]:
items: list[str] = []
depth_obj = depth_arr = 0
in_str: str | None = None
token_start = 0
i = 0
while i < len(obj_body):
ch = obj_body[i]
if in_str:
if ch == "\\":
i += 2
continue
if ch == in_str:
in_str = None
i += 1
continue
if ch in {"'", '"'}:
in_str = ch
elif ch == "{":
depth_obj += 1
elif ch == "}":
depth_obj -= 1
elif ch == "[":
depth_arr += 1
elif ch == "]":
depth_arr -= 1
elif ch == "," and depth_obj == 0 and depth_arr == 0:
piece = obj_body[token_start:i].strip()
if piece:
items.append(piece)
token_start = i + 1
i += 1
tail = obj_body[token_start:].strip()
if tail:
items.append(tail)
return items
def _resolve_import_path(repo_root: Path, current_file: Path, module_path: str) -> Path | None:
candidates: list[Path] = []
if module_path.startswith("@app/"):
rel = module_path[len("@app/") :]
candidates.extend(
[
repo_root / "frontend/src/core" / f"{rel}.ts",
repo_root / "frontend/src/core" / f"{rel}.tsx",
repo_root / "frontend/src/saas" / f"{rel}.ts",
repo_root / "frontend/src/saas" / f"{rel}.tsx",
repo_root / "frontend/src" / f"{rel}.ts",
repo_root / "frontend/src" / f"{rel}.tsx",
]
)
elif module_path.startswith("."):
base = (current_file.parent / module_path).resolve()
candidates.extend([Path(f"{base}.ts"), Path(f"{base}.tsx")])
for candidate in candidates:
if candidate.exists():
return candidate
return None
def _parse_literal_value(value: str, resolver: Callable[[str], dict[str, Any] | None]) -> Any:
value = value.strip()
if not value:
return None
if value.startswith("{") and value.endswith("}"):
return _parse_object_literal(value, resolver)
if value.startswith("[") and value.endswith("]"):
inner = value[1:-1].strip()
if not inner:
return []
return [_parse_literal_value(item, resolver) for item in _split_top_level_items(inner)]
if value.startswith(("'", '"')) and value.endswith(("'", '"')):
return value[1:-1]
if value in {"true", "false"}:
return value == "true"
if value == "null":
return None
if re.fullmatch(r"-?\d+", value):
return int(value)
if re.fullmatch(r"-?\d+\.\d+", value):
return float(value)
resolved = resolver(value)
if resolved is not None:
return resolved
return None
def _parse_object_literal(obj_text: str, resolver: Callable[[str], dict[str, Any] | None]) -> dict[str, Any]:
body = obj_text.strip()[1:-1]
result: dict[str, Any] = {}
for item in _split_top_level_items(body):
if item.startswith("..."):
spread_name = item[3:].strip()
spread = resolver(spread_name)
if isinstance(spread, dict):
result.update(spread)
continue
if ":" not in item:
continue
key, raw_value = item.split(":", 1)
key = key.strip().strip("'\"")
result[key] = _parse_literal_value(raw_value.strip(), resolver)
return result
def _extract_imports(source: str) -> dict[str, str]:
imports: dict[str, str] = {}
for names, module_path in IMPORT_RE.findall(source):
for part in names.split(","):
segment = part.strip()
if not segment:
continue
if " as " in segment:
original, alias = [x.strip() for x in segment.split(" as ", 1)]
imports[alias] = module_path
imports[original] = module_path
else:
imports[segment] = module_path
return imports
def _resolve_object_identifier(repo_root: Path, file_path: Path, source: str, identifier: str) -> dict[str, Any] | None:
var_pattern = VAR_OBJ_RE_TEMPLATE.format(name=re.escape(identifier))
block = _extract_block(source, var_pattern)
imports = _extract_imports(source)
def resolver(name: str) -> dict[str, Any] | None:
local_block = _extract_block(source, VAR_OBJ_RE_TEMPLATE.format(name=re.escape(name)))
if local_block:
return _parse_object_literal(local_block, resolver)
import_path = imports.get(name)
if not import_path:
return None
resolved_file = _resolve_import_path(repo_root, file_path, import_path)
if not resolved_file:
return None
imported_source = resolved_file.read_text(encoding="utf-8")
return _resolve_object_identifier(repo_root, resolved_file, imported_source, name)
if block:
return _parse_object_literal(block, resolver)
import_path = imports.get(identifier)
if not import_path:
return None
resolved_file = _resolve_import_path(repo_root, file_path, import_path)
if not resolved_file:
return None
imported_source = resolved_file.read_text(encoding="utf-8")
return _resolve_object_identifier(repo_root, resolved_file, imported_source, identifier)
def _infer_py_type(value: Any) -> str:
if isinstance(value, bool):
return "bool"
if isinstance(value, int):
return "int"
if isinstance(value, float):
return "float"
if isinstance(value, str):
return "str"
if isinstance(value, list):
return "list[Any]"
if isinstance(value, dict):
return "dict[str, Any]"
return "Any"
def _spec_is_none(spec: dict[str, Any]) -> bool:
return spec.get("kind") == "null"
def _py_type_from_spec(spec: dict[str, Any]) -> str:
kind = spec.get("kind")
if kind == "string":
return "str"
if kind == "number":
return "float"
if kind == "boolean":
return "bool"
if kind == "date":
return "str"
if kind == "enum":
values = spec.get("values")
if isinstance(values, list) and values:
literal_values = ", ".join(_py_repr(v) for v in values)
return f"Literal[{literal_values}]"
if kind == "ref":
ref_name = spec.get("name")
if isinstance(ref_name, str) and ref_name.endswith("Parameters"):
return f"{ref_name[:-10]}Params"
if kind == "array":
element = spec.get("element")
inner = _py_type_from_spec(element) if isinstance(element, dict) else "Any"
return f"list[{inner}]"
if kind == "object":
dict_value = spec.get("dictValue")
if isinstance(dict_value, dict):
inner = _py_type_from_spec(dict_value)
return f"dict[str, {inner}]"
properties = spec.get("properties")
if isinstance(properties, dict) and properties:
property_types = {_py_type_from_spec(p) for p in properties.values() if isinstance(p, dict)}
if len(property_types) == 1:
inner = next(iter(property_types))
return f"dict[str, {inner}]"
return "dict[str, Any]"
if kind in {"null"}:
return "Any"
return "Any"
def _to_class_name(tool_id: str) -> str:
cleaned = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", tool_id)
cleaned = re.sub(r"[^A-Za-z0-9]+", " ", cleaned)
parts = [part.capitalize() for part in cleaned.split() if part]
return "".join(parts) + "Params"
def _to_snake_case(name: str) -> str:
snake = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", name)
snake = re.sub(r"[^A-Za-z0-9]+", "_", snake).strip("_").lower()
if not snake:
snake = "param"
if snake[0].isdigit():
snake = f"param_{snake}"
if keyword.iskeyword(snake):
snake = f"{snake}_"
return snake
def _build_field_name_map(params: dict[str, Any]) -> dict[str, str]:
field_map: dict[str, str] = {}
used: set[str] = set()
for original_key in sorted(params):
base_name = _to_snake_case(original_key)
candidate = base_name
suffix = 2
while candidate in used:
candidate = f"{base_name}_{suffix}"
suffix += 1
used.add(candidate)
field_map[original_key] = candidate
return field_map
def _to_enum_member_name(tool_id: str) -> str:
return _to_snake_case(tool_id).upper()
def _build_enum_member_map(specs: list[ToolModelSpec]) -> dict[str, str]:
member_map: dict[str, str] = {}
used: set[str] = set()
for spec in specs:
base_name = _to_enum_member_name(spec.tool_id)
candidate = base_name
suffix = 2
while candidate in used:
candidate = f"{base_name}_{suffix}"
suffix += 1
used.add(candidate)
member_map[spec.tool_id] = candidate
return member_map
def _py_repr(value: Any) -> str:
return (
json.dumps(value, ensure_ascii=True).replace("true", "True").replace("false", "False").replace("null", "None")
# Namespaces exposed to the LLM as callable tools. Largely matches ``InternalApiClient.java``.
# Note: ``/api/v1/filter/`` is intentionally excluded because those APIs are for pipeline processing,
# not tool execution.
ALLOWED_PATH_PREFIXES = (
"/api/v1/general/",
"/api/v1/misc/",
"/api/v1/security/",
"/api/v1/convert/",
)
def __init__(self, spec: dict[str, Any]):
resource = Resource.from_contents(spec, default_specification=DRAFT202012)
self.resolver = Registry().with_resource("", resource).resolver()
self.spec = spec
def discover_tool_specs(repo_root: Path) -> list[ToolModelSpec]:
frontend_dir = repo_root / "frontend"
extractor = frontend_dir / "scripts/export-tool-specs.ts"
command = ["node", "--import", "tsx", str(extractor)]
result = subprocess.run(
command,
check=True,
capture_output=True,
text=True,
cwd=str(frontend_dir),
def discover(self) -> DiscoveryResult:
tools: list[ToolSpec] = []
defs: dict[str, Any] = {}
used_enum: set[str] = set()
used_class: set[str] = set()
for path, path_item in sorted(self.spec.get("paths", {}).items()):
if "{" in path or not any(path.startswith(p) for p in self.ALLOWED_PATH_PREFIXES):
continue
properties = self._get_request_properties(path_item)
if not properties:
continue
clean_props = self._filter_properties(properties)
if not clean_props:
continue
enum_name = _deduplicate(_path_to_enum_name(path), used_enum)
class_name = _deduplicate(_path_to_class_name(path), used_class)
defs[class_name] = {"type": "object", "properties": clean_props}
tools.append(ToolSpec(path, enum_name, class_name))
combined_schema: dict[str, Any] = {
"$defs": defs,
"anyOf": [{"$ref": f"#/$defs/{t.class_name}"} for t in tools],
}
return DiscoveryResult(tools=tools, combined_schema=combined_schema)
def _resolve_ref(self, schema: dict[str, Any]) -> dict[str, Any]:
if "$ref" in schema:
return self.resolver.lookup(schema["$ref"]).contents
return schema
def _get_request_properties(self, path_item: dict[str, Any]) -> dict[str, Any] | None:
post = path_item.get("post")
if not post:
return None
content = post.get("requestBody", {}).get("content", {})
for media_type in ("multipart/form-data", "application/json"):
if media_type in content:
schema = content[media_type].get("schema")
if schema:
return self._resolve_ref(schema).get("properties")
return None
def _filter_properties(self, properties: dict[str, Any]) -> dict[str, Any]:
"""Remove base-class fields and binary upload fields, resolving any $refs."""
clean: dict[str, Any] = {}
for name, prop in properties.items():
if name in BASE_CLASS_FIELDS:
continue
prop = self._resolve_ref(prop)
if prop.get("type") == "string" and prop.get("format") == "binary":
continue
clean[name] = prop
return clean
def _tool_name_segments(path: str) -> str:
"""Extract a descriptive name from the endpoint path.
Converters use two segments (e.g. /api/v1/convert/cbr/pdf → cbr-to-pdf).
Other tools use the last segment (e.g. /api/v1/misc/compress-pdf → compress-pdf).
"""
parts = path.rstrip("/").split("/")
if "/api/v1/convert/" in path and len(parts) >= 6:
return f"{parts[-2]}-to-{parts[-1]}"
return parts[-1]
def _path_to_enum_name(path: str) -> str:
return _tool_name_segments(path).replace("-", "_").upper()
def _path_to_class_name(path: str) -> str:
return "".join(p.capitalize() for p in _tool_name_segments(path).split("-")) + "Params"
def _deduplicate(name: str, used: set[str]) -> str:
"""Return name, appending 2, 3, ... if already in used. Adds result to used."""
candidate = name
n = 2
while candidate in used:
candidate = f"{name}{n}"
n += 1
used.add(candidate)
return candidate
def generate_models_code(combined_schema: dict[str, Any]) -> str:
"""Run datamodel-code-generator once on the combined schema."""
code = generate(
input_=json.dumps(combined_schema, sort_keys=True),
input_file_type=InputFileType.JsonSchema,
output_model_type=DataModelType.PydanticV2BaseModel,
target_python_version=PythonVersion.PY_313,
snake_case_field=True,
base_class="stirling.models.base.ApiModel",
field_constraints=True,
no_alias=True,
set_default_enum_member=True,
additional_imports=["enum.StrEnum"],
enable_version_header=False,
custom_file_header=_FILE_HEADER,
formatters=[Formatter.RUFF_FORMAT, Formatter.RUFF_CHECK],
settings_path=_ENGINE_ROOT / "pyproject.toml",
)
raw = json.loads(result.stdout)
specs: list[ToolModelSpec] = []
for item in raw:
tool_id = item.get("tool_id")
if not isinstance(tool_id, str) or not tool_id:
continue
params = item.get("params")
param_types = item.get("param_types")
specs.append(
ToolModelSpec(
tool_id=tool_id,
params=params if isinstance(params, dict) else {},
param_types=param_types if isinstance(param_types, dict) else {},
)
)
return sorted(specs, key=lambda spec: spec.tool_id)
return str(code or "")
def write_models_module(out_path: Path, specs: list[ToolModelSpec]) -> None:
lines: list[str] = [
TOOL_MODELS_HEADER,
"from __future__ import annotations\n\n",
"from enum import StrEnum\n",
"from typing import Any, Literal\n\n",
"from stirling.models.base import ApiModel\n",
def write_output(out_path: Path, tools: list[ToolSpec], models_code: str) -> None:
union_lines = ["type ParamToolModel = ("]
for i, tool in enumerate(tools):
prefix = " | " if i > 0 else " "
union_lines.append(f"{prefix}{tool.class_name}")
union_lines.append(")")
union_lines.append("type ParamToolModelType = type[ParamToolModel]")
enum_lines = [
"class ToolEndpoint(StrEnum):",
*(f' {t.enum_name} = "{t.path}"' for t in tools),
]
class_names: dict[str, str] = {spec.tool_id: _to_class_name(spec.tool_id) for spec in specs}
class_name_to_tool_id = {name: tool_id for tool_id, name in class_names.items()}
ops_lines = [
"OPERATIONS: dict[ToolEndpoint, ParamToolModelType] = {",
*(f" ToolEndpoint.{t.enum_name}: {t.class_name}," for t in tools),
"}",
]
def extract_class_dependencies(spec: ToolModelSpec) -> set[str]:
deps: set[str] = set()
if not isinstance(spec.param_types, dict):
return deps
for entry in spec.param_types.values():
if not isinstance(entry, dict):
continue
type_spec = entry
if "type" in entry and isinstance(entry.get("type"), dict):
type_spec = entry["type"]
if not isinstance(type_spec, dict):
continue
if type_spec.get("kind") != "ref":
continue
ref_name = type_spec.get("name")
if isinstance(ref_name, str) and ref_name.endswith("Parameters"):
ref_class = f"{ref_name[:-10]}Params"
if ref_class in class_name_to_tool_id:
deps.add(ref_class)
return deps
dependencies_by_class: dict[str, set[str]] = {}
for spec in specs:
class_name = class_names[spec.tool_id]
dependencies_by_class[class_name] = extract_class_dependencies(spec)
remaining = set(class_names.values())
ordered_class_names: list[str] = []
while remaining:
progress = False
for class_name in sorted(remaining):
deps = dependencies_by_class.get(class_name, set())
if deps.issubset(set(ordered_class_names)):
ordered_class_names.append(class_name)
remaining.remove(class_name)
progress = True
break
if not progress:
ordered_class_names.extend(sorted(remaining))
break
ordered_specs = [next(spec for spec in specs if class_names[spec.tool_id] == name) for name in ordered_class_names]
for spec in ordered_specs:
class_name = class_names[spec.tool_id]
lines.append(f"class {class_name}(ApiModel):\n")
all_param_keys = set(spec.params)
if isinstance(spec.param_types, dict):
all_param_keys.update(spec.param_types.keys())
if not all_param_keys:
lines.append(" pass\n\n\n")
continue
field_name_map = _build_field_name_map({key: True for key in all_param_keys})
for key in sorted(all_param_keys):
field_name = field_name_map[key]
value = spec.params.get(key)
type_spec = spec.param_types.get(key) if isinstance(spec.param_types, dict) else None
if isinstance(type_spec, dict):
py_type = _py_type_from_spec(type_spec)
else:
py_type = _infer_py_type(value)
if value is None and (isinstance(type_spec, dict) and _spec_is_none(type_spec)):
if py_type != "Any" and "| None" not in py_type:
py_type = f"{py_type} | None"
lines.append(f" {field_name}: {py_type} = None\n")
elif value is None:
lines.append(f" {field_name}: {py_type} | None = None\n")
else:
if isinstance(type_spec, dict) and type_spec.get("kind") == "ref" and isinstance(value, dict):
lines.append(f" {field_name}: {py_type} = {py_type}.model_validate({_py_repr(value)})\n")
continue
lines.append(f" {field_name}: {py_type} = {_py_repr(value)}\n")
lines.append("\n\n")
if class_names:
union_members = " | ".join(class_names[tool_id] for tool_id in sorted(class_names))
lines.append(f"type ParamToolModel = {union_members}\n")
lines.append("type ParamToolModelType = type[ParamToolModel]\n\n")
else:
lines.append("type ParamToolModel = ApiModel\n")
lines.append("type ParamToolModelType = type[ParamToolModel]\n\n")
enum_member_map = _build_enum_member_map(specs)
lines.append("class OperationId(StrEnum):\n")
for spec in specs:
lines.append(f" {enum_member_map[spec.tool_id]} = {spec.tool_id!r}\n")
lines.extend(
[
"\n\n",
"OPERATIONS: dict[OperationId, ParamToolModelType] = {\n",
]
)
for spec in specs:
model_name = _to_class_name(spec.tool_id)
lines.append(f" OperationId.{enum_member_map[spec.tool_id]}: {model_name},\n")
lines.append("}\n")
out_path.write_text("".join(lines), encoding="utf-8")
parts = [models_code, "\n", *union_lines, "\n", *enum_lines, "\n", *ops_lines, ""]
out_path.write_text("\n".join(parts), encoding="utf-8")
def main() -> None:
parser = argparse.ArgumentParser(description="Generate tool models from frontend TypeScript tool definitions")
parser.add_argument("--spec", help="Deprecated (ignored)", default="")
parser.add_argument("--output", default="", help="Path to tool_models.py")
parser.add_argument("--ai-output", default="", help="Deprecated (ignored)")
parser = argparse.ArgumentParser(description="Generate Python tool models from Java OpenAPI spec")
parser.add_argument("--spec", required=True, help="Path to SwaggerDoc.json")
parser.add_argument("--output", required=True, help="Path to output tool_models.py")
args = parser.parse_args()
repo_root = Path(__file__).resolve().parents[3]
specs = discover_tool_specs(repo_root)
spec_path = Path(args.spec)
if not spec_path.exists():
raise SystemExit(f"OpenAPI spec not found at {spec_path}\nRun 'task engine:tool-models' to generate it.")
output_path = Path(args.output)
output_path = Path(args.output) if args.output else (repo_root / "src/stirling/models/tool_models.py")
with open(spec_path) as f:
spec = json.load(f)
write_models_module(output_path, specs)
print(f"Wrote {len(specs)} tool model specs")
result = ToolDiscovery(spec).discover()
models_code = generate_models_code(result.combined_schema)
write_output(output_path, result.tools, models_code)
print(f"Generated {len(result.tools)} tool models from {spec_path.name}")
for tool in result.tools:
print(f" {tool.enum_name}: {tool.path} → {tool.class_name}")
if __name__ == "__main__":
engine/scripts/setup_env.py
+13 -37
View File
@@ -1,48 +1,24 @@
"""
Copies .env from .env.example if missing, and errors if any keys from the example
are absent from the actual .env file.
Ensures `.env.local` exists so developers have a place to put overrides
(API keys, local model choices, etc.) without touching the committed `.env`.
Usage:
uv run scripts/setup_env.py
"""
import os
import shutil
import sys
from pathlib import Path
from dotenv import dotenv_values
ROOT = Path(__file__).parent.parent
EXAMPLE_FILE = ROOT / "config" / ".env.example"
ENV_FILE = ROOT / ".env"
ENV_LOCAL_FILE = ROOT / ".env.local"
print("setup-env: see engine/config/.env.example for documentation")
TEMPLATE = """\
###############################################################################
# Local overrides for `engine/.env`
# Put API keys and machine-specific settings here. Any variable defined here
# takes precedence over the committed `.env`
###############################################################################
"""
if not EXAMPLE_FILE.exists():
print(f"setup-env: {EXAMPLE_FILE.name} not found, skipping", file=sys.stderr)
sys.exit(0)
if not ENV_FILE.exists():
shutil.copy(EXAMPLE_FILE, ENV_FILE)
print("setup-env: created .env from .env.example")
env_keys = set(dotenv_values(ENV_FILE).keys()) | set(os.environ.keys())
example_keys = set(dotenv_values(EXAMPLE_FILE).keys())
missing = sorted(example_keys - env_keys)
if missing:
sys.exit(
"setup-env: .env is missing keys from .env.example:\n"
+ "\n".join(f" {k}" for k in missing)
+ "\n Add them manually or delete your local .env to re-copy from config/.env.example."
)
extra = sorted(k for k in dotenv_values(ENV_FILE) if k.startswith("STIRLING_") and k not in example_keys)
if extra:
print(
"setup-env: .env contains STIRLING_ keys not in config/.env.example:\n"
+ "\n".join(f" {k}" for k in extra)
+ "\n Add them to config/.env.example if they are intentional.",
file=sys.stderr,
)
if not ENV_LOCAL_FILE.exists():
ENV_LOCAL_FILE.write_text(TEMPLATE)
print("setup-env: created empty .env.local for local overrides")
engine/src/stirling/agents/ledger/__init__.py
+5
View File
@@ -0,0 +1,5 @@
"""Math Auditor Agent (mathAuditorAgent) — AI-powered math validation for PDF documents."""
from .agent import MathAuditorAgent
__all__ = ["MathAuditorAgent"]
engine/src/stirling/agents/ledger/agent.py
+550
View File
@@ -0,0 +1,550 @@
"""
Math Auditor Agent (mathAuditorAgent) — pydantic-ai agents for PDF math validation.
Examiner (Round 1, /api/v1/ai/math-auditor-agent/examine)
Receives a FolioManifest and returns a Requisition declaring what
Java must extract before validation can begin.
Audit pipeline (Round 2, /api/v1/ai/math-auditor-agent/deliberate)
Processes Evidence per-page:
1. Deterministic pass — ArithmeticScanner on every folio
2. Fast-model pass — extract named figures from each page
3. FigureTracker — cross-page consistency check
4. Fast-model call — generate human-readable summary
5. Assemble Verdict programmatically
Neither agent ever touches a PDF file. All content arrives pre-extracted
by Java, which owns the PDF from start to finish.
"""
from __future__ import annotations
import asyncio
import logging
from collections.abc import Coroutine
from decimal import Decimal, InvalidOperation
from typing import Any
from pydantic import BaseModel, Field
from pydantic_ai import Agent
from pydantic_ai.exceptions import AgentRunError
from stirling.contracts.ledger import (
Discrepancy,
DiscrepancyKind,
Evidence,
Folio,
FolioManifest,
Requisition,
Severity,
Verdict,
)
from stirling.logging import Pretty
from stirling.services import AppRuntime
from .prompts import (
EXAMINER_SYSTEM_PROMPT,
FIGURE_EXTRACTOR_PROMPT,
STATEMENT_VERIFIER_PROMPT,
SUMMARY_PROMPT,
TABLE_FORMULA_PROMPT,
)
from .validators import ArithmeticScanner, FigureTracker, FormulaEvaluator
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Structured output models for the per-page figure extractor
# ---------------------------------------------------------------------------
class ExtractedFigure(BaseModel):
"""A single named figure found on a page."""
label: str = Field(description="Normalised name, e.g. 'Total Revenue', 'VAT'.")
value: str = Field(description="Numeric value as a string, e.g. '1200.00'.")
raw: str = Field(description="Original text from the document, e.g. '£1,200.00'.")
class FigureExtractionResult(BaseModel):
"""All named figures found on a single page."""
figures: list[ExtractedFigure] = Field(default_factory=list)
class FormulaCheck(BaseModel):
"""One verifiable mathematical relationship in a table."""
description: str = Field(description="Human-readable, e.g. 'Line Total = Qty × Unit Price'")
formula: str = Field(description="Expression: 'col3 = col1 * col2' or 'cell(4,3) = sum(col3, 1-3)'")
scope: str = Field(description="'each_row' | 'column_total' | 'single_cell'")
row_range: list[int] | None = Field(default=None, description="Data rows to check (for each_row scope)")
target_row: int | None = Field(default=None, description="Row index of total (for column_total/single_cell)")
target_col: int | None = Field(default=None, description="Column index (for column_total/single_cell)")
class TableFormulas(BaseModel):
"""All verifiable formulas found in one table."""
formulas: list[FormulaCheck] = Field(default_factory=list)
class StatementCheck(BaseModel):
"""One prose claim and its verification result."""
claim: str = Field(description="The exact text of the claim")
verification: str = Field(description="Type: percentage_change, comparison, ratio, trend, average, other")
values_referenced: list[str] = Field(default_factory=list, description="Numbers used in the check")
expected_result: str = Field(description="What the calculation actually yields")
actual_claim: str = Field(description="What the text claims")
is_valid: bool = Field(description="True if the claim is correct within tolerance")
explanation: str = Field(description="One-line working showing the calculation")
class StatementsResult(BaseModel):
"""All verifiable prose claims found on a page."""
statements: list[StatementCheck] = Field(default_factory=list)
# ---------------------------------------------------------------------------
# MathAuditorAgent — main entry point, instantiated once at startup
# ---------------------------------------------------------------------------
class MathAuditorAgent:
"""
Encapsulates the Ledger Auditor pipeline.
Instantiated once at app startup with an AppRuntime, which provides
pre-built Model objects and ModelSettings.
"""
def __init__(self, runtime: AppRuntime) -> None:
fast_model = runtime.fast_model
model_settings = runtime.fast_model_settings
self._runtime = runtime
self._examiner = Agent(
model=fast_model,
deps_type=FolioManifest,
output_type=Requisition,
system_prompt=EXAMINER_SYSTEM_PROMPT,
model_settings=model_settings,
)
self._figure_extractor = Agent(
model=fast_model,
output_type=FigureExtractionResult,
system_prompt=FIGURE_EXTRACTOR_PROMPT,
model_settings=model_settings,
)
self._table_analyser = Agent(
model=fast_model,
output_type=TableFormulas,
system_prompt=TABLE_FORMULA_PROMPT,
model_settings=model_settings,
)
self._statement_verifier = Agent(
model=fast_model,
output_type=StatementsResult,
system_prompt=STATEMENT_VERIFIER_PROMPT,
model_settings=model_settings,
)
self._summary_agent = Agent(
model=fast_model,
output_type=str,
system_prompt=SUMMARY_PROMPT,
model_settings=model_settings,
)
self._llm_semaphore = asyncio.Semaphore(10)
# ------------------------------------------------------------------
# Round 1: Examine
# ------------------------------------------------------------------
async def examine(self, manifest: FolioManifest) -> Requisition:
"""Inspect a FolioManifest and declare the Requisition."""
logger.info(
"[math-auditor-agent] session=%s round=%d examining %d folios",
manifest.session_id,
manifest.round,
manifest.page_count,
)
user_prompt = "Examine this folio manifest and declare your requisition:\n" + manifest.model_dump_json()
logger.debug("REQUEST (examine)\n%s", Pretty({"user_prompt": user_prompt}))
result = await self._examiner.run(user_prompt, deps=manifest)
req = result.output
logger.debug("RESPONSE (examine)\n%s", Pretty(req.model_dump()))
logger.info(
"[math-auditor-agent] session=%s requisition: text=%s tables=%s ocr=%s",
manifest.session_id,
req.need_text,
req.need_tables,
req.need_ocr,
)
return req
# ------------------------------------------------------------------
# Round 2: Deliberate (deterministic-first pipeline)
# ------------------------------------------------------------------
async def audit(self, evidence: Evidence, tolerance: Decimal = Decimal("0.01")) -> Verdict:
"""
Audit the evidence using a deterministic-first pipeline:
1. Run ArithmeticScanner on every folio (no LLM)
2. Extract named figures per-page with fast model
3. Run FigureTracker cross-page consistency check (no LLM)
4. Generate human summary with fast model
5. Assemble Verdict
"""
return await self._audit_inner(evidence, tolerance)
async def _audit_inner(
self,
evidence: Evidence,
tolerance: Decimal,
) -> Verdict:
logger.info(
"[math-auditor-agent] session=%s round=%d auditing %d folios (final=%s)",
evidence.session_id,
evidence.round,
len(evidence.folios),
evidence.final_round,
)
all_discrepancies: list[Discrepancy] = []
pages_examined: list[int] = []
figure_tracker = FigureTracker(tolerance=tolerance)
# Step 1: Arithmetic scanning (deterministic, instant)
arithmetic_scanner = ArithmeticScanner(tolerance=tolerance)
for folio in evidence.folios:
pages_examined.append(folio.page)
text = folio.readable_text
if text and text.strip():
results = arithmetic_scanner.scan(folio.page, text)
all_discrepancies.extend(results)
logger.debug(
"TOOL (scan_arithmetic)\nArgs: %s\nResult: %s",
Pretty({"page": folio.page, "text_length": len(text)}),
Pretty([d.model_dump() for d in results]),
)
# Step 2: Parallel LLM calls — formula inference + figure extraction
# These are independent per-page so we fire them all concurrently.
formula_evaluator = FormulaEvaluator(tolerance=tolerance)
folios_with_text = [f for f in evidence.folios if f.readable_text.strip()]
# Collect all tables as (page, csv) pairs for formula inference
table_tasks: list[tuple[int, str]] = []
for folio in evidence.folios:
if folio.tables:
for table_csv in folio.tables:
table_tasks.append((folio.page, table_csv))
logger.info(
"[math-auditor-agent] session=%s step 2: %d formula + %d figure LLM calls (parallel)",
evidence.session_id,
len(table_tasks),
len(folios_with_text),
)
# Fire all LLM calls concurrently (bounded by _llm_semaphore)
formula_coros = [self._throttled(self._infer_formulas(csv)) for _, csv in table_tasks]
figure_coros = [self._throttled(self._extract_figures_for_page(f)) for f in folios_with_text]
statement_coros = [self._throttled(self._verify_statements(f)) for f in folios_with_text]
all_results = await asyncio.gather(
*formula_coros,
*figure_coros,
*statement_coros,
return_exceptions=True,
)
n_formulas = len(table_tasks)
n_figures = len(folios_with_text)
# Process formula results
for i, (page, table_csv) in enumerate(table_tasks):
result = all_results[i]
if isinstance(result, BaseException):
logger.warning("[math-auditor-agent] formula inference failed for page %d: %s", page, result)
continue
assert isinstance(result, TableFormulas)
formulas = result
if not formulas.formulas:
logger.info("[math-auditor-agent] page %d: no verifiable formulas found", page)
continue
for fc in formulas.formulas:
checked = formula_evaluator.evaluate(
page=page,
table_csv=table_csv,
formula=fc.formula,
scope=fc.scope,
description=fc.description,
row_range=fc.row_range,
target_row=fc.target_row,
target_col=fc.target_col,
)
all_discrepancies.extend(checked)
logger.debug(
"TOOL (check_formula)\nArgs: %s\nResult: %s",
Pretty({"page": page, "formula": fc.formula, "scope": fc.scope, "description": fc.description}),
Pretty([d.model_dump() for d in checked]),
)
# Process figure results
for i, folio in enumerate(folios_with_text):
result = all_results[n_formulas + i]
if isinstance(result, BaseException):
logger.warning("[math-auditor-agent] figure extraction failed for page %d: %s", folio.page, result)
continue
assert isinstance(result, list)
for fig, page in result:
try:
decimal_value = Decimal(fig.value.replace(",", "").strip())
except (InvalidOperation, ValueError):
logger.warning(
"[math-auditor-agent] skipping figure %r on page %d: non-numeric value %r",
fig.label,
page,
fig.value,
)
continue
figure_tracker.record(
label=fig.label,
value=decimal_value,
page=page,
raw=fig.raw,
)
# Process statement verification results
for i, folio in enumerate(folios_with_text):
result = all_results[n_formulas + n_figures + i]
if isinstance(result, BaseException):
logger.warning("[math-auditor-agent] statement verification failed for page %d: %s", folio.page, result)
continue
assert isinstance(result, StatementsResult)
stmts = result
for sc in stmts.statements:
if not sc.is_valid:
all_discrepancies.append(
Discrepancy(
page=folio.page,
kind=DiscrepancyKind.STATEMENT,
severity=Severity.ERROR,
description=f"{sc.claim}: {sc.explanation}",
stated=sc.actual_claim,
expected=sc.expected_result,
context=sc.claim,
)
)
logger.debug(
"TOOL (verify_statement)\nArgs: %s\nResult: %s",
Pretty({"page": folio.page, "claim": sc.claim}),
Pretty(sc.model_dump()),
)
logger.info(
"[math-auditor-agent] session=%s step 2 complete: %d figures registered",
evidence.session_id,
figure_tracker.entry_count,
)
# Step 3: Cross-page consistency — deterministic
consistency_discrepancies = figure_tracker.conflicts()
all_discrepancies.extend(consistency_discrepancies)
if consistency_discrepancies:
logger.debug(
"TOOL (check_figure_consistency)\nResult: %s",
Pretty([d.model_dump() for d in consistency_discrepancies]),
)
# Step 4: Summary — fast model, small payload
# Collect verification stats for the summary
total_tables = sum(len(f.tables) for f in evidence.folios if f.tables)
total_formulas_checked = sum(len(r.formulas) for r in all_results[:n_formulas] if isinstance(r, TableFormulas))
total_statements_checked = sum(
len(r.statements) for r in all_results[n_formulas + n_figures :] if isinstance(r, StatementsResult)
)
verification_stats = (
f"Verified: {len(pages_examined)} pages, {total_tables} tables "
f"({total_formulas_checked} formulas), "
f"{figure_tracker.entry_count} figures tracked, "
f"{total_statements_checked} prose claims checked."
)
logger.info(
"[math-auditor-agent] session=%s step 4: generating summary (%d discrepancies)",
evidence.session_id,
len(all_discrepancies),
)
pages_examined.sort()
summary = await self._generate_summary(
all_discrepancies,
pages_examined,
evidence.unauditable_pages,
verification_stats,
)
# Step 5: Assemble Verdict
error_count = sum(1 for d in all_discrepancies if d.severity == Severity.ERROR)
verdict = Verdict(
session_id=evidence.session_id,
discrepancies=all_discrepancies,
pages_examined=pages_examined,
rounds_taken=evidence.round,
summary=summary,
clean=error_count == 0,
unauditable_pages=evidence.unauditable_pages,
)
logger.debug("RESPONSE (deliberate)\n%s", Pretty(verdict.model_dump()))
logger.info(
"[math-auditor-agent] session=%s verdict: %d errors, %d warnings, clean=%s",
evidence.session_id,
verdict.error_count,
verdict.warning_count,
verdict.clean,
)
return verdict
# ------------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------------
async def _throttled[T](self, coro: Coroutine[Any, Any, T]) -> T:
"""Wrap a coroutine with the LLM concurrency semaphore."""
async with self._llm_semaphore:
return await coro
async def _infer_formulas(self, table_csv: str) -> TableFormulas:
"""Ask the fast model to infer verifiable formulas from a CSV table."""
try:
result = await self._table_analyser.run(f"CSV table:\n{table_csv}")
formulas = result.output
except AgentRunError:
logger.warning("[math-auditor-agent] formula inference failed, skipping table", exc_info=True)
formulas = TableFormulas(formulas=[])
logger.debug(
"TOOL (infer_formulas)\nArgs: %s\nResult: %s",
Pretty({"table_csv": table_csv[:300]}),
Pretty(formulas.model_dump()),
)
return formulas
async def _verify_statements(
self,
folio: Folio,
) -> StatementsResult:
"""Ask the fast model to find and verify prose claims on a page."""
text = folio.readable_text
if not text or not text.strip():
return StatementsResult(statements=[])
# Build context: page text + any table CSVs
prompt = f"Page {folio.page + 1} text:\n{text}"
if folio.tables:
prompt += "\n\nTable data on this page:\n"
for i, csv in enumerate(folio.tables):
prompt += f"\nTable {i + 1}:\n{csv}"
try:
result = await self._statement_verifier.run(prompt)
stmts = result.output
except AgentRunError:
logger.warning("[math-auditor-agent] statement verification failed for page %d", folio.page, exc_info=True)
stmts = StatementsResult(statements=[])
if stmts.statements:
logger.debug(
"TOOL (verify_statements)\nArgs: %s\nResult: %s",
Pretty({"page": folio.page, "text_length": len(text), "n_tables": len(folio.tables or [])}),
Pretty([s.model_dump() for s in stmts.statements]),
)
return stmts
async def _extract_figures_for_page(
self,
folio: Folio,
) -> list[tuple[ExtractedFigure, int]]:
text = folio.readable_text
if not text or not text.strip():
return []
logger.info("[math-auditor-agent] extracting figures from page %d (%d chars)", folio.page, len(text))
prompt = f"Page {folio.page + 1} text:\n{text}"
try:
result = await self._figure_extractor.run(prompt)
figures = result.output.figures
except AgentRunError:
logger.warning(
"[math-auditor-agent] figure extraction failed for page %d, skipping",
folio.page,
exc_info=True,
)
figures = []
logger.debug(
"TOOL (extract_figures)\nArgs: %s\nResult: %s",
Pretty({"page": folio.page, "text_length": len(text)}),
Pretty([f.model_dump() for f in figures]),
)
return [(fig, folio.page) for fig in figures]
async def _generate_summary(
self,
discrepancies: list[Discrepancy],
pages_examined: list[int],
unauditable_pages: list[int],
verification_stats: str,
) -> str:
error_count = sum(1 for d in discrepancies if d.severity == Severity.ERROR)
warning_count = sum(1 for d in discrepancies if d.severity == Severity.WARNING)
prompt = (
f"{verification_stats}\n"
f"Errors: {error_count}, Warnings: {warning_count}, "
f"Pages examined: {len(pages_examined)}, "
f"Unauditable pages: {unauditable_pages or 'none'}.\n"
)
if discrepancies:
prompt += "Discrepancies:\n"
for d in discrepancies:
prompt += f" - [{d.severity}] p{d.page + 1}: {d.description}\n"
try:
result = await self._summary_agent.run(prompt)
summary = result.output
except AgentRunError:
logger.warning("[math-auditor-agent] summary generation failed, using fallback", exc_info=True)
summary = self._fallback_summary(error_count, warning_count, pages_examined, unauditable_pages)
logger.debug("RESPONSE (summary)\n%s", Pretty({"summary": summary}))
return summary
@staticmethod
def _fallback_summary(
error_count: int,
warning_count: int,
pages_examined: list[int],
unauditable_pages: list[int],
) -> str:
parts = []
if error_count == 0 and warning_count == 0:
parts.append(f"No mathematical errors found across {len(pages_examined)} pages.")
else:
if error_count:
parts.append(f"Found {error_count} error{'s' if error_count != 1 else ''}.")
if warning_count:
parts.append(f"Found {warning_count} warning{'s' if warning_count != 1 else ''}.")
if unauditable_pages:
parts.append(
f"Pages {', '.join(str(p + 1) for p in unauditable_pages)} could not be audited (OCR unavailable)."
)
return " ".join(parts)
engine/src/stirling/agents/ledger/prompts.py
+147
View File
@@ -0,0 +1,147 @@
"""
Ledger Auditor — system prompts.
One prompt per role; keep them short and directive. Each agent is a
specialist with a narrow remit, not a general assistant.
"""
EXAMINER_SYSTEM_PROMPT = """\
You are the Examiner, the first stage of the Ledger Auditor pipeline.
You receive a FolioManifest: a list of page types (text / image / mixed) \
for a PDF document. Your sole task is to declare exactly which pages you \
need Java to extract content from so that the Auditor can verify the \
document's mathematics.
Rules:
- Request BOTH text AND table extraction for every 'text' or 'mixed' page. \
Tables are critical — the Auditor cannot verify totals without them. \
Tabula extraction is cheap; missing a table is not.
- Request OCR for any page classified as 'image' or 'mixed' (PDFBox cannot \
read image-only content).
- Be conservative — if in doubt, request the page. False negatives \
(missed errors) are worse than false positives (wasted extraction).
- Do not request pages that are clearly decorative (cover pages, blank pages) \
unless you cannot tell from the manifest alone.
- Return a Requisition with your page lists and a plain-English rationale \
that will appear in server logs.
"""
FIGURE_EXTRACTOR_PROMPT = """\
You are a figure extractor for financial document auditing.
You receive the text content of a single PDF page. Your task is to \
identify every significant named numeric figure on the page.
A "named figure" is a labelled number that could appear elsewhere in \
the document under the same name — for example:
"Total Revenue: £1,200,000"
"Net Profit $45,000"
"VAT (20%): 240.00"
"Subtotal ......... 3,500"
For each figure, return:
- label: a normalised name (e.g. "Total Revenue", "Net Profit", "VAT")
- value: the numeric value as a plain decimal string (e.g. "1200000")
- raw: the original text as it appears in the document
Rules:
- Only extract figures that have a clear label/name attached.
- Do not extract bare numbers without context.
- Strip currency symbols and thousands separators from value.
- If a figure appears multiple times on the same page, extract each.
- Return an empty list if no named figures are found.
- Be precise — do not invent figures that are not in the text.
"""
TABLE_FORMULA_PROMPT = """\
You are a table formula analyser for financial document auditing.
You receive a CSV table extracted from a PDF. Your task is to identify \
every verifiable mathematical relationship between cells.
Relationships fall into three scopes:
1. "each_row" — a formula that should hold for every data row.
Example: "col3 = col1 * col2" (Line Total = Qty × Unit Price)
2. "column_total" — a total row where cells = sum of the column above.
Example: a Subtotal row where each cell sums the column.
3. "single_cell" — one specific cell computed from others.
Example: "cell(5,3) = cell(4,3) * 0.1" (Tax = Subtotal × 10%)
Formula syntax (use exactly this):
- Column references: col0, col1, col2 ... (0-indexed)
- Cell references: cell(row, col) — 0-indexed, header is row 0
- Operators: + - * /
- sum(colN, start-end) — sum of colN from row start to row end inclusive
- Decimal numbers: 0.1, 100, etc.
Rules:
- Row 0 is the header. First data row is row 1.
- Include the left-hand side: "col3 = col1 * col2" not just "col1 * col2"
- For column_total scope, set target_row to the total row index. \
Set target_col to a specific column or null to check all numeric columns.
- For each_row scope, set row_range to the data rows (exclude header \
and total rows).
- Only return formulas you are confident about. Skip columns/rows \
where the relationship is unclear.
- Return an empty list if the table has no verifiable math.
"""
STATEMENT_VERIFIER_PROMPT = """\
You are a statement verifier for financial document auditing.
You receive the text of a single PDF page, plus any table data from \
that page. Your task is to find prose claims that make mathematical \
assertions, and verify whether each claim is correct.
A "verifiable claim" is a sentence that states a mathematical fact \
about numbers present on the page or derivable from the data. Examples:
- "Revenue grew 15% year-over-year"
- "Costs decreased month on month"
- "Department A represents 40% of total spend"
- "Net margin improved to 12.4%"
- "Average transaction value was $250"
For each claim you find:
1. Identify the numbers referenced in the claim
2. Perform the calculation yourself using the data on the page
3. Compare your result to what the claim states
4. Determine if the claim is valid (within reasonable rounding)
Return:
- claim: the exact text of the claim
- verification: the type — "percentage_change", "comparison", \
"ratio", "trend", "average", or "other"
- values_referenced: the specific numbers used in your check
- expected_result: what the calculation actually yields
- actual_claim: what the text claims
- is_valid: true if the claim is correct within 1% tolerance
- explanation: show your working, one line
Rules:
- Only check claims that can be verified from data on this page.
- If a claim references data not on the page, skip it.
- "Decreased month on month" means EVERY consecutive pair decreased.
- Percentage claims allow 1% absolute tolerance (14.8% ≈ 15%).
- Return an empty list if there are no verifiable claims.
- Do not fabricate claims that are not in the text.
"""
SUMMARY_PROMPT = """\
You are a summary writer for a PDF math audit tool.
You receive a list of discrepancies (errors and warnings) found in a \
document, plus coverage statistics and a breakdown of what was verified. \
Write a two to three sentence summary suitable for an end user.
Rules:
- Start with what was verified: e.g. "Audited 6 pages: checked 4 tables \
(12 formulas), scanned 6 pages for arithmetic, extracted 20 figures \
for cross-page consistency, and verified 3 prose claims."
- Then state the outcome: errors found or clean.
- Mention unauditable pages if any exist.
- Be concise and factual. Do not repeat individual discrepancy details.
"""
engine/src/stirling/agents/ledger/validators/__init__.py
+5
View File
@@ -0,0 +1,5 @@
from .arithmetic import ArithmeticScanner
from .figures import FigureTracker
from .formula import FormulaEvaluator
__all__ = ["ArithmeticScanner", "FigureTracker", "FormulaEvaluator"]
engine/src/stirling/agents/ledger/validators/_parsing.py
+31
View File
@@ -0,0 +1,31 @@
"""Shared parsing helpers for ledger validators."""
from __future__ import annotations
import csv
import io
import re
from decimal import Decimal, InvalidOperation
# Strip common currency symbols and thousands separators before parsing.
STRIP_PATTERN = re.compile(r"[£$€¥,\s]")
def to_decimal(raw: str) -> Decimal | None:
"""Parse a cell value to Decimal, returning None for non-numeric cells."""
cleaned = STRIP_PATTERN.sub("", raw.strip())
if not cleaned or cleaned in {"-", "—", "n/a", "N/A", "na", "NA"}:
return None
# Handle parenthesised negatives: (123.45) → -123.45
if cleaned.startswith("(") and cleaned.endswith(")"):
cleaned = "-" + cleaned[1:-1]
try:
return Decimal(cleaned)
except InvalidOperation:
return None
def parse_csv(table_csv: str) -> list[list[str]]:
"""Parse a CSV string into rows, dropping completely empty rows."""
reader = csv.reader(io.StringIO(table_csv.strip()))
return [row for row in reader if any(cell.strip() for cell in row)]
engine/src/stirling/agents/ledger/validators/arithmetic.py
+152
View File
@@ -0,0 +1,152 @@
"""
ArithmeticScanner — finds and verifies inline arithmetic expressions in text.
Targets patterns commonly found in financial documents:
"100 + 200 + 150 = 450"
"Total: 1,250 (500 + 400 + 350)"
"Net profit of £1,200 (£2,000 revenue less £800 costs)"
All arithmetic is performed in Decimal. The scanner does not use an LLM —
it is a deterministic regex-and-eval pipeline.
"""
from __future__ import annotations
import logging
import re
from decimal import Decimal
from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity
from ._parsing import STRIP_PATTERN as _STRIP
from ._parsing import to_decimal as _to_decimal
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Regex patterns
# ---------------------------------------------------------------------------
# Currency / number token: optional sign, optional currency symbol,
# digits with optional thousands separator and decimal point.
_NUM = r"[£$€¥]?-?[\d,]+(?:\.\d+)?"
# "A + B + C = D" or "A + B + C = D" with arbitrary spacing
_EQUALS_EXPR = re.compile(
rf"({_NUM}(?:\s*[+\-]\s*{_NUM})+)\s*=\s*({_NUM})",
re.IGNORECASE,
)
# "Total: X (A + B + C)" — the total comes before the addends
_TOTAL_THEN_ADDENDS = re.compile(
rf"(?:total|sum|grand total|subtotal)\s*[:\-]?\s*({_NUM})\s*\(({_NUM}(?:\s*[+\-]\s*{_NUM})+)\)",
re.IGNORECASE,
)
def _parse(token: str) -> Decimal | None:
"""Parse a regex-matched token to Decimal."""
return _to_decimal(token)
def _eval_expression(expr: str) -> Decimal | None:
"""
Evaluate a simple additive expression of the form A +/- B +/- C ...
Returns None if the expression cannot be parsed.
"""
# Tokenise: split on + or -, keep the operator.
tokens = re.split(r"([+\-])", _STRIP.sub("", expr.strip()))
result = Decimal(0)
operator = "+"
for token in tokens:
token = token.strip()
if not token:
continue # skip empty tokens (e.g. from leading negative)
if token in ("+", "-"):
operator = token
continue
val = _parse(token)
if val is None:
return None
result = result + val if operator == "+" else result - val
return result
class ArithmeticScanner:
"""
Scans a block of text for arithmetic expressions and checks them.
Parameters
----------
tolerance:
Maximum absolute difference before an expression is flagged as wrong.
"""
def __init__(self, tolerance: Decimal = Decimal("0.01")) -> None:
self.tolerance = tolerance
def scan(self, page: int, text: str) -> list[Discrepancy]:
"""
Find all verifiable arithmetic expressions in *text* and return
a Discrepancy for each one that does not balance within tolerance.
"""
discrepancies: list[Discrepancy] = []
discrepancies.extend(self._check_equals_expressions(page, text))
discrepancies.extend(self._check_total_then_addends(page, text))
return discrepancies
# ------------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------------
def _check_equals_expressions(self, page: int, text: str) -> list[Discrepancy]:
"""Handle patterns like '100 + 200 = 300'."""
found: list[Discrepancy] = []
for match in _EQUALS_EXPR.finditer(text):
expr_str = match.group(1)
stated_str = match.group(2)
computed = _eval_expression(expr_str)
stated = _parse(stated_str)
if computed is None or stated is None:
continue
if abs(computed - stated) > self.tolerance:
found.append(
Discrepancy(
page=page,
kind=DiscrepancyKind.ARITHMETIC,
severity=Severity.ERROR,
description=f"Arithmetic error: {expr_str.strip()} should equal {computed}, not {stated}",
stated=str(stated),
expected=str(computed),
context=match.group(0),
)
)
return found
def _check_total_then_addends(self, page: int, text: str) -> list[Discrepancy]:
"""Handle patterns like 'Total: 450 (100 + 200 + 150)'."""
found: list[Discrepancy] = []
for match in _TOTAL_THEN_ADDENDS.finditer(text):
stated_str = match.group(1)
expr_str = match.group(2)
stated = _parse(stated_str)
computed = _eval_expression(expr_str)
if stated is None or computed is None:
continue
if abs(computed - stated) > self.tolerance:
found.append(
Discrepancy(
page=page,
kind=DiscrepancyKind.ARITHMETIC,
severity=Severity.ERROR,
description=f"Stated total {stated} does not match addends ({expr_str.strip()} = {computed})",
stated=str(stated),
expected=str(computed),
context=match.group(0),
)
)
return found
engine/src/stirling/agents/ledger/validators/figures.py
+98
View File
@@ -0,0 +1,98 @@
"""
FigureTracker — cross-page consistency checker for named figures.
Collects named numeric figures as the auditor encounters them (e.g.
"Total Revenue: £1,200,000") and surfaces any that appear under the same
label but with a different value on another page — a classic symptom of
copy-paste errors or stale data in executive summaries.
The tracker is intentionally simple: normalise labels, compare values
within tolerance, emit Discrepancy for each conflict.
"""
from __future__ import annotations
import logging
import re
from decimal import Decimal
from pydantic import BaseModel
from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity
logger = logging.getLogger(__name__)
class FigureRecord(BaseModel):
"""A named numeric figure seen on a specific page."""
label: str
value: Decimal
page: int
raw: str
# Strip punctuation that varies between contexts ("revenue:" vs "revenue —")
_LABEL_NOISE = re.compile(r"[:\-—\s]+")
def _normalise_label(label: str) -> str:
return _LABEL_NOISE.sub(" ", label.lower()).strip()
class FigureTracker:
"""
Accumulates named figures during an audit and checks them for consistency.
Typical usage:
tracker = FigureTracker()
tracker.record("Net Profit", Decimal("1200.00"), page=3, raw="£1,200.00")
tracker.record("Net Profit", Decimal("1250.00"), page=7, raw="£1,250.00")
discrepancies = tracker.conflicts() # returns one Discrepancy
"""
def __init__(self, tolerance: Decimal = Decimal("0.01")) -> None:
self.tolerance = tolerance
self._ledger: dict[str, list[FigureRecord]] = {}
def record(self, label: str, value: Decimal, page: int, raw: str) -> None:
"""Register a named figure sighting."""
key = _normalise_label(label)
self._ledger.setdefault(key, []).append(FigureRecord(label=key, value=value, page=page, raw=raw))
def conflicts(self) -> list[Discrepancy]:
"""
Return a Discrepancy for every label that has sightings whose value
differs from the first-seen (canonical) value by more than tolerance.
O(n) per label — each record is compared against the canonical only.
"""
discrepancies: list[Discrepancy] = []
for label, records in self._ledger.items():
if len(records) < 2:
continue
canonical = records[0]
for other in records[1:]:
if abs(canonical.value - other.value) > self.tolerance:
discrepancies.append(
Discrepancy(
page=other.page,
kind=DiscrepancyKind.CONSISTENCY,
severity=Severity.WARNING,
description=(
f'"{label}" stated as {canonical.raw} on page'
f" {canonical.page + 1}"
f" but {other.raw} on page {other.page + 1}"
),
stated=other.raw,
expected=canonical.raw,
context=(f"First seen: page {canonical.page + 1} | Later: page {other.page + 1}"),
)
)
return discrepancies
@property
def entry_count(self) -> int:
return sum(len(v) for v in self._ledger.values())
engine/src/stirling/agents/ledger/validators/formula.py
+375
View File
@@ -0,0 +1,375 @@
"""
FormulaEvaluator — verifies LLM-inferred formulas against CSV table data.
Supports a safe expression syntax:
- Column refs: col0, col1, col2 ...
- Cell refs: cell(row, col)
- Operators: + - * /
- Functions: sum(colN, rows start-end)
All arithmetic is Decimal. No eval(), no arbitrary code execution.
"""
from __future__ import annotations
import logging
import re
from decimal import Decimal, InvalidOperation
from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity
from ._parsing import parse_csv as _parse_csv
from ._parsing import to_decimal as _to_decimal
logger = logging.getLogger(__name__)
class FormulaEvaluator:
"""
Evaluates formula expressions against parsed CSV rows.
Formulas use a simple syntax:
"col3 = col1 * col2" — per-row check
"cell(4,3) = sum(col3, 1-3)" — single cell check
"""
def __init__(self, tolerance: Decimal = Decimal("0.01")) -> None:
self.tolerance = tolerance
def evaluate(
self,
page: int,
table_csv: str,
formula: str,
scope: str,
description: str,
row_range: list[int] | None = None,
target_row: int | None = None,
target_col: int | None = None,
) -> list[Discrepancy]:
"""
Evaluate a formula against table data.
scope: "each_row" | "column_total" | "single_cell"
"""
rows = _parse_csv(table_csv)
if len(rows) < 2:
return []
if scope == "each_row":
return self._check_each_row(page, rows, formula, description, row_range)
elif scope == "column_total":
return self._check_column_total(page, rows, formula, description, target_row, target_col)
elif scope == "single_cell":
return self._check_single_cell(page, rows, formula, description, target_row, target_col)
else:
logger.warning("[formula] unknown scope %r, skipping", scope)
return []
def _check_each_row(
self,
page: int,
rows: list[list[str]],
formula: str,
description: str,
row_range: list[int] | None,
) -> list[Discrepancy]:
"""Verify formula holds for each data row."""
discrepancies: list[Discrepancy] = []
# Parse "colX = expr" format
parts = formula.split("=", 1)
if len(parts) != 2:
return []
lhs = parts[0].strip()
rhs = parts[1].strip()
lhs_col = self._parse_col_ref(lhs)
if lhs_col is None:
return []
check_rows = row_range if row_range else list(range(1, len(rows)))
for row_idx in check_rows:
if row_idx >= len(rows):
continue
row = rows[row_idx]
stated = self._get_cell(row, lhs_col)
if stated is None:
continue
computed = self._eval_row_expr(rhs, row, rows)
if computed is None:
continue
if abs(stated - computed) > self.tolerance:
discrepancies.append(
Discrepancy(
page=page,
kind=DiscrepancyKind.TALLY,
severity=Severity.ERROR,
description=f"{description}: row {row_idx} — stated {stated}, expected {computed}",
stated=str(stated),
expected=str(computed),
context=f"row {row_idx}, {formula}",
)
)
return discrepancies
def _check_column_total(
self,
page: int,
rows: list[list[str]],
formula: str,
description: str,
target_row: int | None,
target_col: int | None,
) -> list[Discrepancy]:
"""Verify that a total row contains correct column sums."""
if target_row is None or target_row >= len(rows):
return []
discrepancies: list[Discrepancy] = []
total_row = rows[target_row]
# Determine which columns to check
cols_to_check: list[int] = []
if target_col is not None:
cols_to_check = [target_col]
else:
# Check all numeric columns in the total row
cols_to_check = list(range(len(total_row)))
# Determine addend rows (all rows between header and total row)
addend_rows = list(range(1, target_row))
for col in cols_to_check:
stated = self._get_cell(total_row, col)
if stated is None:
continue
computed = Decimal(0)
has_addends = False
for r_idx in addend_rows:
if r_idx >= len(rows):
continue
val = self._get_cell(rows[r_idx], col)
if val is not None:
computed += val
has_addends = True
if not has_addends:
continue
if abs(stated - computed) > self.tolerance:
discrepancies.append(
Discrepancy(
page=page,
kind=DiscrepancyKind.TALLY,
severity=Severity.ERROR,
description=f"{description}: column {col} — stated {stated}, expected {computed}",
stated=str(stated),
expected=str(computed),
context=f"column {col}, total row {target_row}",
)
)
return discrepancies
def _check_single_cell(
self,
page: int,
rows: list[list[str]],
formula: str,
description: str,
target_row: int | None,
target_col: int | None,
) -> list[Discrepancy]:
"""Verify a single cell formula (e.g. Grand Total = Subtotal + Tax)."""
parts = formula.split("=", 1)
if len(parts) != 2:
return []
# Parse target from LHS cell(r,c) if not provided explicitly
if target_row is None or target_col is None:
lhs_match = re.match(r"cell\(\s*(\d+)\s*,\s*(\d+)\s*\)", parts[0].strip())
if lhs_match:
target_row = int(lhs_match.group(1))
target_col = int(lhs_match.group(2))
else:
return []
if target_row >= len(rows):
return []
rhs = parts[1].strip()
stated = self._get_cell(rows[target_row], target_col)
if stated is None:
return []
computed = self._eval_row_expr(rhs, rows[target_row], rows)
if computed is None:
return []
if abs(stated - computed) > self.tolerance:
return [
Discrepancy(
page=page,
kind=DiscrepancyKind.TALLY,
severity=Severity.ERROR,
description=f"{description}: stated {stated}, expected {computed}",
stated=str(stated),
expected=str(computed),
context=f"cell({target_row},{target_col}), {formula}",
)
]
return []
# ------------------------------------------------------------------
# Expression evaluation
# ------------------------------------------------------------------
def _eval_row_expr(self, expr: str, row: list[str], all_rows: list[list[str]]) -> Decimal | None:
"""
Evaluate an expression in the context of a specific row.
Supports: colN refs, cell(r,c) refs, +, -, *, /
Also supports: sum(colN, start-end)
"""
# Handle sum() function first
sum_pattern = re.compile(r"sum\(\s*col(\d+)\s*,\s*(\d+)\s*-\s*(\d+)\s*\)")
resolved = expr
for match in sum_pattern.finditer(expr):
col = int(match.group(1))
start = int(match.group(2))
end = int(match.group(3))
total = Decimal(0)
for r_idx in range(start, end + 1):
if r_idx < len(all_rows):
val = self._get_cell(all_rows[r_idx], col)
if val is not None:
total += val
resolved = resolved.replace(match.group(0), str(total))
# Handle cell(r, c) references
cell_pattern = re.compile(r"cell\(\s*(\d+)\s*,\s*(\d+)\s*\)")
for match in cell_pattern.finditer(resolved):
r = int(match.group(1))
c = int(match.group(2))
if r < len(all_rows):
val = self._get_cell(all_rows[r], c)
if val is not None:
resolved = resolved.replace(match.group(0), str(val))
else:
return None
else:
return None
# Replace colN references with values from the current row.
# Use re.sub with word boundaries to avoid col1 corrupting col12.
_failed = False
def _col_replacer(m: re.Match[str]) -> str:
nonlocal _failed
col_idx = int(m.group(1))
val = self._get_cell(row, col_idx)
if val is None:
_failed = True
return m.group(0)
return str(val)
resolved = re.sub(r"\bcol(\d+)\b", _col_replacer, resolved)
if _failed:
return None
# Evaluate the resulting arithmetic expression safely
return self._safe_eval(resolved)
def _safe_eval(self, expr: str) -> Decimal | None:
"""
Evaluate a simple arithmetic expression containing only
numbers and +, -, *, / operators. Respects standard operator
precedence (* and / bind tighter than + and -). No eval().
"""
try:
raw = re.findall(r"\d+(?:\.\d+)?|[+\-*/]", expr.strip())
if not raw:
return None
# Build (values, ops) lists, merging a leading '-' or an
# operator-adjacent '-' into the next number token.
values: list[Decimal] = []
ops: list[str] = []
i = 0
while i < len(raw):
tok = raw[i]
if tok in "+-*/" and not values and tok == "-":
# Leading negative: merge with next number
i += 1
if i >= len(raw):
return None
values.append(Decimal("-" + raw[i]))
elif tok in "+-*/":
# Operator followed by '-' → negative operand
if (
tok in "+-*/"
and i + 1 < len(raw)
and raw[i + 1] == "-"
and i + 2 < len(raw)
and raw[i + 2] not in "+-*/"
):
ops.append(tok)
values.append(Decimal("-" + raw[i + 2]))
i += 2 # skip the '-' and the number
else:
ops.append(tok)
else:
values.append(Decimal(tok))
i += 1
if not values:
return None
# Pass 1: evaluate * and /
j = 0
while j < len(ops):
if ops[j] in ("*", "/"):
if ops[j] == "*":
values[j] = values[j] * values[j + 1]
else:
if values[j + 1] == 0:
return None
values[j] = values[j] / values[j + 1]
values.pop(j + 1)
ops.pop(j)
else:
j += 1
# Pass 2: evaluate + and -
result = values[0]
for j, op in enumerate(ops):
if op == "+":
result += values[j + 1]
elif op == "-":
result -= values[j + 1]
return result
except (InvalidOperation, IndexError, ValueError):
return None
# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------
@staticmethod
def _parse_col_ref(ref: str) -> int | None:
match = re.match(r"col(\d+)", ref.strip())
return int(match.group(1)) if match else None
@staticmethod
def _get_cell(row: list[str], col: int) -> Decimal | None:
if col >= len(row):
return None
return _to_decimal(row[col])
engine/src/stirling/agents/orchestrator.py
+111 -19
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import assert_never
from pydantic_ai import Agent
from pydantic_ai.output import ToolOutput
@@ -13,15 +14,20 @@ from stirling.agents.user_spec import UserSpecAgent
from stirling.contracts import (
AgentDraftRequest,
AgentDraftWorkflowResponse,
ExtractedTextArtifact,
OrchestratorRequest,
OrchestratorResponse,
PdfEditRequest,
PdfEditResponse,
PdfQuestionRequest,
PdfQuestionResponse,
SupportedCapability,
ToolOperationStep,
UnsupportedCapabilityResponse,
format_conversation_history,
)
from stirling.contracts.form_fill import KnowledgeUpdateResponse
from stirling.contracts.pdf_edit import EditPlanResponse
from stirling.models.agent_tool_models import AgentToolId, MathAuditorAgentParams
from stirling.services import AppRuntime
@@ -57,6 +63,14 @@ class OrchestratorAgent:
name="delegate_user_spec",
description="Delegate requests to create or revise a user agent spec and return the draft result.",
),
ToolOutput(
self.math_auditor_agent,
name="math_auditor_agent",
description=(
"Delegate requests to check arithmetic, validate table totals, "
"audit financial calculations, or verify mathematical accuracy in PDFs."
),
),
ToolOutput(
self.unsupported_capability,
name="unsupported_capability",
@@ -67,44 +81,91 @@ class OrchestratorAgent:
system_prompt=(
"You are the top-level orchestrator. "
"Choose exactly one output function that best handles the request. "
"Use delegate_pdf_edit for requested PDF modifications. "
"Use delegate_pdf_question for questions about the contents of a PDF. "
"Use delegate_form_fill for requests to extract personal information for form filling. "
"Use delegate_pdf_edit for requested modifications of single or multiple PDFs. "
"Use delegate_pdf_question for questions about PDF contents. "
"Use delegate_user_spec for requests to create or define an agent spec. "
"Use math_auditor_agent for requests to check arithmetic, validate "
"table totals, audit financial calculations, or verify math in PDFs. "
"Use unsupported_capability only when none of the other outputs fit."
),
model_settings=runtime.fast_model_settings,
)
async def handle(self, request: OrchestratorRequest) -> OrchestratorResponse:
if request.resume_with is not None:
return await self._resume(request, request.resume_with)
result = await self.agent.run(
request.user_message,
self._build_prompt(request),
deps=OrchestratorDeps(runtime=self.runtime, request=request),
)
return result.output
async def delegate_form_fill(self, ctx: RunContext[OrchestratorDeps]) -> KnowledgeUpdateResponse:
request = ctx.deps.request
return await DocumentExtractorAgent(ctx.deps.runtime).extract_single(
document_text=request.user_message,
user_message=request.user_message,
)
async def _resume(self, request: OrchestratorRequest, capability: SupportedCapability) -> OrchestratorResponse:
"""Fast-path to get back to the correct endpoint without having to call AI."""
match capability:
case SupportedCapability.PDF_QUESTION:
return await self._run_pdf_question(request)
case SupportedCapability.PDF_EDIT:
return await self._run_pdf_edit(request)
case SupportedCapability.AGENT_DRAFT:
return await self._run_agent_draft(request)
case (
SupportedCapability.ORCHESTRATE
| SupportedCapability.AGENT_REVISE
| SupportedCapability.AGENT_NEXT_ACTION
| SupportedCapability.MATH_AUDITOR_AGENT
):
raise ValueError(f"Cannot resume orchestrator with capability: {capability}")
case _ as unreachable:
assert_never(unreachable)
async def delegate_pdf_edit(self, ctx: RunContext[OrchestratorDeps]) -> PdfEditResponse:
request = ctx.deps.request
return await PdfEditAgent(ctx.deps.runtime).handle(
PdfEditRequest(user_message=request.user_message, conversation_id=request.conversation_id)
return await self._run_pdf_edit(ctx.deps.request)
async def _run_pdf_edit(self, request: OrchestratorRequest) -> PdfEditResponse:
return await PdfEditAgent(self.runtime).handle(
PdfEditRequest(
user_message=request.user_message,
file_names=request.file_names,
conversation_history=request.conversation_history,
)
)
async def delegate_pdf_question(self, ctx: RunContext[OrchestratorDeps]) -> PdfQuestionResponse:
request = ctx.deps.request
return await PdfQuestionAgent(ctx.deps.runtime).handle(
PdfQuestionRequest(question=request.user_message, conversation_id=request.conversation_id)
return await self._run_pdf_question(ctx.deps.request)
async def _run_pdf_question(self, request: OrchestratorRequest) -> PdfQuestionResponse:
extracted_text = self._get_extracted_text_artifact(request)
return await PdfQuestionAgent(self.runtime).handle(
PdfQuestionRequest(
question=request.user_message,
file_names=request.file_names,
page_text=extracted_text.files if extracted_text is not None else [],
conversation_history=request.conversation_history,
)
)
async def delegate_user_spec(self, ctx: RunContext[OrchestratorDeps]) -> AgentDraftWorkflowResponse:
request = ctx.deps.request
return await UserSpecAgent(ctx.deps.runtime).draft(AgentDraftRequest(user_message=request.user_message))
return await self._run_agent_draft(ctx.deps.request)
async def _run_agent_draft(self, request: OrchestratorRequest) -> AgentDraftWorkflowResponse:
return await UserSpecAgent(self.runtime).draft(
AgentDraftRequest(
user_message=request.user_message,
conversation_history=request.conversation_history,
)
)
async def math_auditor_agent(self, ctx: RunContext[OrchestratorDeps]) -> EditPlanResponse:
return EditPlanResponse(
summary="Validate mathematical calculations in the document.",
steps=[
ToolOperationStep(
tool=AgentToolId.MATH_AUDITOR_AGENT,
parameters=MathAuditorAgentParams(),
)
],
)
async def unsupported_capability(
self,
@@ -113,3 +174,34 @@ class OrchestratorAgent:
message: str,
) -> UnsupportedCapabilityResponse:
return UnsupportedCapabilityResponse(capability=capability, message=message)
def _get_extracted_text_artifact(self, request: OrchestratorRequest) -> ExtractedTextArtifact | None:
for artifact in request.artifacts:
if isinstance(artifact, ExtractedTextArtifact):
return artifact
return None
def _build_prompt(self, request: OrchestratorRequest) -> str:
artifact_summary = self._describe_artifacts(request)
file_names = ", ".join(request.file_names) if request.file_names else "Unknown files"
history = format_conversation_history(request.conversation_history)
return (
f"Conversation history:\n{history}\n"
f"User message: {request.user_message}\n"
f"Files: {file_names}\n"
f"Available artifacts:\n{artifact_summary}"
)
def _describe_artifacts(self, request: OrchestratorRequest) -> str:
if not request.artifacts:
return "- none"
descriptions: list[str] = []
for artifact in request.artifacts:
if isinstance(artifact, ExtractedTextArtifact):
total_pages = sum(len(f.pages) for f in artifact.files)
file_names = [f.file_name for f in artifact.files]
descriptions.append(f"- extracted_text: {total_pages} pages from {file_names}")
continue
descriptions.append("- unknown artifact")
return "\n".join(descriptions)
engine/src/stirling/agents/pdf_edit.py
+10 -8
View File
@@ -13,14 +13,15 @@ from stirling.contracts import (
PdfEditRequest,
PdfEditResponse,
ToolOperationStep,
format_conversation_history,
)
from stirling.models import OPERATIONS, ApiModel, OperationId, ParamToolModel
from stirling.models import OPERATIONS, ApiModel, ParamToolModel, ToolEndpoint
from stirling.services import AppRuntime
class PdfEditPlanSelection(ApiModel):
outcome: Literal["plan"] = "plan"
operations: list[OperationId] = Field(min_length=1)
operations: list[ToolEndpoint] = Field(min_length=1)
summary: str
rationale: str | None = None
@@ -41,7 +42,7 @@ class PdfEditParameterSelector:
async def select(
self,
request: PdfEditRequest,
operation_plan: list[OperationId],
operation_plan: list[ToolEndpoint],
operation_index: int,
generated_steps: list[ToolOperationStep],
) -> ParamToolModel:
@@ -51,7 +52,7 @@ class PdfEditParameterSelector:
self._build_parameter_prompt(request, operation_plan, operation_index, generated_steps),
output_type=NativeOutput(parameter_model),
instructions=(
f"Generate only the parameters for the PDF operation `{operation_id.value}`. "
f"Generate only the parameters for the PDF operation `{operation_id.name}`. "
"Do not include fields from any other operation."
),
)
@@ -60,12 +61,12 @@ class PdfEditParameterSelector:
def _build_parameter_prompt(
self,
request: PdfEditRequest,
operation_plan: list[OperationId],
operation_plan: list[ToolEndpoint],
operation_index: int,
generated_steps: list[ToolOperationStep],
) -> str:
operation_id = operation_plan[operation_index]
operation_list = ", ".join(operation.value for operation in operation_plan)
operation_list = ", ".join(operation.name for operation in operation_plan)
file_names = ", ".join(request.file_names) if request.file_names else "No file names were provided."
generated_steps_text = (
"\n".join(
@@ -79,7 +80,7 @@ class PdfEditParameterSelector:
f"Files: {file_names}\n"
f"Operation plan: {operation_list}\n"
f"Selected operation index: {operation_index + 1} of {len(operation_plan)}\n"
f"Selected operation: {operation_id.value}\n"
f"Selected operation: {operation_id.name}\n"
f"Already generated steps:\n{generated_steps_text}\n"
"Return only the parameter object for the selected operation."
)
@@ -146,6 +147,7 @@ class PdfEditAgent:
def _build_selection_prompt(self, request: PdfEditRequest) -> str:
file_names = ", ".join(request.file_names) if request.file_names else "No file names were provided."
return (
f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n"
f"User request: {request.user_message}\n"
f"Files: {file_names}\n"
f"Supported operations: {self._supported_operations_prompt()}\n"
@@ -153,4 +155,4 @@ class PdfEditAgent:
)
def _supported_operations_prompt(self) -> str:
return ", ".join(operation_id.value for operation_id in self.supported_operations)
return ", ".join(f"{op.name} ({op.value})" for op in self.supported_operations)
engine/src/stirling/agents/pdf_questions.py
+42 -8
View File
@@ -4,18 +4,26 @@ from pydantic_ai import Agent
from pydantic_ai.output import NativeOutput
from stirling.contracts import (
ExtractedFileText,
NeedContentFileRequest,
PdfContentType,
PdfQuestionAnswerResponse,
PdfQuestionNeedTextResponse,
PdfQuestionNeedContentResponse,
PdfQuestionNotFoundResponse,
PdfQuestionRequest,
PdfQuestionResponse,
format_conversation_history,
)
from stirling.services import AppRuntime
class PdfQuestionAgent:
DEFAULT_MAX_PAGES = 12
DEFAULT_MAX_CHARACTERS = 24_000
def __init__(self, runtime: AppRuntime) -> None:
self.runtime = runtime
rag = runtime.rag_capability
self.agent = Agent(
model=runtime.smart_model,
output_type=NativeOutput(
@@ -25,18 +33,29 @@ class PdfQuestionAgent:
]
),
system_prompt=(
"Answer questions about a PDF using only the extracted text provided in the prompt. "
"Answer questions about PDFs using only the extracted page text provided in the prompt. "
"Do not guess or use outside knowledge. "
"If the answer is not supported by the provided text, return not_found. "
"When answering, include a short list of evidence snippets copied from the provided text."
"When answering, include a short list of evidence snippets with their page numbers."
),
instructions=rag.instructions,
toolsets=[rag.toolset],
model_settings=runtime.smart_model_settings,
)
async def handle(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
if not request.extracted_text.strip():
return PdfQuestionNeedTextResponse(
reason="No extracted PDF text was provided, so the question cannot be answered yet."
if not self._has_page_text(request.page_text):
return PdfQuestionNeedContentResponse(
reason="No extracted PDF page text was provided, so the question cannot be answered yet.",
files=[
NeedContentFileRequest(
file_name=file_name,
content_types=[PdfContentType.PAGE_TEXT],
)
for file_name in request.file_names
],
max_pages=self.DEFAULT_MAX_PAGES,
max_characters=self.DEFAULT_MAX_CHARACTERS,
)
return await self._run_answer_agent(request)
@@ -45,5 +64,20 @@ class PdfQuestionAgent:
return result.output
def _build_prompt(self, request: PdfQuestionRequest) -> str:
file_name = request.file_name or "Unknown file"
return f"File: {file_name}\nQuestion: {request.question}\nExtracted text:\n{request.extracted_text}"
file_names = ", ".join(request.file_names) if request.file_names else "Unknown files"
sections = [
f"[File: {file_text.file_name}, Page {selection.page_number or '?'}]\n{selection.text}"
for file_text in request.page_text
for selection in file_text.pages
]
pages = "\n\n".join(sections)
history = format_conversation_history(request.conversation_history)
return (
f"Conversation history:\n{history}\n"
f"Files: {file_names}\n"
f"Question: {request.question}\n"
f"Extracted page text:\n{pages}"
)
def _has_page_text(self, page_text: list[ExtractedFileText]) -> bool:
return any(selection.text.strip() for file_text in page_text for selection in file_text.pages)
engine/src/stirling/agents/user_spec.py
+10 -10
View File
@@ -18,6 +18,7 @@ from stirling.contracts import (
EditClarificationRequest,
EditPlanResponse,
PdfEditRequest,
format_conversation_history,
)
from stirling.models import ApiModel
from stirling.services import AppRuntime
@@ -45,14 +46,15 @@ class UserSpecAgent:
)
async def draft(self, request: AgentDraftRequest) -> AgentDraftWorkflowResponse:
edit_plan = await self._build_edit_plan(request.user_message)
edit_plan = await self._build_edit_plan(request.user_message, request.conversation_history)
if not isinstance(edit_plan, EditPlanResponse):
return edit_plan
return AgentDraftResponse(draft=await self._run_draft_agent(request, edit_plan))
async def revise(self, request: AgentRevisionRequest) -> AgentRevisionWorkflowResponse:
edit_plan = await self._build_edit_plan(
f"Current objective: {request.current_draft.objective}\nRevision request: {request.user_message}"
f"Current objective: {request.current_draft.objective}\nRevision request: {request.user_message}",
request.conversation_history,
)
if not isinstance(edit_plan, EditPlanResponse):
return edit_plan
@@ -80,7 +82,7 @@ class UserSpecAgent:
def _build_draft_prompt(self, request: AgentDraftRequest, edit_plan: EditPlanResponse) -> str:
return (
f"User request:\n{request.user_message}\n\n"
f"Conversation history:\n{self._format_conversation_history(request.conversation_history)}\n\n"
f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n\n"
f"Edit plan summary:\n{edit_plan.summary}\n\n"
f"Edit plan rationale:\n{edit_plan.rationale or 'None'}\n\n"
f"Edit plan steps:\n{edit_plan.model_dump_json(indent=2)}"
@@ -89,20 +91,18 @@ class UserSpecAgent:
def _build_revision_prompt(self, request: AgentRevisionRequest, edit_plan: EditPlanResponse) -> str:
return (
f"Revision request:\n{request.user_message}\n\n"
f"Conversation history:\n{self._format_conversation_history(request.conversation_history)}\n\n"
f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n\n"
f"Current draft:\n{request.current_draft.model_dump_json(indent=2)}\n\n"
f"Edit plan summary:\n{edit_plan.summary}\n\n"
f"Edit plan rationale:\n{edit_plan.rationale or 'None'}\n\n"
f"Edit plan steps:\n{edit_plan.model_dump_json(indent=2)}"
)
def _format_conversation_history(self, conversation_history: list[ConversationMessage]) -> str:
if not conversation_history:
return "None"
return "\n".join(f"- {message.role}: {message.content}" for message in conversation_history)
async def _build_edit_plan(
self,
user_message: str,
conversation_history: list[ConversationMessage],
) -> EditPlanResponse | EditClarificationRequest | EditCannotDoResponse:
return await self.pdf_edit_agent.handle(PdfEditRequest(user_message=user_message))
return await self.pdf_edit_agent.handle(
PdfEditRequest(user_message=user_message, conversation_history=conversation_history)
)
engine/src/stirling/api/app.py
+16 -1
View File
@@ -4,6 +4,8 @@ from contextlib import asynccontextmanager
from typing import Annotated
from fastapi import Depends, FastAPI
from pydantic_ai import Agent
from pydantic_ai.models.instrumented import InstrumentationSettings
from stirling.agents import (
DocumentExtractorAgent,
@@ -15,17 +17,21 @@ from stirling.agents import (
PdfQuestionAgent,
UserSpecAgent,
)
from stirling.agents.ledger import MathAuditorAgent
from stirling.api.middleware import UserIdMiddleware
from stirling.api.routes import (
agent_draft_router,
execution_router,
form_fill_router,
ledger_router,
orchestrator_router,
pdf_edit_router,
pdf_question_router,
rag_router,
)
from stirling.config import AppSettings, load_settings
from stirling.contracts import HealthResponse
from stirling.services import build_runtime
from stirling.services import build_runtime, setup_posthog_tracking
def _load_startup_settings(fast_api: FastAPI) -> AppSettings:
@@ -46,18 +52,27 @@ async def lifespan(fast_api: FastAPI):
fast_api.state.pdf_question_agent = PdfQuestionAgent(runtime)
fast_api.state.user_spec_agent = UserSpecAgent(runtime)
fast_api.state.execution_planning_agent = ExecutionPlanningAgent(runtime)
fast_api.state.math_auditor_agent = MathAuditorAgent(runtime)
fast_api.state.form_analyser_agent = FormAnalyserAgent(runtime)
fast_api.state.form_filler_agent = FormFillerAgent(runtime)
fast_api.state.document_extractor_agent = DocumentExtractorAgent(runtime)
tracer_provider = setup_posthog_tracking(settings)
if tracer_provider:
Agent.instrument_all(InstrumentationSettings(tracer_provider=tracer_provider))
yield
if tracer_provider:
tracer_provider.shutdown()
app = FastAPI(title="Stirling AI Engine", lifespan=lifespan, version="0.1.0")
app.add_middleware(UserIdMiddleware)
app.include_router(orchestrator_router)
app.include_router(pdf_edit_router)
app.include_router(pdf_question_router)
app.include_router(agent_draft_router)
app.include_router(execution_router)
app.include_router(rag_router)
app.include_router(ledger_router)
app.include_router(form_fill_router)
engine/src/stirling/api/dependencies.py
+14
View File
@@ -12,6 +12,8 @@ from stirling.agents import (
PdfQuestionAgent,
UserSpecAgent,
)
from stirling.agents.ledger import MathAuditorAgent
from stirling.rag import RagService
from stirling.services import AppRuntime
@@ -39,6 +41,18 @@ def get_execution_planning_agent(request: Request) -> ExecutionPlanningAgent:
return request.app.state.execution_planning_agent
def get_rag_service(request: Request) -> RagService:
return request.app.state.runtime.rag_service
def get_rag_embedding_model(request: Request) -> str:
return request.app.state.runtime.settings.rag_embedding_model
def get_math_auditor_agent(request: Request) -> MathAuditorAgent:
return request.app.state.math_auditor_agent
def get_form_analyser_agent(request: Request) -> FormAnalyserAgent:
return request.app.state.form_analyser_agent
engine/src/stirling/api/middleware.py
+23
View File
@@ -0,0 +1,23 @@
from __future__ import annotations
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response
from stirling.services.tracking import current_user_id
_USER_ID_HEADER = "X-User-Id"
class UserIdMiddleware(BaseHTTPMiddleware):
"""Extract X-User-Id header and set it as the current user for PostHog tracking."""
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
user_id = request.headers.get(_USER_ID_HEADER)
if user_id:
token = current_user_id.set(user_id)
try:
return await call_next(request)
finally:
current_user_id.reset(token)
return await call_next(request)
engine/src/stirling/api/routes/__init__.py
+4
View File
@@ -1,15 +1,19 @@
from .agent_drafts import router as agent_draft_router
from .execution import router as execution_router
from .form_fill import router as form_fill_router
from .ledger import router as ledger_router
from .orchestrator import router as orchestrator_router
from .pdf_edit import router as pdf_edit_router
from .pdf_questions import router as pdf_question_router
from .rag import router as rag_router
__all__ = [
"agent_draft_router",
"execution_router",
"form_fill_router",
"ledger_router",
"orchestrator_router",
"pdf_edit_router",
"pdf_question_router",
"rag_router",
]
engine/src/stirling/api/routes/ledger.py
+60
View File
@@ -0,0 +1,60 @@
"""
Math Auditor Agent (mathAuditorAgent) — FastAPI routes.
Two internal endpoints, called only by the Java MathAuditorOrchestrator:
POST /api/v1/ai/math-auditor-agent/examine
Java sends a FolioManifest (cheap page classification).
Python returns a Requisition (what Java must extract).
POST /api/v1/ai/math-auditor-agent/deliberate
Java sends Evidence (fulfilled extraction results).
Python returns a Verdict directly.
"""
from __future__ import annotations
import logging
from decimal import Decimal, InvalidOperation
from typing import Annotated
from fastapi import APIRouter, Depends, HTTPException, Query
from stirling.agents.ledger import MathAuditorAgent
from stirling.api.dependencies import get_math_auditor_agent
from stirling.contracts.ledger import (
Evidence,
FolioManifest,
Requisition,
Verdict,
)
logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/v1/ai/math-auditor-agent", tags=["math-auditor-agent"])
@router.post("/examine", response_model=Requisition)
async def examine_endpoint(
manifest: FolioManifest,
agent: Annotated[MathAuditorAgent, Depends(get_math_auditor_agent)],
) -> Requisition:
"""Round 1: Java presents a FolioManifest; Python declares its Requisition."""
return await agent.examine(manifest)
@router.post("/deliberate", response_model=Verdict)
async def deliberate_endpoint(
evidence: Evidence,
agent: Annotated[MathAuditorAgent, Depends(get_math_auditor_agent)],
tolerance: str = Query(default="0.01"),
) -> Verdict:
"""Round 2: Java presents fulfilled Evidence; Python returns a Verdict."""
try:
tol = Decimal(tolerance)
if tol < 0:
raise HTTPException(status_code=400, detail="tolerance must be non-negative")
except InvalidOperation:
raise HTTPException(status_code=400, detail=f"Invalid tolerance value: {tolerance!r}")
return await agent.audit(evidence, tol)
engine/src/stirling/api/routes/rag.py
+78
View File
@@ -0,0 +1,78 @@
from __future__ import annotations
from typing import Annotated
from fastapi import APIRouter, Depends
from stirling.api.dependencies import get_rag_embedding_model, get_rag_service
from stirling.contracts import (
RagCollectionsResponse,
RagDeleteCollectionResponse,
RagIndexRequest,
RagIndexResponse,
RagSearchRequest,
RagSearchResponse,
RagSearchResultItem,
RagStatusResponse,
)
from stirling.rag import RagService
router = APIRouter(prefix="/api/v1/rag", tags=["rag"])
@router.get("/status", response_model=RagStatusResponse)
async def rag_status(
rag: Annotated[RagService, Depends(get_rag_service)],
embedding_model: Annotated[str, Depends(get_rag_embedding_model)],
) -> RagStatusResponse:
collections = await rag.list_collections()
return RagStatusResponse(embedding_model=embedding_model, collections=collections)
@router.post("/index", response_model=RagIndexResponse)
async def rag_index(
request: RagIndexRequest,
rag: Annotated[RagService, Depends(get_rag_service)],
) -> RagIndexResponse:
count = await rag.index_text(
collection=request.collection,
text=request.text,
source=request.source,
metadata=request.metadata,
)
return RagIndexResponse(collection=request.collection, chunks_indexed=count)
@router.post("/search", response_model=RagSearchResponse)
async def rag_search(
request: RagSearchRequest,
rag: Annotated[RagService, Depends(get_rag_service)],
) -> RagSearchResponse:
results = await rag.search(query=request.query, collection=request.collection, top_k=request.top_k)
items = [
RagSearchResultItem(
text=r.document.text,
source=r.document.metadata.get("source", ""),
chunk_id=r.document.metadata.get("chunk_index", ""),
score=r.score,
)
for r in results
]
return RagSearchResponse(query=request.query, results=items)
@router.get("/collections", response_model=RagCollectionsResponse)
async def rag_collections(
rag: Annotated[RagService, Depends(get_rag_service)],
) -> RagCollectionsResponse:
collections = await rag.list_collections()
return RagCollectionsResponse(collections=collections)
@router.delete("/collections/{name}", response_model=RagDeleteCollectionResponse)
async def rag_delete_collection(
name: str,
rag: Annotated[RagService, Depends(get_rag_service)],
) -> RagDeleteCollectionResponse:
await rag.delete_collection(name)
return RagDeleteCollectionResponse(status="deleted", collection=name)
engine/src/stirling/config/__init__.py
+3 -1
View File
@@ -1,8 +1,10 @@
"""Configuration models and loaders for the Stirling AI service."""
from .settings import AppSettings, load_settings
from .settings import ENGINE_ROOT, AppSettings, RagBackend, load_settings
__all__ = [
"ENGINE_ROOT",
"AppSettings",
"RagBackend",
"load_settings",
]
engine/src/stirling/config/settings.py
+57 -2
View File
@@ -1,5 +1,8 @@
from __future__ import annotations
import logging
import logging.handlers
from enum import StrEnum
from functools import lru_cache
from pathlib import Path
@@ -9,18 +12,70 @@ from pydantic_settings import BaseSettings, SettingsConfigDict
ENGINE_ROOT = Path(__file__).resolve().parents[3]
ENV_FILE = ENGINE_ROOT / ".env"
ENV_LOCAL_FILE = ENGINE_ROOT / ".env.local"
class RagBackend(StrEnum):
SQLITE = "sqlite"
PGVECTOR = "pgvector"
class AppSettings(BaseSettings):
model_config = SettingsConfigDict(env_file=ENV_FILE, extra="ignore", populate_by_name=True)
model_config = SettingsConfigDict(env_file=(ENV_FILE, ENV_LOCAL_FILE), extra="ignore", populate_by_name=True)
smart_model_name: str = Field(validation_alias="STIRLING_SMART_MODEL")
fast_model_name: str = Field(validation_alias="STIRLING_FAST_MODEL")
smart_model_max_tokens: int = Field(validation_alias="STIRLING_SMART_MODEL_MAX_TOKENS")
fast_model_max_tokens: int = Field(validation_alias="STIRLING_FAST_MODEL_MAX_TOKENS")
# RAG settings — always on; the backend picks between embedded sqlite-vec and external pgvector.
rag_backend: RagBackend = Field(validation_alias="STIRLING_RAG_BACKEND")
rag_embedding_model: str = Field(validation_alias="STIRLING_RAG_EMBEDDING_MODEL")
rag_store_path: Path = Field(validation_alias="STIRLING_RAG_STORE_PATH")
rag_pgvector_dsn: str = Field(validation_alias="STIRLING_RAG_PGVECTOR_DSN")
rag_chunk_size: int = Field(validation_alias="STIRLING_RAG_CHUNK_SIZE")
rag_chunk_overlap: int = Field(validation_alias="STIRLING_RAG_CHUNK_OVERLAP")
rag_default_top_k: int = Field(validation_alias="STIRLING_RAG_TOP_K")
log_level: str = Field(default="INFO", validation_alias="STIRLING_LOG_LEVEL")
log_file: str = Field(default="", validation_alias="STIRLING_LOG_FILE")
posthog_enabled: bool = Field(validation_alias="STIRLING_POSTHOG_ENABLED")
posthog_api_key: str = Field(validation_alias="STIRLING_POSTHOG_API_KEY")
posthog_host: str = Field(validation_alias="STIRLING_POSTHOG_HOST")
def _configure_logging(level_name: str, log_file: str) -> None:
"""Configure the ``stirling`` logger hierarchy."""
level = logging.getLevelNamesMapping().get(level_name.upper())
if level is None:
logging.getLogger("stirling").warning(
"Unknown STIRLING_LOG_LEVEL %r, defaulting to INFO",
level_name,
)
level = logging.INFO
root = logging.getLogger("stirling")
root.setLevel(level)
if log_file:
log_path = Path(log_file)
log_path.parent.mkdir(parents=True, exist_ok=True)
fh = logging.handlers.TimedRotatingFileHandler(
log_path,
when="midnight",
backupCount=1,
encoding="utf-8",
)
fh.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(name)s [%(funcName)s] %(message)s"))
fh.setLevel(level)
root.addHandler(fh)
@lru_cache(maxsize=1)
def load_settings() -> AppSettings:
load_dotenv(ENV_FILE)
return AppSettings.model_validate({})
load_dotenv(ENV_LOCAL_FILE, override=True)
settings = AppSettings.model_validate({})
_configure_logging(settings.log_level, settings.log_file)
return settings
engine/src/stirling/contracts/__init__.py
+85 -18
View File
@@ -8,7 +8,18 @@ from .agent_drafts import (
AgentRevisionWorkflowResponse,
)
from .agent_specs import AgentSpec, AgentSpecStep, AiToolAgentStep
from .common import ConversationMessage, PdfTextSelection, ToolOperationStep
from .common import (
ArtifactKind,
ConversationMessage,
ExtractedFileText,
PdfContentType,
PdfTextSelection,
StepKind,
SupportedCapability,
ToolOperationStep,
WorkflowOutcome,
format_conversation_history,
)
from .execution import (
AgentExecutionRequest,
CannotContinueExecutionAction,
@@ -41,7 +52,24 @@ from .form_fill import (
ProposedProfile,
)
from .health import HealthResponse
from .orchestrator import OrchestratorRequest, OrchestratorResponse, SupportedCapability, UnsupportedCapabilityResponse
from .ledger import (
Discrepancy,
DiscrepancyKind,
Evidence,
Folio,
FolioManifest,
FolioType,
Requisition,
Severity,
Verdict,
)
from .orchestrator import (
ExtractedTextArtifact,
OrchestratorRequest,
OrchestratorResponse,
UnsupportedCapabilityResponse,
WorkflowArtifact,
)
from .pdf_edit import (
EditCannotDoResponse,
EditClarificationRequest,
@@ -50,14 +78,27 @@ from .pdf_edit import (
PdfEditResponse,
)
from .pdf_questions import (
NeedContentFileRequest,
PdfQuestionAnswerResponse,
PdfQuestionNeedTextResponse,
PdfQuestionNeedContentResponse,
PdfQuestionNotFoundResponse,
PdfQuestionRequest,
PdfQuestionResponse,
)
from .rag import (
MAX_INDEX_TEXT_LENGTH,
RagCollectionsResponse,
RagDeleteCollectionResponse,
RagIndexRequest,
RagIndexResponse,
RagSearchRequest,
RagSearchResponse,
RagSearchResultItem,
RagStatusResponse,
)
__all__ = [
"MAX_INDEX_TEXT_LENGTH",
"AgentDraft",
"AgentDraftRequest",
"AgentDraftResponse",
@@ -69,48 +110,74 @@ __all__ = [
"AgentSpec",
"AgentSpecStep",
"AiToolAgentStep",
"AnalysedFileResult",
"ArtifactKind",
"CannotContinueExecutionAction",
"CleanedLabel",
"CompletedExecutionAction",
"ConversationMessage",
"CrossFileRole",
"DetectedRole",
"Discrepancy",
"DiscrepancyKind",
"DocumentExtractionRequest",
"DocumentExtractionResponse",
"DocumentText",
"FileFieldSet",
"FileFillRequest",
"FileFillResult",
"FormAnalysisRequest",
"FormAnalysisResponse",
"AnalysedFileResult",
"CrossFileRole",
"FormFillBatchRequest",
"FormFillBatchResponse",
"CompletedExecutionAction",
"FieldMapping",
"FormField",
"EditCannotDoResponse",
"EditClarificationRequest",
"EditPlanResponse",
"Evidence",
"ExecutionContext",
"ExecutionStepResult",
"ExtractedFileText",
"ExtractedTextArtifact",
"FieldMapping",
"FileFieldSet",
"FileFillRequest",
"FileFillResult",
"Folio",
"FolioManifest",
"FolioType",
"FormAnalysisRequest",
"FormAnalysisResponse",
"FormField",
"FormFillBatchRequest",
"FormFillBatchResponse",
"format_conversation_history",
"HealthResponse",
"KnowledgeEntry",
"KnowledgeUpdateResponse",
"MultiProfileExtractionResponse",
"NeedContentFileRequest",
"ProposedProfile",
"NextExecutionAction",
"OrchestratorRequest",
"OrchestratorResponse",
"ProposedProfile",
"PdfContentType",
"PdfEditRequest",
"PdfEditResponse",
"PdfQuestionAnswerResponse",
"PdfQuestionNeedContentResponse",
"PdfQuestionNotFoundResponse",
"PdfQuestionNeedTextResponse",
"PdfQuestionRequest",
"PdfQuestionResponse",
"PdfTextSelection",
"RagCollectionsResponse",
"RagDeleteCollectionResponse",
"RagIndexRequest",
"RagIndexResponse",
"RagSearchRequest",
"RagSearchResponse",
"RagSearchResultItem",
"RagStatusResponse",
"Requisition",
"Severity",
"StepKind",
"SupportedCapability",
"ToolOperationStep",
"ToolCallExecutionAction",
"ToolOperationStep",
"UnsupportedCapabilityResponse",
"Verdict",
"WorkflowArtifact",
"WorkflowOutcome",
]
engine/src/stirling/contracts/agent_drafts.py
+4 -4
View File
@@ -7,12 +7,12 @@ from pydantic import Field
from stirling.models import ApiModel
from .agent_specs import AgentSpecStep
from .common import ConversationMessage
from .common import ConversationMessage, StepKind, WorkflowOutcome
from .pdf_edit import EditCannotDoResponse, EditClarificationRequest
class AgentDraftStep(ApiModel):
kind: Literal["tool", "ai_tool"]
kind: Literal[StepKind.TOOL, StepKind.AI_TOOL]
title: str
description: str
@@ -30,7 +30,7 @@ class AgentDraftRequest(ApiModel):
class AgentDraftResponse(ApiModel):
outcome: Literal["draft"] = "draft"
outcome: Literal[WorkflowOutcome.DRAFT] = WorkflowOutcome.DRAFT
draft: AgentDraft
@@ -41,7 +41,7 @@ class AgentRevisionRequest(ApiModel):
class AgentRevisionResponse(ApiModel):
outcome: Literal["draft"] = "draft"
outcome: Literal[WorkflowOutcome.DRAFT] = WorkflowOutcome.DRAFT
draft: AgentDraft
engine/src/stirling/contracts/agent_specs.py
+4 -4
View File
@@ -4,16 +4,16 @@ from typing import Annotated, Literal
from pydantic import Field
from stirling.models import ApiModel, OperationId
from stirling.models import ApiModel, ToolEndpoint
from .common import ToolOperationStep
from .common import StepKind, ToolOperationStep
class AiToolAgentStep(ApiModel):
kind: Literal["ai_tool"] = "ai_tool"
kind: Literal[StepKind.AI_TOOL] = StepKind.AI_TOOL
title: str
description: str
tool: OperationId
tool: ToolEndpoint
instruction: str
engine/src/stirling/contracts/common.py
+104 -9
View File
@@ -1,10 +1,89 @@
from __future__ import annotations
from typing import Literal
from enum import StrEnum
from typing import Literal, assert_never
from pydantic import model_validator
from pydantic import Field, model_validator
from stirling.models import OPERATIONS, ApiModel, OperationId, ParamToolModel
from stirling.models import OPERATIONS, ApiModel, ToolEndpoint
from stirling.models.agent_tool_models import AGENT_OPERATIONS, AgentToolId, AnyParamModel, AnyToolId
class PdfContentType(StrEnum):
"""Types of content that can be extracted from a PDF and sent to the AI.
Java counterpart: AiPdfContentType.java - values must stay in sync.
"""
# Document-level structured data
PAGE_LAYOUT = "page_layout"
DOCUMENT_METADATA = "document_metadata"
ENCRYPTION_INFO = "encryption_info"
BOOKMARKS = "bookmarks"
LAYERS = "layers"
EMBEDDED_FILES = "embedded_files"
JAVASCRIPT = "javascript"
LINKS = "links"
IMAGE_INFO = "image_info"
FONTS = "fonts"
# Text and content
PAGE_TEXT = "page_text"
FULL_TEXT = "full_text"
FORM_FIELDS = "form_fields"
ANNOTATIONS = "annotations"
SIGNATURES = "signatures"
STRUCTURE_TREE = "structure_tree"
XMP_METADATA = "xmp_metadata"
# Heavy content
COMPLIANCE = "compliance"
IMAGES = "images"
class WorkflowOutcome(StrEnum):
"""Discriminator values for all workflow response unions (outcome field).
Java counterpart: AiWorkflowOutcome.java - values must stay in sync.
"""
ANSWER = "answer"
NEED_CONTENT = "need_content"
NOT_FOUND = "not_found"
PLAN = "plan"
NEED_CLARIFICATION = "need_clarification"
CANNOT_DO = "cannot_do"
DRAFT = "draft"
TOOL_CALL = "tool_call"
COMPLETED = "completed"
CANNOT_CONTINUE = "cannot_continue"
UNSUPPORTED_CAPABILITY = "unsupported_capability"
class ArtifactKind(StrEnum):
"""Discriminator values for WorkflowArtifact unions (kind field).
Java counterpart: PdfContentExtractor.ArtifactKind - values must stay in sync.
"""
EXTRACTED_TEXT = "extracted_text"
class StepKind(StrEnum):
"""Discriminator values for AgentSpecStep unions (kind field)."""
TOOL = "tool"
AI_TOOL = "ai_tool"
class SupportedCapability(StrEnum):
ORCHESTRATE = "orchestrate"
PDF_EDIT = "pdf_edit"
PDF_QUESTION = "pdf_question"
AGENT_DRAFT = "agent_draft"
AGENT_REVISE = "agent_revise"
AGENT_NEXT_ACTION = "agent_next_action"
MATH_AUDITOR_AGENT = "math_auditor_agent"
class ConversationMessage(ApiModel):
@@ -12,21 +91,37 @@ class ConversationMessage(ApiModel):
content: str
def format_conversation_history(conversation_history: list[ConversationMessage]) -> str:
if not conversation_history:
return "None"
return "\n".join(f"- {message.role}: {message.content}" for message in conversation_history)
class PdfTextSelection(ApiModel):
page_number: int | None = None
text: str
class ExtractedFileText(ApiModel):
file_name: str
pages: list[PdfTextSelection] = Field(default_factory=list)
class ToolOperationStep(ApiModel):
kind: Literal["tool"] = "tool"
tool: OperationId
parameters: ParamToolModel
kind: Literal[StepKind.TOOL] = StepKind.TOOL
tool: AnyToolId
parameters: AnyParamModel
@model_validator(mode="after")
def validate_tool_parameter_pairing(self) -> ToolOperationStep:
expected_type = OPERATIONS[self.tool]
if isinstance(self.tool, AgentToolId):
expected_type = AGENT_OPERATIONS[self.tool]
elif isinstance(self.tool, ToolEndpoint):
expected_type = OPERATIONS[self.tool]
else:
assert_never(self.tool)
if not isinstance(self.parameters, expected_type):
actual_type = type(self.parameters).__name__
expected_type_name = expected_type.__name__
raise ValueError(f"Parameters for tool {self.tool.value} must be {expected_type_name}, got {actual_type}.")
raise ValueError(f"Parameters for tool {self.tool} must be {expected_type.__name__}, got {actual_type}.")
return self
engine/src/stirling/contracts/execution.py
+7 -6
View File
@@ -4,14 +4,15 @@ from typing import Annotated, Any, Literal
from pydantic import Field
from stirling.models import ApiModel, OperationId, ParamToolModel
from stirling.models import ApiModel, ParamToolModel, ToolEndpoint
from .agent_specs import AgentSpec
from .common import WorkflowOutcome
class ExecutionStepResult(ApiModel):
step_index: int
tool: OperationId | None = None
tool: ToolEndpoint | None = None
success: bool
output_summary: str | None = None
output_data: dict[str, Any] = Field(default_factory=dict)
@@ -31,19 +32,19 @@ class AgentExecutionRequest(ApiModel):
class ToolCallExecutionAction(ApiModel):
outcome: Literal["tool_call"] = "tool_call"
tool: OperationId
outcome: Literal[WorkflowOutcome.TOOL_CALL] = WorkflowOutcome.TOOL_CALL
tool: ToolEndpoint
parameters: ParamToolModel
rationale: str | None = None
class CompletedExecutionAction(ApiModel):
outcome: Literal["completed"] = "completed"
outcome: Literal[WorkflowOutcome.COMPLETED] = WorkflowOutcome.COMPLETED
summary: str
class CannotContinueExecutionAction(ApiModel):
outcome: Literal["cannot_continue"] = "cannot_continue"
outcome: Literal[WorkflowOutcome.CANNOT_CONTINUE] = WorkflowOutcome.CANNOT_CONTINUE
reason: str
engine/src/stirling/contracts/ledger.py
+196
View File
@@ -0,0 +1,196 @@
"""
Ledger Auditor — shared models for the Java-Python protocol.
Every struct that crosses the wire lives here so the contract is
impossible to miss or partially implement.
"""
from __future__ import annotations
from enum import StrEnum
from typing import Literal
from pydantic import Field
from stirling.models import ApiModel
# ---------------------------------------------------------------------------
# Page classification — Java's side of the conversation
# ---------------------------------------------------------------------------
class FolioType(StrEnum):
"""How Java classifies each page after a cheap PDFBox scan.
Java counterpart: FolioType.java - values must stay in sync.
"""
TEXT = "text" # selectable text layer present
IMAGE = "image" # image-only, will need OCR
MIXED = "mixed" # partial text layer + embedded images
class FolioManifest(ApiModel):
"""
Java's opening move: a fast, cheap page classification with no OCR or
table extraction — just PDFBox character counts and image detection.
Python inspects this and returns a Requisition declaring what it needs.
"""
session_id: str = Field(description="Opaque handle Java uses to find the PDF on disk.")
page_count: int = Field(ge=1)
folio_types: list[FolioType] = Field(description="One entry per page (0-indexed). len(folio_types) == page_count.")
round: int = Field(default=1, ge=1, le=3, description="Which negotiation round this is.")
# ---------------------------------------------------------------------------
# Requisition — Python's declaration of what it needs
# ---------------------------------------------------------------------------
class Requisition(ApiModel):
"""
Python's reply to a FolioManifest: a precise shopping list of what Java
must extract before the auditor can form an opinion.
Java fulfils this and sends back an Evidence payload.
"""
type: Literal["requisition"] = "requisition"
need_text: list[int] = Field(
default_factory=list,
description="0-indexed page numbers. Java runs PDFBox text extraction on these.",
)
need_tables: list[int] = Field(
default_factory=list,
description="0-indexed page numbers. Java runs Tabula CSV extraction on these.",
)
need_ocr: list[int] = Field(
default_factory=list,
description="0-indexed page numbers. Java runs OCRmyPDF on these.",
)
rationale: str = Field(description="Plain-language reason, written for log readability, not the client.")
# ---------------------------------------------------------------------------
# Evidence — Java's fulfilment of a Requisition
# ---------------------------------------------------------------------------
class Folio(ApiModel):
"""
One page's worth of extracted content — whatever Java was able to provide
in response to the Requisition for that page.
"""
page: int = Field(ge=0, description="0-indexed page number.")
text: str | None = Field(default=None, description="PDFBox plain-text extraction.")
tables: list[str] | None = Field(default=None, description="Tabula CSV strings, one per table found on the page.")
ocr_text: str | None = Field(default=None, description="OCRmyPDF output text.")
ocr_confidence: float | None = Field(
default=None, ge=0.0, le=1.0, description="Mean character confidence from OCRmyPDF."
)
@property
def readable_text(self) -> str:
"""Best available text for this folio — OCR wins over digital when present."""
return self.ocr_text or self.text or ""
class Evidence(ApiModel):
"""
Java's fulfilment package: the extracted content Python asked for.
Java may also set final_round=True on the last allowable round to signal
that the auditor must return a Verdict regardless of remaining questions.
"""
session_id: str
folios: list[Folio]
round: int = Field(ge=1, le=3)
final_round: bool = Field(
default=False,
description="When True, Java will not honour further Requisitions. "
"The auditor must return a Verdict this round.",
)
unauditable_pages: list[int] = Field(
default_factory=list,
description=(
"Pages that were requested in the Requisition but could not be fulfilled — "
"e.g. OCR was asked for but is not wired. The Auditor echoes these into "
"Verdict.unauditable_pages so the client knows coverage is incomplete."
),
)
# ---------------------------------------------------------------------------
# Findings — what the auditor discovers
# ---------------------------------------------------------------------------
class DiscrepancyKind(StrEnum):
"""Java counterpart: DiscrepancyKind.java - values must stay in sync."""
TALLY = "tally" # a row/column sum is wrong
ARITHMETIC = "arithmetic" # an inline calculation is wrong
CONSISTENCY = "consistency" # the same figure is stated differently elsewhere
STATEMENT = "statement" # a prose claim contradicts the numbers
class Severity(StrEnum):
"""Java counterpart: AuditSeverity.java - values must stay in sync."""
ERROR = "error" # definite arithmetic mistake
WARNING = "warning" # possible rounding or ambiguity
class Discrepancy(ApiModel):
"""A single mathematical error found in the document."""
page: int = Field(ge=0)
kind: DiscrepancyKind
severity: Severity
description: str = Field(description="Human-readable explanation of the error.")
stated: str = Field(description="The value as it appears in the document.")
expected: str = Field(description="The value the auditor calculated.")
context: str = Field(
default="",
description="Surrounding text or table fragment for traceability.",
)
# ---------------------------------------------------------------------------
# Verdict — the final report
# ---------------------------------------------------------------------------
class Verdict(ApiModel):
"""
The auditor's final opinion on the document's mathematical integrity.
Returned to Java as the terminal message in the negotiation.
"""
type: Literal["verdict"] = "verdict"
session_id: str
discrepancies: list[Discrepancy] = Field(default_factory=list)
pages_examined: list[int] = Field(description="0-indexed page numbers the auditor actually inspected.")
rounds_taken: int = Field(ge=1, le=3)
summary: str = Field(description="One or two sentences summarising the audit outcome for the client.")
clean: bool = Field(description="True iff no errors were found (warnings are tolerated).")
unauditable_pages: list[int] = Field(
default_factory=list,
description=(
"0-indexed pages that could not be audited — typically because OCR was "
"requested but is not yet wired. Java populates this by omitting the folio "
"and the Auditor echoes the page number here so the client knows coverage "
"is incomplete."
),
)
@property
def error_count(self) -> int:
return sum(1 for d in self.discrepancies if d.severity == Severity.ERROR)
@property
def warning_count(self) -> int:
return sum(1 for d in self.discrepancies if d.severity == Severity.WARNING)
engine/src/stirling/contracts/orchestrator.py
+18 -11
View File
@@ -1,6 +1,5 @@
from __future__ import annotations
from enum import StrEnum
from typing import Annotated, Literal
from pydantic import Field
@@ -8,29 +7,37 @@ from pydantic import Field
from stirling.models import ApiModel
from .agent_drafts import AgentDraftResponse
from .common import (
ArtifactKind,
ConversationMessage,
ExtractedFileText,
SupportedCapability,
WorkflowOutcome,
)
from .execution import NextExecutionAction
from .form_fill import KnowledgeUpdateResponse
from .pdf_edit import PdfEditResponse
from .pdf_questions import PdfQuestionResponse
class SupportedCapability(StrEnum):
ORCHESTRATE = "orchestrate"
PDF_EDIT = "pdf_edit"
PDF_QUESTION = "pdf_question"
AGENT_DRAFT = "agent_draft"
AGENT_REVISE = "agent_revise"
AGENT_NEXT_ACTION = "agent_next_action"
FORM_FILL = "form_fill"
class ExtractedTextArtifact(ApiModel):
kind: Literal[ArtifactKind.EXTRACTED_TEXT] = ArtifactKind.EXTRACTED_TEXT
files: list[ExtractedFileText] = Field(default_factory=list)
WorkflowArtifact = Annotated[ExtractedTextArtifact, Field(discriminator="kind")]
class OrchestratorRequest(ApiModel):
user_message: str
conversation_id: str | None = None
file_names: list[str]
conversation_history: list[ConversationMessage] = Field(default_factory=list)
artifacts: list[WorkflowArtifact] = Field(default_factory=list)
resume_with: SupportedCapability | None = None
class UnsupportedCapabilityResponse(ApiModel):
outcome: Literal["unsupported_capability"] = "unsupported_capability"
outcome: Literal[WorkflowOutcome.UNSUPPORTED_CAPABILITY] = WorkflowOutcome.UNSUPPORTED_CAPABILITY
capability: str
message: str
engine/src/stirling/contracts/pdf_edit.py
+5 -5
View File
@@ -6,30 +6,30 @@ from pydantic import Field
from stirling.models import ApiModel
from .common import ToolOperationStep
from .common import ConversationMessage, ToolOperationStep, WorkflowOutcome
class PdfEditRequest(ApiModel):
user_message: str
conversation_id: str | None = None
file_names: list[str] = Field(default_factory=list)
conversation_history: list[ConversationMessage] = Field(default_factory=list)
class EditPlanResponse(ApiModel):
outcome: Literal["plan"] = "plan"
outcome: Literal[WorkflowOutcome.PLAN] = WorkflowOutcome.PLAN
summary: str
rationale: str | None = None
steps: list[ToolOperationStep]
class EditClarificationRequest(ApiModel):
outcome: Literal["need_clarification"] = "need_clarification"
outcome: Literal[WorkflowOutcome.NEED_CLARIFICATION] = WorkflowOutcome.NEED_CLARIFICATION
question: str
reason: str
class EditCannotDoResponse(ApiModel):
outcome: Literal["cannot_do"] = "cannot_do"
outcome: Literal[WorkflowOutcome.CANNOT_DO] = WorkflowOutcome.CANNOT_DO
reason: str
engine/src/stirling/contracts/pdf_questions.py
+27 -9
View File
@@ -6,31 +6,49 @@ from pydantic import Field
from stirling.models import ApiModel
from .common import (
ConversationMessage,
ExtractedFileText,
PdfContentType,
SupportedCapability,
WorkflowOutcome,
)
class PdfQuestionRequest(ApiModel):
question: str
conversation_id: str | None = None
extracted_text: str = ""
file_name: str | None = None
page_text: list[ExtractedFileText] = Field(default_factory=list)
file_names: list[str]
conversation_history: list[ConversationMessage] = Field(default_factory=list)
class PdfQuestionAnswerResponse(ApiModel):
outcome: Literal["answer"] = "answer"
outcome: Literal[WorkflowOutcome.ANSWER] = WorkflowOutcome.ANSWER
answer: str
evidence: list[str] = Field(default_factory=list)
evidence: list[ExtractedFileText] = Field(default_factory=list)
class PdfQuestionNeedTextResponse(ApiModel):
outcome: Literal["need_text"] = "need_text"
class NeedContentFileRequest(ApiModel):
file_name: str
page_numbers: list[int] = Field(default_factory=list)
content_types: list[PdfContentType]
class PdfQuestionNeedContentResponse(ApiModel):
outcome: Literal[WorkflowOutcome.NEED_CONTENT] = WorkflowOutcome.NEED_CONTENT
resume_with: SupportedCapability = SupportedCapability.PDF_QUESTION
reason: str
files: list[NeedContentFileRequest] = Field(default_factory=list)
max_pages: int
max_characters: int
class PdfQuestionNotFoundResponse(ApiModel):
outcome: Literal["not_found"] = "not_found"
outcome: Literal[WorkflowOutcome.NOT_FOUND] = WorkflowOutcome.NOT_FOUND
reason: str
PdfQuestionResponse = Annotated[
PdfQuestionAnswerResponse | PdfQuestionNeedTextResponse | PdfQuestionNotFoundResponse,
PdfQuestionAnswerResponse | PdfQuestionNeedContentResponse | PdfQuestionNotFoundResponse,
Field(discriminator="outcome"),
]
engine/src/stirling/contracts/rag.py
+51
View File
@@ -0,0 +1,51 @@
from __future__ import annotations
from pydantic import Field
from stirling.models import ApiModel
MAX_INDEX_TEXT_LENGTH = 1_000_000 # 1MB text limit per index request
class RagStatusResponse(ApiModel):
embedding_model: str
collections: list[str]
class RagIndexRequest(ApiModel):
collection: str = Field(min_length=1)
text: str = Field(max_length=MAX_INDEX_TEXT_LENGTH)
source: str = ""
metadata: dict[str, str] = Field(default_factory=dict)
class RagIndexResponse(ApiModel):
collection: str
chunks_indexed: int
class RagSearchRequest(ApiModel):
query: str
collection: str | None = Field(default=None, min_length=1)
top_k: int = 5
class RagSearchResultItem(ApiModel):
text: str
source: str
chunk_id: str
score: float
class RagSearchResponse(ApiModel):
query: str
results: list[RagSearchResultItem]
class RagCollectionsResponse(ApiModel):
collections: list[str]
class RagDeleteCollectionResponse(ApiModel):
status: str
collection: str
engine/src/stirling/logging.py
+22
View File
@@ -0,0 +1,22 @@
"""Shared logging utilities for the Stirling AI engine."""
from __future__ import annotations
import json
class Pretty:
"""Lazy JSON formatter — only serialises when ``str()`` is called.
Designed for use with ``logging``'s ``%s`` formatting so that the
JSON serialisation is skipped entirely when the log message is
never emitted.
"""
__slots__ = ("_obj",)
def __init__(self, obj: object) -> None:
self._obj = obj
def __str__(self) -> str:
return json.dumps(self._obj, indent=2, default=str, ensure_ascii=True)
engine/src/stirling/models/__init__.py
+2 -2
View File
@@ -1,11 +1,11 @@
from . import tool_models
from .base import ApiModel
from .tool_models import OPERATIONS, OperationId, ParamToolModel
from .tool_models import OPERATIONS, ParamToolModel, ToolEndpoint
__all__ = [
"ApiModel",
"OPERATIONS",
"OperationId",
"ParamToolModel",
"ToolEndpoint",
"tool_models",
]
engine/src/stirling/models/agent_tool_models.py
+30
View File
@@ -0,0 +1,30 @@
"""Agent tool IDs, parameter models, and registry.
tool_models.py is auto-generated from the Java OpenAPI spec. This file is its
manually-maintained counterpart for tools backed by AI agent pipelines.
"""
from __future__ import annotations
from enum import StrEnum
from stirling.models.base import ApiModel
from stirling.models.tool_models import ParamToolModel, ToolEndpoint
class AgentToolId(StrEnum):
MATH_AUDITOR_AGENT = "mathAuditorAgent"
class MathAuditorAgentParams(ApiModel):
tolerance: str = "0.01"
type AgentParamModel = MathAuditorAgentParams
type AnyToolId = ToolEndpoint | AgentToolId
type AnyParamModel = ParamToolModel | AgentParamModel
AGENT_OPERATIONS: dict[AgentToolId, type[AgentParamModel]] = {
AgentToolId.MATH_AUDITOR_AGENT: MathAuditorAgentParams,
}
engine/src/stirling/models/tool_models.py
+1292 -351
View File
File diff suppressed because it is too large Load Diff
engine/src/stirling/rag/README.md
+88
View File
@@ -0,0 +1,88 @@
# RAG Integration Guide
## Adding RAG to an Agent
```python
from pydantic_ai import Agent
from stirling.services import AppRuntime
class MyAgent:
def __init__(self, runtime: AppRuntime) -> None:
rag = runtime.rag_capability
self.agent = Agent(
model=runtime.smart_model,
system_prompt="Your prompt here...",
instructions=rag.instructions,
toolsets=[rag.toolset],
)
```
That's it. The agent gets a `search_knowledge` tool it can call autonomously.
## Scoping to Specific Collections
Collections are named buckets of indexed documents — think folders. By default an agent searches everything in the store. Pass `collections=` to restrict it to only the docs indexed under those names.
```python
from stirling.rag import RagCapability
# Only searches docs indexed under "company-docs" — ignores everything else
scoped = RagCapability(runtime.rag_service, collections=["company-docs"], top_k=3)
# Searches multiple collections
multi = RagCapability(runtime.rag_service, collections=["company-docs", "product-specs"])
# No collections arg = searches all collections in the store
everything = RagCapability(runtime.rag_service)
```
## Config
Non-secret defaults live in the committed `engine/.env`:
```
STIRLING_RAG_BACKEND=sqlite # or "pgvector"
STIRLING_RAG_EMBEDDING_MODEL=voyageai:voyage-4
STIRLING_RAG_STORE_PATH=data/rag.db # used when backend=sqlite
STIRLING_RAG_PGVECTOR_DSN= # used when backend=pgvector
STIRLING_RAG_CHUNK_SIZE=512
STIRLING_RAG_CHUNK_OVERLAP=64
STIRLING_RAG_TOP_K=5
```
Provider credentials (and any local overrides) go in the uncommitted `engine/.env.local`:
```
VOYAGE_API_KEY=your-key
```
## Backends
**`sqlite`** — Embedded sqlite-vec. Single `.db` file, zero ops. Ideal for dev and self-hosted deployments.
**`pgvector`** — External PostgreSQL with the `vector` extension. Point `STIRLING_RAG_PGVECTOR_DSN` at your Postgres instance.
Both backends implement the same `VectorStore` interface, so agents and the RAG service work identically regardless of which you pick.
For a self-hosted embedding server (e.g. Ollama, TEI, vLLM) set the model string accordingly and point at the server via its native env var:
```
# Ollama running on another machine
STIRLING_RAG_EMBEDDING_MODEL=ollama:nomic-embed-text
OLLAMA_HOST=http://192.168.1.50:11434
# Any OpenAI-compatible embedding server
STIRLING_RAG_EMBEDDING_MODEL=openai:my-model
OPENAI_BASE_URL=http://192.168.1.50:8080/v1
```
## API Endpoints
| Method | Endpoint | Purpose |
|--------|----------|---------|
| GET | `/api/v1/rag/status` | Report embedding model and existing collections |
| POST | `/api/v1/rag/index` | Index text into a collection |
| POST | `/api/v1/rag/search` | Search a collection |
| GET | `/api/v1/rag/collections` | List collections |
| DELETE | `/api/v1/rag/collections/{name}` | Delete a collection |
engine/src/stirling/rag/__init__.py
+19
View File
@@ -0,0 +1,19 @@
from __future__ import annotations
from stirling.rag.capability import RagCapability
from stirling.rag.embedder import EmbeddingService
from stirling.rag.pgvector_store import PgVectorStore
from stirling.rag.service import RagService
from stirling.rag.sqlite_vec_store import SqliteVecStore
from stirling.rag.store import Document, SearchResult, VectorStore
__all__ = [
"Document",
"EmbeddingService",
"PgVectorStore",
"RagCapability",
"RagService",
"SearchResult",
"SqliteVecStore",
"VectorStore",
]
engine/src/stirling/rag/capability.py
+108
View File
@@ -0,0 +1,108 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable
from pydantic_ai import FunctionToolset
from pydantic_ai.toolsets import AbstractToolset
from stirling.rag.service import RagService
class RagCapability:
"""Bundles RAG instructions and the ``search_knowledge`` toolset for agent injection.
Agents consume this as::
rag = runtime.rag_capability
Agent(
...,
instructions=rag.instructions,
toolsets=[rag.toolset],
)
When no collections are pinned, the instructions are generated dynamically at
run time so the agent sees the current list of collections in the store.
"""
def __init__(
self,
rag_service: RagService,
collections: list[str] | None = None,
top_k: int = 5,
) -> None:
self._rag_service = rag_service
self._collections = collections
self._top_k = top_k
toolset: FunctionToolset[None] = FunctionToolset()
toolset.add_function(self._search_knowledge, name="search_knowledge")
self._toolset = toolset
@property
def instructions(self) -> str | Callable[[], Awaitable[str]]:
if self._collections:
return self._static_instructions_text(self._collections)
return self._dynamic_instructions
@property
def toolset(self) -> AbstractToolset[None]:
return self._toolset
@staticmethod
def _static_instructions_text(collections: list[str]) -> str:
collection_desc = f"collections: {', '.join(collections)}"
return (
"You have access to a knowledge base search tool called 'search_knowledge'. "
f"It searches {collection_desc} for relevant information. "
"Use it when the provided context is insufficient to answer the user's question, "
"or when you think additional background information would improve your answer. "
"You do not have to use it if the answer is already clear from the provided text."
)
async def _dynamic_instructions(self) -> str:
collections = await self._rag_service.list_collections()
if collections:
names = ", ".join(collections)
collection_desc = f"the following knowledge base collections: {names}"
else:
collection_desc = "the knowledge base (currently empty — no collections indexed yet)"
return (
"You have access to a knowledge base search tool called 'search_knowledge'. "
f"It searches {collection_desc} for relevant information. "
"Use it when the provided context is insufficient to answer the user's question, "
"or when you think additional background information would improve your answer. "
"You do not have to use it if the answer is already clear from the provided text."
)
async def _search_knowledge(self, query: str, max_results: int | None = None) -> str:
"""Search the knowledge base for information relevant to the query.
Args:
query: The search query describing what information you need.
max_results: Maximum number of results to return.
Returns:
Formatted text with the most relevant knowledge base excerpts.
"""
k = max_results if max_results is not None else self._top_k
if self._collections:
all_results = []
for col in self._collections:
col_results = await self._rag_service.search(query, collection=col, top_k=k)
all_results.extend(col_results)
all_results.sort(key=lambda r: r.score, reverse=True)
results = all_results[:k]
else:
results = await self._rag_service.search(query, top_k=k)
if not results:
return "No relevant results found in the knowledge base."
sections = []
for i, result in enumerate(results, 1):
source = result.document.metadata.get("source", "unknown")
chunk_idx = result.document.metadata.get("chunk_index", "?")
score = f"{result.score:.3f}"
sections.append(
f"[Result {i} | source: {source}, chunk: {chunk_idx}, relevance: {score}]\n{result.document.text}"
)
return "\n\n---\n\n".join(sections)
engine/src/stirling/rag/chunker.py
+118
View File
@@ -0,0 +1,118 @@
from __future__ import annotations
import re
# TODO: replace with pydantic-ai's built-in chunking once
# https://github.com/pydantic/pydantic-ai/issues/3962 lands.
def chunk_text(text: str, chunk_size: int = 512, overlap: int = 64) -> list[str]:
"""Split text into chunks of approximately chunk_size characters with overlap.
Splits on paragraph then sentence boundaries to avoid cutting mid-thought.
Returns an empty list for empty/whitespace-only input.
"""
text = text.strip()
if not text:
return []
paragraphs = _split_paragraphs(text)
chunks: list[str] = []
current: list[str] = []
current_len = 0
for para in paragraphs:
para_len = len(para)
if current_len + para_len <= chunk_size:
current.append(para)
current_len += para_len
continue
# If the current buffer has content, flush it
if current:
chunks.append("\n\n".join(current))
# If this paragraph alone exceeds chunk_size, split it by sentences
if para_len > chunk_size:
sentence_chunks = _split_long_paragraph(para, chunk_size, overlap)
chunks.extend(sentence_chunks)
current = []
current_len = 0
else:
# Start new chunk with overlap from previous chunk
overlap_text = _get_overlap(chunks, overlap) if chunks else ""
if overlap_text:
current = [overlap_text, para]
current_len = len(overlap_text) + para_len
else:
current = [para]
current_len = para_len
if current:
chunks.append("\n\n".join(current))
return [c.strip() for c in chunks if c.strip()]
def _split_paragraphs(text: str) -> list[str]:
"""Split text into paragraphs on double newlines."""
return [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
def _split_sentences(text: str) -> list[str]:
"""Split text into sentences, keeping the delimiter attached."""
parts = re.split(r"(?<=[.!?])\s+", text)
return [s.strip() for s in parts if s.strip()]
def _split_long_paragraph(paragraph: str, chunk_size: int, overlap: int) -> list[str]:
"""Split a single long paragraph into sentence-boundary chunks."""
sentences = _split_sentences(paragraph)
chunks: list[str] = []
current: list[str] = []
current_len = 0
for sentence in sentences:
sent_len = len(sentence)
if current_len + sent_len <= chunk_size:
current.append(sentence)
current_len += sent_len + 1 # +1 for space
continue
if current:
chunks.append(" ".join(current))
# If a single sentence exceeds chunk_size, force-split it
if sent_len > chunk_size:
for i in range(0, sent_len, chunk_size - overlap):
chunks.append(sentence[i : i + chunk_size])
current = []
current_len = 0
else:
overlap_text = _get_overlap(chunks, overlap) if chunks else ""
if overlap_text:
current = [overlap_text, sentence]
current_len = len(overlap_text) + sent_len + 1
else:
current = [sentence]
current_len = sent_len
if current:
chunks.append(" ".join(current))
return chunks
def _get_overlap(chunks: list[str], overlap: int) -> str:
"""Extract the last ~`overlap` characters from the most recent chunk, snapped to a word boundary."""
if not chunks or overlap <= 0:
return ""
last = chunks[-1]
tail = last[-overlap:] if len(last) > overlap else last
# Snap to the nearest word boundary to avoid starting mid-word
space_idx = tail.find(" ")
if space_idx > 0:
tail = tail[space_idx + 1 :]
return tail
engine/src/stirling/rag/embedder.py
+47
View File
@@ -0,0 +1,47 @@
from __future__ import annotations
from pydantic_ai import Embedder
from stirling.rag.chunker import chunk_text
from stirling.rag.store import Document
class EmbeddingService:
"""Wraps Pydantic AI's Embedder to provide document chunking and embedding."""
def __init__(self, model_name: str, chunk_size: int = 512, chunk_overlap: int = 64) -> None:
self._embedder = Embedder(model_name)
self._chunk_size = chunk_size
self._chunk_overlap = chunk_overlap
async def embed_query(self, text: str) -> list[float]:
"""Embed a search query, optimised for retrieval."""
result = await self._embedder.embed_query(text)
return list(result.embeddings[0])
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
"""Embed multiple document texts for indexing."""
if not texts:
return []
result = await self._embedder.embed_documents(texts)
return [list(emb) for emb in result.embeddings]
def chunk_and_prepare(
self,
text: str,
source: str = "",
base_metadata: dict[str, str] | None = None,
) -> list[Document]:
"""Chunk text and return Document objects ready for embedding.
Each chunk gets a unique ID based on source and chunk index.
"""
chunks = chunk_text(text, self._chunk_size, self._chunk_overlap)
documents: list[Document] = []
for i, chunk in enumerate(chunks):
meta = dict(base_metadata) if base_metadata else {}
meta["source"] = source
meta["chunk_index"] = str(i)
doc_id = f"{source}:chunk:{i}" if source else f"chunk:{i}"
documents.append(Document(id=doc_id, text=chunk, metadata=meta))
return documents
engine/src/stirling/rag/pgvector_store.py
+133
View File
@@ -0,0 +1,133 @@
from __future__ import annotations
import json
import psycopg
from pgvector.psycopg import register_vector_async
from stirling.rag.store import Document, SearchResult, VectorStore
class PgVectorStore(VectorStore):
"""PostgreSQL + pgvector backed store.
Connects to an external Postgres instance (DSN provided via config) and uses the
`vector` extension for similarity search. The schema is created on first use.
"""
def __init__(self, dsn: str) -> None:
if not dsn:
raise ValueError("pgvector backend requires a non-empty DSN (STIRLING_RAG_PGVECTOR_DSN)")
self._dsn = dsn
self._initialized = False
async def _connect(self) -> psycopg.AsyncConnection:
conn = await psycopg.AsyncConnection.connect(self._dsn)
await register_vector_async(conn)
return conn
async def _ensure_schema(self) -> None:
if self._initialized:
return
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
await cur.execute(
"""
CREATE TABLE IF NOT EXISTS rag_documents (
id TEXT NOT NULL,
collection TEXT NOT NULL,
text TEXT NOT NULL,
metadata JSONB NOT NULL DEFAULT '{}'::jsonb,
embedding vector NOT NULL,
PRIMARY KEY (id, collection)
)
"""
)
await cur.execute("CREATE INDEX IF NOT EXISTS idx_rag_collection ON rag_documents(collection)")
await conn.commit()
self._initialized = True
async def add_documents(
self,
collection: str,
documents: list[Document],
embeddings: list[list[float]],
) -> None:
if len(documents) != len(embeddings):
raise ValueError(f"Got {len(documents)} documents but {len(embeddings)} embeddings")
if not documents:
return
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
for doc, emb in zip(documents, embeddings):
await cur.execute(
"""
INSERT INTO rag_documents (id, collection, text, metadata, embedding)
VALUES (%s, %s, %s, %s::jsonb, %s)
ON CONFLICT (id, collection)
DO UPDATE SET
text = EXCLUDED.text,
metadata = EXCLUDED.metadata,
embedding = EXCLUDED.embedding
""",
(doc.id, collection, doc.text, json.dumps(doc.metadata), emb),
)
await conn.commit()
async def search(
self,
collection: str,
query_embedding: list[float],
top_k: int = 5,
) -> list[SearchResult]:
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute(
"""
SELECT id, text, metadata, 1 - (embedding <=> %s) AS score
FROM rag_documents
WHERE collection = %s
ORDER BY embedding <=> %s
LIMIT %s
""",
(query_embedding, collection, query_embedding, top_k),
)
rows = await cur.fetchall()
return [
SearchResult(
document=Document(id=r[0], text=r[1], metadata=r[2] or {}),
score=float(r[3]),
)
for r in rows
]
async def delete_collection(self, collection: str) -> None:
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute("DELETE FROM rag_documents WHERE collection = %s", (collection,))
await conn.commit()
async def list_collections(self) -> list[str]:
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute("SELECT DISTINCT collection FROM rag_documents ORDER BY collection")
rows = await cur.fetchall()
return [r[0] for r in rows]
async def has_collection(self, collection: str) -> bool:
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute(
"SELECT 1 FROM rag_documents WHERE collection = %s LIMIT 1",
(collection,),
)
row = await cur.fetchone()
return row is not None
engine/src/stirling/rag/service.py
+80
View File
@@ -0,0 +1,80 @@
from __future__ import annotations
import logging
from stirling.rag.embedder import EmbeddingService
from stirling.rag.store import Document, SearchResult, VectorStore
logger = logging.getLogger(__name__)
class RagService:
"""Orchestrates embedding and vector storage for RAG workflows."""
def __init__(self, embedder: EmbeddingService, store: VectorStore, default_top_k: int = 5) -> None:
self._embedder = embedder
self._store = store
self._default_top_k = default_top_k
async def index_text(
self,
collection: str,
text: str,
source: str = "",
metadata: dict[str, str] | None = None,
) -> int:
"""Chunk, embed, and store text. Returns the number of chunks indexed."""
documents = self._embedder.chunk_and_prepare(text, source=source, base_metadata=metadata)
if not documents:
return 0
embeddings = await self._embedder.embed_documents([doc.text for doc in documents])
await self._store.add_documents(collection, documents, embeddings)
return len(documents)
async def index_documents(self, collection: str, documents: list[Document]) -> int:
"""Embed and store pre-chunked documents. Returns the number stored."""
if not documents:
return 0
embeddings = await self._embedder.embed_documents([doc.text for doc in documents])
await self._store.add_documents(collection, documents, embeddings)
return len(documents)
async def search(
self,
query: str,
collection: str | None = None,
top_k: int | None = None,
) -> list[SearchResult]:
"""Embed query and search across one or all collections.
If collection is None, searches all available collections and merges results.
"""
k = top_k if top_k is not None else self._default_top_k
query_embedding = await self._embedder.embed_query(query)
if collection is not None:
if not await self._store.has_collection(collection):
return []
return await self._store.search(collection, query_embedding, k)
# Search all collections, skipping any that error (e.g. dimension mismatch)
collections = await self._store.list_collections()
all_results: list[SearchResult] = []
for col_name in collections:
try:
results = await self._store.search(col_name, query_embedding, k)
all_results.extend(results)
except Exception: # noqa: BLE001 — any backend error on one collection should not stop the others
logger.warning("Skipping collection %s during cross-collection search", col_name, exc_info=True)
# Sort by score descending, return top_k across all collections
all_results.sort(key=lambda r: r.score, reverse=True)
return all_results[:k]
async def delete_collection(self, collection: str) -> None:
"""Remove a collection and all its documents."""
await self._store.delete_collection(collection)
async def list_collections(self) -> list[str]:
"""List all available collections."""
return await self._store.list_collections()
engine/src/stirling/rag/sqlite_vec_store.py
+227
View File
@@ -0,0 +1,227 @@
from __future__ import annotations
import asyncio
import json
import math
import re
import sqlite3
from pathlib import Path
import sqlite_vec
from stirling.rag.store import Document, SearchResult, VectorStore
class SqliteVecStore(VectorStore):
"""sqlite-vec backed vector store. Single-file SQLite database, embedded, no server.
Each collection gets its own `vec0` virtual table with a fixed embedding dimension
(detected on first insert). Document metadata lives in a regular table joined by rowid.
"""
def __init__(self, db_path: str | Path) -> None:
is_memory = str(db_path) == ":memory:"
self._db_path: Path | None = None if is_memory else Path(db_path)
if self._db_path is not None:
self._db_path.parent.mkdir(parents=True, exist_ok=True)
conn = sqlite3.connect(str(self._db_path), check_same_thread=False)
else:
conn = sqlite3.connect(":memory:", check_same_thread=False)
conn.enable_load_extension(True)
sqlite_vec.load(conn)
conn.enable_load_extension(False)
if self._db_path is not None:
conn.execute("PRAGMA journal_mode=WAL")
self._conn = conn
self._lock = asyncio.Lock()
self._init_schema()
@classmethod
def ephemeral(cls) -> SqliteVecStore:
"""In-memory store for testing."""
return cls(":memory:")
def _init_schema(self) -> None:
self._conn.execute(
"""
CREATE TABLE IF NOT EXISTS collections (
name TEXT PRIMARY KEY,
dim INTEGER NOT NULL,
table_name TEXT NOT NULL
)
"""
)
self._conn.execute(
"""
CREATE TABLE IF NOT EXISTS documents (
id TEXT NOT NULL,
collection TEXT NOT NULL,
text TEXT NOT NULL,
metadata TEXT NOT NULL DEFAULT '{}',
vec_rowid INTEGER NOT NULL,
PRIMARY KEY (id, collection)
)
"""
)
self._conn.execute("CREATE INDEX IF NOT EXISTS idx_doc_collection ON documents(collection)")
self._conn.commit()
@staticmethod
def _sanitize_table_name(collection: str) -> str:
safe = re.sub(r"[^a-zA-Z0-9_]", "_", collection)
return f"vec_{safe}"
@staticmethod
def _normalize(vector: list[float]) -> list[float]:
norm = math.sqrt(sum(x * x for x in vector))
if norm == 0:
return list(vector)
return [x / norm for x in vector]
async def add_documents(
self,
collection: str,
documents: list[Document],
embeddings: list[list[float]],
) -> None:
if len(documents) != len(embeddings):
raise ValueError(f"Got {len(documents)} documents but {len(embeddings)} embeddings")
if not documents:
return
async with self._lock:
await asyncio.to_thread(self._sync_add, collection, documents, embeddings)
def _sync_add(
self,
collection: str,
documents: list[Document],
embeddings: list[list[float]],
) -> None:
dim = len(embeddings[0])
row = self._conn.execute("SELECT dim, table_name FROM collections WHERE name = ?", (collection,)).fetchone()
if row is None:
table_name = self._sanitize_table_name(collection)
self._conn.execute(f"CREATE VIRTUAL TABLE IF NOT EXISTS {table_name} USING vec0(embedding float[{dim}])")
self._conn.execute(
"INSERT INTO collections(name, dim, table_name) VALUES (?, ?, ?)",
(collection, dim, table_name),
)
else:
existing_dim, table_name = row
if existing_dim != dim:
raise ValueError(f"Collection {collection} has dim {existing_dim}, got embedding of dim {dim}")
# Upsert: delete existing docs with matching IDs first
ids = [doc.id for doc in documents]
placeholders = ",".join("?" * len(ids))
existing = self._conn.execute(
f"SELECT vec_rowid FROM documents WHERE collection = ? AND id IN ({placeholders})",
(collection, *ids),
).fetchall()
if existing:
vec_rowids = [r[0] for r in existing]
row_placeholders = ",".join("?" * len(vec_rowids))
self._conn.execute(
f"DELETE FROM {table_name} WHERE rowid IN ({row_placeholders})",
vec_rowids,
)
self._conn.execute(
f"DELETE FROM documents WHERE collection = ? AND id IN ({placeholders})",
(collection, *ids),
)
for doc, emb in zip(documents, embeddings):
normalized = self._normalize(list(emb))
cursor = self._conn.execute(
f"INSERT INTO {table_name}(embedding) VALUES (?)",
(sqlite_vec.serialize_float32(normalized),),
)
vec_rowid = cursor.lastrowid
self._conn.execute(
"INSERT INTO documents(id, collection, text, metadata, vec_rowid) VALUES (?, ?, ?, ?, ?)",
(doc.id, collection, doc.text, json.dumps(doc.metadata), vec_rowid),
)
self._conn.commit()
async def search(
self,
collection: str,
query_embedding: list[float],
top_k: int = 5,
) -> list[SearchResult]:
async with self._lock:
return await asyncio.to_thread(self._sync_search, collection, query_embedding, top_k)
def _sync_search(
self,
collection: str,
query_embedding: list[float],
top_k: int,
) -> list[SearchResult]:
row = self._conn.execute("SELECT table_name, dim FROM collections WHERE name = ?", (collection,)).fetchone()
if row is None:
return []
table_name, dim = row
if len(query_embedding) != dim:
raise ValueError(f"Query embedding dim {len(query_embedding)} does not match collection dim {dim}")
normalized = self._normalize(list(query_embedding))
query_blob = sqlite_vec.serialize_float32(normalized)
results = self._conn.execute(
f"""
SELECT d.id, d.text, d.metadata, v.distance
FROM {table_name} v
JOIN documents d ON d.vec_rowid = v.rowid AND d.collection = ?
WHERE v.embedding MATCH ? AND k = ?
ORDER BY v.distance
""",
(collection, query_blob, top_k),
).fetchall()
return [
SearchResult(
document=Document(
id=r[0],
text=r[1],
metadata=json.loads(r[2]) if r[2] else {},
),
# For normalized vectors: cosine_sim = 1 - (L2^2 / 2)
score=max(0.0, 1.0 - (r[3] ** 2) / 2.0),
)
for r in results
]
async def delete_collection(self, collection: str) -> None:
async with self._lock:
await asyncio.to_thread(self._sync_delete_collection, collection)
def _sync_delete_collection(self, collection: str) -> None:
row = self._conn.execute("SELECT table_name FROM collections WHERE name = ?", (collection,)).fetchone()
if row is None:
return
table_name = row[0]
self._conn.execute(f"DROP TABLE IF EXISTS {table_name}")
self._conn.execute("DELETE FROM documents WHERE collection = ?", (collection,))
self._conn.execute("DELETE FROM collections WHERE name = ?", (collection,))
self._conn.commit()
async def list_collections(self) -> list[str]:
async with self._lock:
return await asyncio.to_thread(self._sync_list_collections)
def _sync_list_collections(self) -> list[str]:
rows = self._conn.execute("SELECT name FROM collections ORDER BY name").fetchall()
return [r[0] for r in rows]
async def has_collection(self, collection: str) -> bool:
async with self._lock:
return await asyncio.to_thread(self._sync_has_collection, collection)
def _sync_has_collection(self, collection: str) -> bool:
row = self._conn.execute("SELECT 1 FROM collections WHERE name = ?", (collection,)).fetchone()
return row is not None
engine/src/stirling/rag/store.py
+59
View File
@@ -0,0 +1,59 @@
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
@dataclass
class Document:
"""A chunk of text with metadata, ready for embedding and storage."""
id: str
text: str
metadata: dict[str, str] = field(default_factory=dict)
@dataclass
class SearchResult:
"""A document returned from a vector search with its relevance score."""
document: Document
score: float
class VectorStore(ABC):
"""Abstract interface for vector storage backends.
Implementations must handle persistence, collection management,
and nearest-neighbor search over pre-computed embeddings.
"""
@abstractmethod
async def add_documents(
self,
collection: str,
documents: list[Document],
embeddings: list[list[float]],
) -> None:
"""Store documents with their embeddings in the named collection."""
@abstractmethod
async def search(
self,
collection: str,
query_embedding: list[float],
top_k: int = 5,
) -> list[SearchResult]:
"""Return the top_k most similar documents from the collection."""
@abstractmethod
async def delete_collection(self, collection: str) -> None:
"""Remove a collection and all its documents."""
@abstractmethod
async def list_collections(self) -> list[str]:
"""Return names of all existing collections."""
@abstractmethod
async def has_collection(self, collection: str) -> bool:
"""Check whether a collection exists."""
engine/src/stirling/services/__init__.py
+2
View File
@@ -1,9 +1,11 @@
"""Shared services used by the Stirling AI runtime."""
from .runtime import AppRuntime, build_model_settings, build_runtime
from .tracking import setup_posthog_tracking
__all__ = [
"AppRuntime",
"build_model_settings",
"build_runtime",
"setup_posthog_tracking",
]
engine/src/stirling/services/runtime.py
+49 -1
View File
@@ -1,11 +1,23 @@
from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import assert_never
from pydantic_ai.models import Model, infer_model
from pydantic_ai.settings import ModelSettings
from stirling.config import AppSettings
from stirling.config import ENGINE_ROOT, AppSettings, RagBackend
from stirling.rag import (
EmbeddingService,
PgVectorStore,
RagCapability,
RagService,
SqliteVecStore,
VectorStore,
)
logger = logging.getLogger(__name__)
@dataclass(frozen=True)
@@ -13,6 +25,8 @@ class AppRuntime:
settings: AppSettings
fast_model: Model
smart_model: Model
rag_service: RagService
rag_capability: RagCapability
@property
def fast_model_settings(self) -> ModelSettings:
@@ -39,13 +53,47 @@ def validate_structured_output_support(model: Model, model_name: str) -> None:
raise ValueError(f"Unsupported model {model_name}. This model does not support structured outputs.")
def _build_vector_store(settings: AppSettings) -> VectorStore:
"""Build the configured vector store backend."""
if settings.rag_backend == RagBackend.SQLITE:
store_path = settings.rag_store_path
# Treat ":memory:" as a special in-process token; otherwise resolve against the engine root.
if str(store_path) != ":memory:" and not store_path.is_absolute():
store_path = ENGINE_ROOT / store_path
logger.info("RAG backend=sqlite, db_path=%s", store_path)
return SqliteVecStore(db_path=store_path)
if settings.rag_backend == RagBackend.PGVECTOR:
logger.info("RAG backend=pgvector, dsn=<configured>")
return PgVectorStore(dsn=settings.rag_pgvector_dsn)
assert_never(settings.rag_backend)
def _build_rag(settings: AppSettings) -> tuple[RagService, RagCapability]:
"""Build the RAG service and capability."""
logger.info("RAG: embedding_model=%s", settings.rag_embedding_model)
embedder = EmbeddingService(
model_name=settings.rag_embedding_model,
chunk_size=settings.rag_chunk_size,
chunk_overlap=settings.rag_chunk_overlap,
)
store = _build_vector_store(settings)
service = RagService(embedder=embedder, store=store, default_top_k=settings.rag_default_top_k)
capability = RagCapability(rag_service=service, top_k=settings.rag_default_top_k)
return service, capability
def build_runtime(settings: AppSettings) -> AppRuntime:
fast_model = infer_model(settings.fast_model_name)
smart_model = infer_model(settings.smart_model_name)
validate_structured_output_support(fast_model, settings.fast_model_name)
validate_structured_output_support(smart_model, settings.smart_model_name)
rag_service, rag_capability = _build_rag(settings)
return AppRuntime(
settings=settings,
fast_model=fast_model,
smart_model=smart_model,
rag_service=rag_service,
rag_capability=rag_capability,
)
engine/src/stirling/services/tracking.py
+233
View File
@@ -0,0 +1,233 @@
from __future__ import annotations
import json
from collections import OrderedDict
from collections.abc import Mapping
from contextvars import ContextVar
from typing import Any
from opentelemetry.context import Context
from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor, TracerProvider
from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( # No public import for these constants yet
GEN_AI_INPUT_MESSAGES,
GEN_AI_OPERATION_NAME,
GEN_AI_OUTPUT_MESSAGES,
GEN_AI_REQUEST_MAX_TOKENS,
GEN_AI_REQUEST_MODEL,
GEN_AI_REQUEST_TEMPERATURE,
GEN_AI_RESPONSE_MODEL,
GEN_AI_SYSTEM,
GEN_AI_TOOL_DEFINITIONS,
GEN_AI_USAGE_INPUT_TOKENS,
GEN_AI_USAGE_OUTPUT_TOKENS,
GenAiOperationNameValues,
)
from opentelemetry.semconv.attributes.server_attributes import SERVER_ADDRESS, SERVER_PORT
from opentelemetry.trace import Span
from posthog.client import Client as PostHogClient
from stirling.config import AppSettings
# Per-request user ID, set by middleware from the X-User-Id header.
# When not set, PostHog generates a random ID and marks the event as personless.
current_user_id: ContextVar[str | None] = ContextVar("current_user_id", default=None)
class LRUSet:
"""Least Recently Used Set: a set with a maximum size that evicts the oldest entries first."""
def __init__(self, max_size: int) -> None:
self._max_size = max_size
self._data: OrderedDict[str, None] = OrderedDict()
def __contains__(self, key: str) -> bool:
return key in self._data
def add(self, key: str) -> None:
self._data[key] = None
if len(self._data) > self._max_size:
self._data.popitem(last=False)
def _parse_json_attr(attrs: Mapping[str, Any], key: str) -> Any | None:
"""Parse a JSON string span attribute, returning None on failure."""
raw = attrs.get(key)
if raw is None:
return None
try:
return json.loads(str(raw))
except (json.JSONDecodeError, TypeError):
return None
def _transform_output_choices(choices: list[Any]) -> list[Any]:
"""Transform Pydantic AI's parts-based output format to PostHog-compatible format.
Pydantic AI emits: ``[{"role": "assistant", "parts": [{"type": "tool_call", "name": "..."}]}]``
PostHog expects: ``[{"role": "assistant", "tool_calls": [{"type": "function", "function": {"name": "..."}}]}]``
"""
for choice in choices:
if not isinstance(choice, dict) or "parts" not in choice:
continue
tool_calls = []
for part in choice.get("parts", []):
if isinstance(part, dict) and part.get("type") == "tool_call":
tool_calls.append(
{
"type": "function",
"id": part.get("id", ""),
"function": {"name": part.get("name", "")},
}
)
if tool_calls:
choice["tool_calls"] = tool_calls
choice["content"] = choice.pop("parts")
return choices
def _extract_user_message(attrs: Mapping[str, Any]) -> str:
"""Extract the last user message text from the input messages span attribute."""
messages = _parse_json_attr(attrs, GEN_AI_INPUT_MESSAGES)
if not isinstance(messages, list):
return ""
for msg in reversed(messages):
if not isinstance(msg, dict):
continue
if msg.get("role") == "user":
for part in msg.get("parts", []):
if isinstance(part, dict) and part.get("type") == "text":
return str(part.get("content", ""))
return ""
# TODO: Replace with an official PostHog integration if one ever exists
class PostHogSpanProcessor(SpanProcessor):
"""Translates Pydantic AI OpenTelemetry spans into PostHog $ai_generation events."""
def __init__(self, client: PostHogClient) -> None:
self._client = client
self._seen_traces = LRUSet(max_size=10_000)
def on_start(self, span: Span, parent_context: Context | None = None) -> None:
pass
def on_end(self, span: ReadableSpan) -> None:
attrs = dict(span.attributes or {})
if attrs.get(GEN_AI_OPERATION_NAME) != GenAiOperationNameValues.CHAT.value:
return
properties = self._build_generation_properties(span, attrs)
self._maybe_emit_trace_event(span, attrs, properties)
self._client.capture(
distinct_id=current_user_id.get(),
event="$ai_generation",
properties=properties,
)
def _build_generation_properties(self, span: ReadableSpan, attrs: Mapping[str, Any]) -> dict[str, object]:
"""Build the $ai_generation event properties from span data."""
properties: dict[str, object] = {
"$ai_provider": attrs.get(GEN_AI_SYSTEM, ""),
"$ai_model": attrs.get(GEN_AI_RESPONSE_MODEL) or attrs.get(GEN_AI_REQUEST_MODEL, ""),
"$ai_input_tokens": attrs.get(GEN_AI_USAGE_INPUT_TOKENS, 0),
"$ai_output_tokens": attrs.get(GEN_AI_USAGE_OUTPUT_TOKENS, 0),
}
if span.context:
properties["$ai_trace_id"] = format(span.context.trace_id, "032x")
properties["$ai_span_id"] = format(span.context.span_id, "016x")
if span.parent and span.parent.span_id:
properties["$ai_parent_id"] = format(span.parent.span_id, "016x")
if span.start_time and span.end_time:
properties["$ai_latency"] = (span.end_time - span.start_time) / 1e9
self._add_message_properties(properties, attrs)
self._add_model_parameters(properties, attrs)
self._add_tool_definitions(properties, attrs)
self._add_base_url(properties, attrs)
return properties
def _maybe_emit_trace_event(
self, span: ReadableSpan, attrs: Mapping[str, Any], properties: dict[str, object]
) -> None:
"""Emit an $ai_trace event for the first span seen per trace ID."""
trace_id = str(properties.get("$ai_trace_id", ""))
if not trace_id or trace_id in self._seen_traces:
return
self._seen_traces.add(trace_id)
trace_properties: dict[str, object] = {
"$ai_trace_id": trace_id,
"$ai_trace_name": _extract_user_message(attrs),
"$ai_provider": attrs.get(GEN_AI_SYSTEM, ""),
}
if span.start_time and span.end_time:
trace_properties["$ai_latency"] = (span.end_time - span.start_time) / 1e9
self._client.capture(
distinct_id=current_user_id.get(),
event="$ai_trace",
properties=trace_properties,
)
@staticmethod
def _add_message_properties(properties: dict[str, object], attrs: Mapping[str, Any]) -> None:
input_messages = _parse_json_attr(attrs, GEN_AI_INPUT_MESSAGES)
if input_messages is not None:
properties["$ai_input"] = input_messages
output_messages = _parse_json_attr(attrs, GEN_AI_OUTPUT_MESSAGES)
if isinstance(output_messages, list):
properties["$ai_output_choices"] = _transform_output_choices(output_messages)
elif output_messages is not None:
properties["$ai_output_choices"] = output_messages
@staticmethod
def _add_model_parameters(properties: dict[str, object], attrs: Mapping[str, Any]) -> None:
model_parameters: dict[str, object] = {}
if GEN_AI_REQUEST_TEMPERATURE in attrs:
model_parameters["temperature"] = attrs[GEN_AI_REQUEST_TEMPERATURE]
if GEN_AI_REQUEST_MAX_TOKENS in attrs:
model_parameters["max_tokens"] = attrs[GEN_AI_REQUEST_MAX_TOKENS]
if model_parameters:
properties["$ai_model_parameters"] = model_parameters
@staticmethod
def _add_tool_definitions(properties: dict[str, object], attrs: Mapping[str, Any]) -> None:
tools = _parse_json_attr(attrs, GEN_AI_TOOL_DEFINITIONS)
if tools is not None:
properties["$ai_tools"] = tools
@staticmethod
def _add_base_url(properties: dict[str, object], attrs: Mapping[str, Any]) -> None:
parts: list[str] = []
if host := attrs.get(SERVER_ADDRESS):
parts.append(str(host))
if port := attrs.get(SERVER_PORT):
parts.append(str(port))
if parts:
properties["$ai_base_url"] = ":".join(parts)
def shutdown(self) -> None:
self._client.shutdown()
def force_flush(self, timeout_millis: int = 30000) -> bool:
self._client.flush()
return True
def setup_posthog_tracking(settings: AppSettings) -> TracerProvider | None:
"""Configure OpenTelemetry with a PostHog span processor for LLM analytics.
Returns the TracerProvider so it can be shut down on app exit,
or None when tracking is disabled.
"""
if not settings.posthog_enabled or not settings.posthog_api_key:
return None
client = PostHogClient(project_api_key=settings.posthog_api_key, host=settings.posthog_host)
processor = PostHogSpanProcessor(client)
provider = TracerProvider()
provider.add_span_processor(processor)
return provider
engine/tests/conftest.py
+46
View File
@@ -0,0 +1,46 @@
from __future__ import annotations
from collections.abc import Iterator
from pathlib import Path
import pytest
from stirling.config import AppSettings, RagBackend, load_settings
from stirling.services import build_runtime
from stirling.services.runtime import AppRuntime
@pytest.fixture(autouse=True)
def clear_settings_cache() -> Iterator[None]:
load_settings.cache_clear()
yield
load_settings.cache_clear()
def build_app_settings() -> AppSettings:
return AppSettings(
smart_model_name="test",
fast_model_name="test",
smart_model_max_tokens=8192,
fast_model_max_tokens=2048,
rag_backend=RagBackend.SQLITE,
rag_embedding_model="voyageai:voyage-4",
rag_store_path=Path(":memory:"),
rag_pgvector_dsn="",
rag_chunk_size=512,
rag_chunk_overlap=64,
rag_default_top_k=5,
posthog_enabled=False,
posthog_api_key="",
posthog_host="https://eu.i.posthog.com",
)
@pytest.fixture
def app_settings() -> AppSettings:
return build_app_settings()
@pytest.fixture
def runtime(app_settings: AppSettings) -> AppRuntime:
return build_runtime(app_settings)
engine/tests/ledger/__init__.py
View File
engine/tests/ledger/test_arithmetic_scanner.py
+133
View File
@@ -0,0 +1,133 @@
"""
ArithmeticScanner — unit tests.
Tests cover the two inline arithmetic patterns the scanner targets:
1. Equals expressions: A + B = C
2. Total-then-addends: Total: C (A + B)
"""
from decimal import Decimal
import pytest
from stirling.agents.ledger.validators.arithmetic import ArithmeticScanner
@pytest.fixture
def scanner() -> ArithmeticScanner:
return ArithmeticScanner(tolerance=Decimal("0.01"))
# ---------------------------------------------------------------------------
# Equals expressions: A + B + C = D
# ---------------------------------------------------------------------------
def test_correct_equals_expression(scanner: ArithmeticScanner) -> None:
"""A correct sum should produce no findings."""
text = "The total cost is 100 + 200 + 150 = 450."
assert scanner.scan(page=0, text=text) == []
def test_wrong_equals_expression(scanner: ArithmeticScanner) -> None:
"""An incorrect sum should produce one error discrepancy."""
text = "Revenue: 500 + 300 = 900" # should be 800
discrepancies = scanner.scan(page=3, text=text)
assert len(discrepancies) == 1
d = discrepancies[0]
assert d.page == 3
assert d.kind == "arithmetic"
assert d.severity == "error"
assert d.stated == "900"
assert d.expected == "800"
def test_subtraction_expression(scanner: ArithmeticScanner) -> None:
"""Subtraction in expressions should be evaluated correctly."""
text = "Net: 1000 - 250 = 750"
assert scanner.scan(page=0, text=text) == []
def test_wrong_subtraction(scanner: ArithmeticScanner) -> None:
text = "Net: 1000 - 250 = 800" # should be 750
discrepancies = scanner.scan(page=0, text=text)
assert len(discrepancies) == 1
assert discrepancies[0].expected == "750"
def test_currency_symbols_stripped(scanner: ArithmeticScanner) -> None:
"""Currency symbols and thousand separators must not break parsing."""
text = "Total: £1,000 + £500 = £1,500"
assert scanner.scan(page=0, text=text) == []
def test_multiple_expressions_in_text(scanner: ArithmeticScanner) -> None:
"""Multiple expressions in the same text should each be evaluated."""
text = (
"Q1 revenue: 100 + 200 = 300. "
"Q2 revenue: 150 + 100 = 350. " # wrong: should be 250
)
discrepancies = scanner.scan(page=0, text=text)
assert len(discrepancies) == 1
assert discrepancies[0].expected == "250"
# ---------------------------------------------------------------------------
# Total-then-addends: "Total: X (A + B + C)"
# ---------------------------------------------------------------------------
def test_correct_total_then_addends(scanner: ArithmeticScanner) -> None:
text = "Grand Total: 750 (300 + 250 + 200)"
assert scanner.scan(page=0, text=text) == []
def test_wrong_total_then_addends(scanner: ArithmeticScanner) -> None:
text = "Grand Total: 900 (300 + 250 + 200)" # addends sum to 750
discrepancies = scanner.scan(page=0, text=text)
assert len(discrepancies) == 1
d = discrepancies[0]
assert d.stated == "900"
assert d.expected == "750"
def test_total_keyword_variations(scanner: ArithmeticScanner) -> None:
"""The pattern must work for 'Sum', 'Subtotal', 'Grand Total' etc."""
cases = [
("Sum: 600 (200 + 200 + 200)", True),
("Subtotal: 600 (200 + 200 + 200)", True),
("Total: 999 (200 + 200 + 200)", False), # wrong
]
for text, should_be_clean in cases:
result = scanner.scan(page=0, text=text)
if should_be_clean:
assert result == [], f"Expected clean for: {text!r}"
else:
assert len(result) == 1, f"Expected error for: {text!r}"
# ---------------------------------------------------------------------------
# Edge cases
# ---------------------------------------------------------------------------
def test_no_expressions_in_text(scanner: ArithmeticScanner) -> None:
text = "This paragraph discusses revenue trends but contains no arithmetic."
assert scanner.scan(page=0, text=text) == []
def test_empty_text(scanner: ArithmeticScanner) -> None:
assert scanner.scan(page=0, text="") == []
def test_leading_negative_expression(scanner: ArithmeticScanner) -> None:
"""Expressions starting with a negative number should evaluate correctly."""
text = "Adjustment: -100 + 250 = 150"
assert scanner.scan(page=0, text=text) == []
def test_leading_negative_wrong(scanner: ArithmeticScanner) -> None:
text = "Adjustment: -100 + 250 = 200" # should be 150
discrepancies = scanner.scan(page=0, text=text)
assert len(discrepancies) == 1
assert discrepancies[0].expected == "150"
engine/tests/ledger/test_figure_tracker.py
+100
View File
@@ -0,0 +1,100 @@
"""
FigureTracker — unit tests.
Tests that named figures are correctly accumulated and that conflicting
sightings (same label, different value) are surfaced as consistency warnings.
"""
from decimal import Decimal
import pytest
from stirling.agents.ledger.validators.figures import FigureTracker
@pytest.fixture
def tracker() -> FigureTracker:
return FigureTracker(tolerance=Decimal("0.01"))
# ---------------------------------------------------------------------------
# No conflicts
# ---------------------------------------------------------------------------
def test_no_conflicts_single_figure(tracker: FigureTracker) -> None:
tracker.record("Net Profit", Decimal("1200.00"), page=3, raw="£1,200.00")
assert tracker.conflicts() == []
def test_no_conflicts_consistent_figure(tracker: FigureTracker) -> None:
"""The same figure cited identically on two pages must not raise a conflict."""
tracker.record("Total Revenue", Decimal("5000.00"), page=1, raw="£5,000")
tracker.record("Total Revenue", Decimal("5000.00"), page=8, raw="£5,000")
assert tracker.conflicts() == []
def test_no_conflicts_within_tolerance(tracker: FigureTracker) -> None:
"""A difference within tolerance must not be flagged."""
tracker.record("VAT", Decimal("100.00"), page=2, raw="£100.00")
tracker.record("VAT", Decimal("100.005"), page=5, raw="£100.005")
assert tracker.conflicts() == []
# ---------------------------------------------------------------------------
# Conflicts
# ---------------------------------------------------------------------------
def test_conflict_different_values(tracker: FigureTracker) -> None:
"""Same label, different value on two pages → one consistency warning."""
tracker.record("Net Profit", Decimal("1200.00"), page=3, raw="£1,200")
tracker.record("Net Profit", Decimal("1250.00"), page=7, raw="£1,250")
conflicts = tracker.conflicts()
assert len(conflicts) == 1
d = conflicts[0]
assert d.kind == "consistency"
assert d.severity == "warning"
assert d.page == 7 # later occurrence is flagged
def test_conflict_three_sightings_two_values(tracker: FigureTracker) -> None:
"""Three sightings where one differs from canonical → 1 conflict."""
tracker.record("Revenue", Decimal("1000"), page=1, raw="£1,000")
tracker.record("Revenue", Decimal("1000"), page=3, raw="£1,000")
tracker.record("Revenue", Decimal("999"), page=5, raw="£999")
conflicts = tracker.conflicts()
# Canonical=p1 (1000). p3 matches, p5 differs → 1 conflict
assert len(conflicts) == 1
assert conflicts[0].page == 5
# ---------------------------------------------------------------------------
# Label normalisation
# ---------------------------------------------------------------------------
def test_label_normalisation_case_insensitive(tracker: FigureTracker) -> None:
"""Labels must be compared case-insensitively."""
tracker.record("Net Profit", Decimal("1200"), page=2, raw="1200")
tracker.record("net profit", Decimal("1100"), page=4, raw="1100")
assert len(tracker.conflicts()) == 1
def test_label_normalisation_punctuation(tracker: FigureTracker) -> None:
"""Colons and dashes in labels must be normalised before comparison."""
tracker.record("Total Revenue:", Decimal("5000"), page=1, raw="5000")
tracker.record("Total Revenue —", Decimal("4000"), page=9, raw="4000")
assert len(tracker.conflicts()) == 1
# ---------------------------------------------------------------------------
# Entry count
# ---------------------------------------------------------------------------
def test_entry_count(tracker: FigureTracker) -> None:
tracker.record("A", Decimal("1"), page=0, raw="1")
tracker.record("A", Decimal("1"), page=1, raw="1")
tracker.record("B", Decimal("2"), page=2, raw="2")
assert tracker.entry_count == 3
engine/tests/ledger/test_formula_evaluator.py
+205
View File
@@ -0,0 +1,205 @@
"""
FormulaEvaluator — unit tests.
Tests cover:
- Operator precedence (* / before + -)
- Column reference replacement (colN with word boundaries)
- Negative number handling
- each_row, column_total, and single_cell scopes
"""
from decimal import Decimal
import pytest
from stirling.agents.ledger.validators.formula import FormulaEvaluator
@pytest.fixture
def evaluator() -> FormulaEvaluator:
return FormulaEvaluator(tolerance=Decimal("0.01"))
# ---------------------------------------------------------------------------
# _safe_eval — operator precedence
# ---------------------------------------------------------------------------
def test_safe_eval_addition(evaluator: FormulaEvaluator) -> None:
assert evaluator._safe_eval("2 + 3") == Decimal("5")
def test_safe_eval_multiplication_before_addition(evaluator: FormulaEvaluator) -> None:
"""2 + 3 * 4 should be 14, not 20."""
assert evaluator._safe_eval("2 + 3 * 4") == Decimal("14")
def test_safe_eval_division_before_subtraction(evaluator: FormulaEvaluator) -> None:
"""10 - 6 / 2 should be 7, not 2."""
assert evaluator._safe_eval("10 - 6 / 2") == Decimal("7")
def test_safe_eval_mixed_precedence(evaluator: FormulaEvaluator) -> None:
"""1 + 2 * 3 - 4 / 2 should be 1 + 6 - 2 = 5."""
assert evaluator._safe_eval("1 + 2 * 3 - 4 / 2") == Decimal("5")
def test_safe_eval_all_multiplication(evaluator: FormulaEvaluator) -> None:
assert evaluator._safe_eval("2 * 3 * 4") == Decimal("24")
def test_safe_eval_division_by_zero(evaluator: FormulaEvaluator) -> None:
assert evaluator._safe_eval("10 / 0") is None
def test_safe_eval_negative_result(evaluator: FormulaEvaluator) -> None:
assert evaluator._safe_eval("3 - 5") == Decimal("-2")
def test_safe_eval_leading_negative(evaluator: FormulaEvaluator) -> None:
"""Expressions starting with a negative number should work."""
result = evaluator._safe_eval("-100 + 200")
assert result == Decimal("100")
def test_safe_eval_empty(evaluator: FormulaEvaluator) -> None:
assert evaluator._safe_eval("") is None
def test_safe_eval_single_number(evaluator: FormulaEvaluator) -> None:
assert evaluator._safe_eval("42") == Decimal("42")
def test_safe_eval_decimal_numbers(evaluator: FormulaEvaluator) -> None:
result = evaluator._safe_eval("1.5 * 2 + 0.5")
assert result == Decimal("3.5")
# ---------------------------------------------------------------------------
# colN replacement — word boundary safety
# ---------------------------------------------------------------------------
def test_col1_does_not_corrupt_col12(evaluator: FormulaEvaluator) -> None:
"""col1 replacement must not alter col12."""
csv = "a,b,c,d,e,f,g,h,i,j,k,l,m\n0,10,0,0,0,0,0,0,0,0,0,0,120\n"
# col1=10, col12=120 → col12 - col1 should be 110
result = evaluator.evaluate(
page=0,
table_csv=csv,
formula="col0 = col12 - col1",
scope="each_row",
description="test",
)
# row 1: col0=0, expected=120-10=110 → discrepancy
assert len(result) == 1
assert result[0].expected == "110"
def test_col_replacement_adjacent_columns(evaluator: FormulaEvaluator) -> None:
"""col1 and col10 should both be replaced correctly."""
csv = "a,b,c,d,e,f,g,h,i,j,k\n55,5,0,0,0,0,0,0,0,0,50\n"
# col0=55, col1=5, col10=50 → col1 + col10 = 55
result = evaluator.evaluate(
page=0,
table_csv=csv,
formula="col0 = col1 + col10",
scope="each_row",
description="test",
)
assert result == [] # 5 + 50 = 55, matches col0
# ---------------------------------------------------------------------------
# each_row scope
# ---------------------------------------------------------------------------
def test_each_row_correct(evaluator: FormulaEvaluator) -> None:
csv = "Item,Qty,Price,Total\nWidget,10,5,50\nGadget,3,20,60\n"
result = evaluator.evaluate(
page=0,
table_csv=csv,
formula="col3 = col1 * col2",
scope="each_row",
description="unit price check",
)
assert result == []
def test_each_row_error(evaluator: FormulaEvaluator) -> None:
csv = "Item,Qty,Price,Total\nWidget,10,5,50\nGadget,3,20,99\n"
result = evaluator.evaluate(
page=0,
table_csv=csv,
formula="col3 = col1 * col2",
scope="each_row",
description="unit price check",
)
assert len(result) == 1
assert result[0].expected == "60"
assert result[0].stated == "99"
# ---------------------------------------------------------------------------
# column_total scope
# ---------------------------------------------------------------------------
def test_column_total_correct(evaluator: FormulaEvaluator) -> None:
csv = "Name,Amount\nA,100\nB,200\nTotal,300\n"
result = evaluator.evaluate(
page=0,
table_csv=csv,
formula="sum",
scope="column_total",
description="total check",
target_row=3,
target_col=1,
)
assert result == []
def test_column_total_error(evaluator: FormulaEvaluator) -> None:
csv = "Name,Amount\nA,100\nB,200\nTotal,400\n"
result = evaluator.evaluate(
page=0,
table_csv=csv,
formula="sum",
scope="column_total",
description="total check",
target_row=3,
target_col=1,
)
assert len(result) == 1
assert result[0].expected == "300"
# ---------------------------------------------------------------------------
# single_cell scope
# ---------------------------------------------------------------------------
def test_single_cell_correct(evaluator: FormulaEvaluator) -> None:
csv = "A,B,C\n10,20,30\n5,15,20\n15,35,50\n"
result = evaluator.evaluate(
page=0,
table_csv=csv,
formula="cell(3,2) = cell(1,2) + cell(2,2)",
scope="single_cell",
description="grand total",
)
assert result == []
def test_single_cell_error(evaluator: FormulaEvaluator) -> None:
csv = "A,B,C\n10,20,30\n5,15,20\n15,35,99\n"
result = evaluator.evaluate(
page=0,
table_csv=csv,
formula="cell(3,2) = cell(1,2) + cell(2,2)",
scope="single_cell",
description="grand total",
)
assert len(result) == 1
assert result[0].expected == "50"
engine/tests/ledger/test_models.py
+144
View File
@@ -0,0 +1,144 @@
"""
Ledger models — unit tests for serialisation and business logic.
These tests confirm the wire contract: models round-trip through JSON
correctly and their helper properties behave as documented.
"""
import pytest
from pydantic import ValidationError
from stirling.contracts.ledger import (
Discrepancy,
DiscrepancyKind,
Evidence,
Folio,
FolioManifest,
FolioType,
Requisition,
Severity,
Verdict,
)
# ---------------------------------------------------------------------------
# FolioManifest
# ---------------------------------------------------------------------------
def test_folio_manifest_round_trip() -> None:
manifest = FolioManifest(
session_id="abc-123",
page_count=3,
folio_types=[FolioType.TEXT, FolioType.IMAGE, FolioType.MIXED],
)
reloaded = FolioManifest.model_validate_json(manifest.model_dump_json())
assert reloaded == manifest
def test_folio_manifest_round_bounds() -> None:
with pytest.raises(ValidationError):
FolioManifest(session_id="x", page_count=1, folio_types=[FolioType.TEXT], round=0)
with pytest.raises(ValidationError):
FolioManifest(session_id="x", page_count=1, folio_types=[FolioType.TEXT], round=4)
# ---------------------------------------------------------------------------
# Requisition
# ---------------------------------------------------------------------------
def test_requisition_empty() -> None:
req = Requisition(rationale="nothing needed")
assert req.need_text == []
assert req.need_tables == []
assert req.need_ocr == []
def test_requisition_type_discriminator() -> None:
req = Requisition(need_text=[0, 1], rationale="needs text")
assert req.type == "requisition"
# ---------------------------------------------------------------------------
# Folio.readable_text
# ---------------------------------------------------------------------------
def test_folio_readable_text_prefers_ocr() -> None:
folio = Folio(page=0, text="digital text", ocr_text="ocr text")
assert folio.readable_text == "ocr text"
def test_folio_readable_text_falls_back_to_text() -> None:
folio = Folio(page=0, text="digital text")
assert folio.readable_text == "digital text"
def test_folio_readable_text_empty_when_none() -> None:
folio = Folio(page=0)
assert folio.readable_text == ""
# ---------------------------------------------------------------------------
# Verdict
# ---------------------------------------------------------------------------
def test_verdict_clean_flag() -> None:
verdict = Verdict(
session_id="s1",
discrepancies=[],
pages_examined=[0, 1],
rounds_taken=2,
summary="All figures balance.",
clean=True,
)
assert verdict.error_count == 0
assert verdict.warning_count == 0
assert verdict.clean is True
def test_verdict_error_and_warning_counts() -> None:
discrepancies = [
Discrepancy(
page=0,
kind=DiscrepancyKind.TALLY,
severity=Severity.ERROR,
description="bad sum",
stated="100",
expected="110",
),
Discrepancy(
page=1,
kind=DiscrepancyKind.CONSISTENCY,
severity=Severity.WARNING,
description="mismatched figure",
stated="500",
expected="550",
),
]
verdict = Verdict(
session_id="s1",
discrepancies=discrepancies,
pages_examined=[0, 1],
rounds_taken=1,
summary="Issues found.",
clean=False,
)
assert verdict.error_count == 1
assert verdict.warning_count == 1
# ---------------------------------------------------------------------------
# Evidence.final_round
# ---------------------------------------------------------------------------
def test_evidence_final_round() -> None:
evidence = Evidence(
session_id="s",
folios=[Folio(page=0, text="hello")],
round=3,
final_round=True,
)
assert evidence.final_round is True
engine/tests/ledger/test_routes.py
+243
View File
@@ -0,0 +1,243 @@
"""
Ledger Auditor — FastAPI route tests.
Uses FastAPI's TestClient with dependency overrides. All LLM calls are
mocked out; these tests exercise HTTP parsing, serialisation, and response
enveloping only — not the agent's reasoning.
"""
from __future__ import annotations
from collections.abc import Iterator
from decimal import Decimal
import pytest
from fastapi.testclient import TestClient
from stirling.api import app
from stirling.api.dependencies import get_math_auditor_agent
from stirling.config import AppSettings, load_settings
from stirling.contracts.ledger import (
Discrepancy,
DiscrepancyKind,
Evidence,
FolioManifest,
Requisition,
Severity,
Verdict,
)
# ---------------------------------------------------------------------------
# Stubs
# ---------------------------------------------------------------------------
class StubSettingsProvider:
def __call__(self) -> AppSettings:
from conftest import build_app_settings
return build_app_settings()
class StubLedgerAgent:
"""Stub that returns canned responses without touching any model."""
def __init__(
self,
requisition: Requisition | None = None,
verdict: Verdict | None = None,
) -> None:
self._requisition = requisition or _stub_requisition()
self._verdict = verdict or _stub_verdict()
self.examine_calls: list[FolioManifest] = []
self.audit_calls: list[tuple[Evidence, Decimal]] = []
async def examine(self, manifest: FolioManifest) -> Requisition:
self.examine_calls.append(manifest)
return self._requisition
async def audit(self, evidence: Evidence, tolerance: Decimal = Decimal("0.01")) -> Verdict:
self.audit_calls.append((evidence, tolerance))
return self._verdict
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _stub_requisition() -> Requisition:
return Requisition(
need_text=[0, 2],
need_tables=[0],
need_ocr=[1],
rationale="Page 1 is image-only; pages 0 and 2 have financial text.",
)
def _stub_verdict(
clean: bool = True,
discrepancies: list[Discrepancy] | None = None,
) -> Verdict:
return Verdict(
session_id="test-session",
discrepancies=discrepancies or [],
pages_examined=[0, 2],
rounds_taken=2,
summary="No errors found." if clean else "1 tally error found.",
clean=clean,
)
def _manifest_body(**overrides: object) -> dict[str, object]:
base: dict[str, object] = {
"sessionId": "test-session",
"pageCount": 3,
"folioTypes": ["text", "image", "mixed"],
"round": 1,
}
return {**base, **overrides}
def _evidence_body(**overrides: object) -> dict[str, object]:
base: dict[str, object] = {
"sessionId": "test-session",
"folios": [
{"page": 0, "text": "Fee: £100\nTax: £20\nTotal: £120"},
{"page": 2, "text": "Summary: all tallies correct"},
],
"round": 2,
"finalRound": False,
}
return {**base, **overrides}
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture
def stub_agent() -> StubLedgerAgent:
return StubLedgerAgent()
@pytest.fixture
def client(stub_agent: StubLedgerAgent) -> Iterator[TestClient]:
app.dependency_overrides[load_settings] = StubSettingsProvider()
app.dependency_overrides[get_math_auditor_agent] = lambda: stub_agent
yield TestClient(app, raise_server_exceptions=False)
app.dependency_overrides.pop(load_settings, None)
app.dependency_overrides.pop(get_math_auditor_agent, None)
# ---------------------------------------------------------------------------
# POST /api/v1/ai/math-auditor-agent/examine
# ---------------------------------------------------------------------------
class TestExamineEndpoint:
"""Tests for POST /api/v1/ai/math-auditor-agent/examine."""
def test_returns_200(self, client: TestClient) -> None:
resp = client.post("/api/v1/ai/math-auditor-agent/examine", json=_manifest_body())
assert resp.status_code == 200
def test_response_is_requisition(self, client: TestClient) -> None:
resp = client.post("/api/v1/ai/math-auditor-agent/examine", json=_manifest_body())
body = resp.json()
assert body["type"] == "requisition"
assert body["needText"] == [0, 2]
assert body["needTables"] == [0]
assert body["needOcr"] == [1]
assert "rationale" in body
def test_examine_called_with_parsed_manifest(
self,
client: TestClient,
stub_agent: StubLedgerAgent,
) -> None:
client.post("/api/v1/ai/math-auditor-agent/examine", json=_manifest_body(sessionId="my-session", pageCount=3))
assert len(stub_agent.examine_calls) == 1
manifest = stub_agent.examine_calls[0]
assert manifest.session_id == "my-session"
assert manifest.page_count == 3
def test_content_type_is_json(self, client: TestClient) -> None:
resp = client.post("/api/v1/ai/math-auditor-agent/examine", json=_manifest_body())
assert "application/json" in resp.headers["content-type"]
# ---------------------------------------------------------------------------
# POST /api/v1/ai/math-auditor-agent/deliberate
# ---------------------------------------------------------------------------
class TestDeliberateEndpoint:
"""Tests for POST /api/v1/ai/math-auditor-agent/deliberate."""
def test_returns_200_clean(self, client: TestClient) -> None:
resp = client.post("/api/v1/ai/math-auditor-agent/deliberate", json=_evidence_body())
assert resp.status_code == 200
def test_response_is_verdict(self, client: TestClient) -> None:
resp = client.post("/api/v1/ai/math-auditor-agent/deliberate", json=_evidence_body())
body = resp.json()
assert body["type"] == "verdict"
assert body["clean"] is True
def test_discrepancies_serialised(self, client: TestClient) -> None:
d = Discrepancy(
page=0,
kind=DiscrepancyKind.TALLY,
severity=Severity.ERROR,
description="Column total wrong",
stated="250",
expected="300",
)
stub = StubLedgerAgent(verdict=_stub_verdict(clean=False, discrepancies=[d]))
app.dependency_overrides[get_math_auditor_agent] = lambda: stub
resp = client.post("/api/v1/ai/math-auditor-agent/deliberate", json=_evidence_body())
body = resp.json()
discrepancies = body["discrepancies"]
assert len(discrepancies) == 1
assert discrepancies[0]["kind"] == "tally"
assert discrepancies[0]["severity"] == "error"
assert discrepancies[0]["stated"] == "250"
assert discrepancies[0]["expected"] == "300"
def test_tolerance_query_param_forwarded(
self,
client: TestClient,
stub_agent: StubLedgerAgent,
) -> None:
client.post("/api/v1/ai/math-auditor-agent/deliberate?tolerance=0.05", json=_evidence_body())
assert len(stub_agent.audit_calls) == 1
_, tolerance = stub_agent.audit_calls[0]
assert tolerance == Decimal("0.05")
def test_default_tolerance_when_omitted(
self,
client: TestClient,
stub_agent: StubLedgerAgent,
) -> None:
client.post("/api/v1/ai/math-auditor-agent/deliberate", json=_evidence_body())
_, tolerance = stub_agent.audit_calls[0]
assert tolerance == Decimal("0.01")
def test_invalid_tolerance_returns_400(
self,
client: TestClient,
stub_agent: StubLedgerAgent,
) -> None:
resp = client.post("/api/v1/ai/math-auditor-agent/deliberate?tolerance=notanumber", json=_evidence_body())
assert resp.status_code == 400
def test_final_round_flag_parsed(
self,
client: TestClient,
stub_agent: StubLedgerAgent,
) -> None:
client.post("/api/v1/ai/math-auditor-agent/deliberate", json=_evidence_body(finalRound=True))
evidence, _ = stub_agent.audit_calls[0]
assert evidence.final_round is True
engine/tests/test_pdf_edit_agent.py
+26 -31
View File
@@ -5,7 +5,6 @@ from dataclasses import dataclass
import pytest
from stirling.agents import PdfEditAgent, PdfEditParameterSelector, PdfEditPlanSelection
from stirling.config import AppSettings
from stirling.contracts import (
EditCannotDoResponse,
EditClarificationRequest,
@@ -13,23 +12,14 @@ from stirling.contracts import (
PdfEditRequest,
ToolOperationStep,
)
from stirling.models.tool_models import CompressParams, OperationId, RotateParams
from stirling.services import build_runtime
def build_test_settings() -> AppSettings:
return AppSettings(
smart_model_name="test",
fast_model_name="test",
smart_model_max_tokens=8192,
fast_model_max_tokens=2048,
)
from stirling.models.tool_models import Angle, FlattenParams, RotatePdfParams, ToolEndpoint
from stirling.services.runtime import AppRuntime
@dataclass(frozen=True)
class ParameterSelectorCall:
request: PdfEditRequest
operation_plan: list[OperationId]
operation_plan: list[ToolEndpoint]
operation_index: int
generated_steps: list[ToolOperationStep]
@@ -41,10 +31,10 @@ class RecordingParameterSelector:
async def select(
self,
request: PdfEditRequest,
operation_plan: list[OperationId],
operation_plan: list[ToolEndpoint],
operation_index: int,
generated_steps: list[ToolOperationStep],
) -> RotateParams | CompressParams:
) -> RotatePdfParams | FlattenParams:
self.calls.append(
ParameterSelectorCall(
request=request,
@@ -54,17 +44,18 @@ class RecordingParameterSelector:
)
)
if operation_index == 0:
return RotateParams(angle=90)
return CompressParams(compression_level=5)
return RotatePdfParams(angle=Angle(90))
return FlattenParams(flatten_only_forms=False, render_dpi=None)
class StubPdfEditAgent(PdfEditAgent):
def __init__(
self,
runtime: AppRuntime,
selection: PdfEditPlanSelection | EditClarificationRequest | EditCannotDoResponse,
parameter_selector: RecordingParameterSelector | PdfEditParameterSelector | None = None,
) -> None:
super().__init__(build_runtime(build_test_settings()))
super().__init__(runtime)
self.selection = selection
if parameter_selector is not None:
self.parameter_selector = parameter_selector
@@ -77,11 +68,12 @@ class StubPdfEditAgent(PdfEditAgent):
@pytest.mark.anyio
async def test_pdf_edit_agent_builds_multi_step_plan() -> None:
async def test_pdf_edit_agent_builds_multi_step_plan(runtime: AppRuntime) -> None:
parameter_selector = RecordingParameterSelector()
agent = StubPdfEditAgent(
runtime,
PdfEditPlanSelection(
operations=[OperationId.ROTATE, OperationId.COMPRESS],
operations=[ToolEndpoint.ROTATE_PDF, ToolEndpoint.FLATTEN],
summary="Rotate the PDF, then compress it.",
rationale="The pages need reorientation before reducing file size.",
),
@@ -98,17 +90,18 @@ async def test_pdf_edit_agent_builds_multi_step_plan() -> None:
assert isinstance(response, EditPlanResponse)
assert response.summary == "Rotate the PDF, then compress it."
assert response.rationale == "The pages need reorientation before reducing file size."
assert [step.tool for step in response.steps] == [OperationId.ROTATE, OperationId.COMPRESS]
assert isinstance(response.steps[0].parameters, RotateParams)
assert isinstance(response.steps[1].parameters, CompressParams)
assert [step.tool for step in response.steps] == [ToolEndpoint.ROTATE_PDF, ToolEndpoint.FLATTEN]
assert isinstance(response.steps[0].parameters, RotatePdfParams)
assert isinstance(response.steps[1].parameters, FlattenParams)
@pytest.mark.anyio
async def test_pdf_edit_agent_passes_previous_steps_to_parameter_selector() -> None:
async def test_pdf_edit_agent_passes_previous_steps_to_parameter_selector(runtime: AppRuntime) -> None:
parameter_selector = RecordingParameterSelector()
agent = StubPdfEditAgent(
runtime,
PdfEditPlanSelection(
operations=[OperationId.ROTATE, OperationId.COMPRESS],
operations=[ToolEndpoint.ROTATE_PDF, ToolEndpoint.FLATTEN],
summary="Rotate the PDF, then compress it.",
),
parameter_selector=parameter_selector,
@@ -127,19 +120,20 @@ async def test_pdf_edit_agent_passes_previous_steps_to_parameter_selector() -> N
assert parameter_selector.calls[1].operation_index == 1
assert parameter_selector.calls[1].generated_steps == [
ToolOperationStep(
tool=OperationId.ROTATE,
parameters=RotateParams(angle=90),
tool=ToolEndpoint.ROTATE_PDF,
parameters=RotatePdfParams(angle=Angle(90)),
)
]
@pytest.mark.anyio
async def test_pdf_edit_agent_returns_clarification_without_partial_plan() -> None:
async def test_pdf_edit_agent_returns_clarification_without_partial_plan(runtime: AppRuntime) -> None:
agent = StubPdfEditAgent(
runtime,
EditClarificationRequest(
question="Which pages should be rotated?",
reason="The request does not say which pages to change.",
)
),
)
response = await agent.handle(PdfEditRequest(user_message="Rotate some pages."))
@@ -148,11 +142,12 @@ async def test_pdf_edit_agent_returns_clarification_without_partial_plan() -> No
@pytest.mark.anyio
async def test_pdf_edit_agent_returns_cannot_do_without_partial_plan() -> None:
async def test_pdf_edit_agent_returns_cannot_do_without_partial_plan(runtime: AppRuntime) -> None:
agent = StubPdfEditAgent(
runtime,
EditCannotDoResponse(
reason="This request requires OCR, which is not part of PDF edit planning.",
)
),
)
response = await agent.handle(PdfEditRequest(user_message="Read this scan and summarize it."))
engine/tests/test_pdf_question_agent.py
+31 -24
View File
@@ -3,19 +3,20 @@ from __future__ import annotations
import pytest
from stirling.agents import PdfQuestionAgent
from stirling.config import AppSettings
from stirling.contracts import (
ExtractedFileText,
PdfQuestionAnswerResponse,
PdfQuestionNeedTextResponse,
PdfQuestionNeedContentResponse,
PdfQuestionNotFoundResponse,
PdfQuestionRequest,
PdfTextSelection,
)
from stirling.services import build_runtime
from stirling.services.runtime import AppRuntime
class StubPdfQuestionAgent(PdfQuestionAgent):
def __init__(self, response: PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse) -> None:
super().__init__(build_runtime(build_test_settings()))
def __init__(self, runtime: AppRuntime, response: PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse) -> None:
super().__init__(runtime)
self.response = response
async def _run_answer_agent(
@@ -25,38 +26,39 @@ class StubPdfQuestionAgent(PdfQuestionAgent):
return self.response
def build_test_settings() -> AppSettings:
return AppSettings(
smart_model_name="test",
fast_model_name="test",
smart_model_max_tokens=8192,
fast_model_max_tokens=2048,
def invoice_page() -> ExtractedFileText:
return ExtractedFileText(
file_name="invoice.pdf",
pages=[PdfTextSelection(page_number=1, text="Invoice total: 120.00")],
)
@pytest.mark.anyio
async def test_pdf_question_agent_requires_extracted_text() -> None:
agent = PdfQuestionAgent(build_runtime(build_test_settings()))
async def test_pdf_question_agent_requires_extracted_text(runtime: AppRuntime) -> None:
agent = PdfQuestionAgent(runtime)
response = await agent.handle(PdfQuestionRequest(question="What is the total?", extracted_text=""))
response = await agent.handle(
PdfQuestionRequest(question="What is the total?", page_text=[], file_names=["test.pdf"])
)
assert isinstance(response, PdfQuestionNeedTextResponse)
assert isinstance(response, PdfQuestionNeedContentResponse)
@pytest.mark.anyio
async def test_pdf_question_agent_returns_grounded_answer() -> None:
async def test_pdf_question_agent_returns_grounded_answer(runtime: AppRuntime) -> None:
agent = StubPdfQuestionAgent(
runtime,
PdfQuestionAnswerResponse(
answer="The invoice total is 120.00.",
evidence=["Invoice total: 120.00"],
)
evidence=[invoice_page()],
),
)
response = await agent.handle(
PdfQuestionRequest(
question="What is the total?",
extracted_text="Invoice total: 120.00",
file_name="invoice.pdf",
page_text=[invoice_page()],
file_names=["invoice.pdf"],
)
)
@@ -65,14 +67,19 @@ async def test_pdf_question_agent_returns_grounded_answer() -> None:
@pytest.mark.anyio
async def test_pdf_question_agent_returns_not_found_when_text_is_insufficient() -> None:
agent = StubPdfQuestionAgent(PdfQuestionNotFoundResponse(reason="The answer is not present in the text."))
async def test_pdf_question_agent_returns_not_found_when_text_is_insufficient(runtime: AppRuntime) -> None:
agent = StubPdfQuestionAgent(runtime, PdfQuestionNotFoundResponse(reason="The answer is not present in the text."))
response = await agent.handle(
PdfQuestionRequest(
question="What is the total?",
extracted_text="This page contains only a shipping address.",
file_name="invoice.pdf",
page_text=[
ExtractedFileText(
file_name="invoice.pdf",
pages=[PdfTextSelection(page_number=1, text="This page contains only a shipping address.")],
)
],
file_names=["invoice.pdf"],
)
)
engine/tests/test_rag.py
+283
View File
@@ -0,0 +1,283 @@
from __future__ import annotations
import pytest
from stirling.rag.capability import RagCapability
from stirling.rag.chunker import chunk_text
from stirling.rag.service import RagService
from stirling.rag.sqlite_vec_store import SqliteVecStore
from stirling.rag.store import Document, SearchResult
# ── chunk_text ──────────────────────────────────────────────────────────
class TestChunkText:
def test_empty_input_returns_empty(self) -> None:
assert chunk_text("") == []
assert chunk_text(" ") == []
def test_short_text_returns_single_chunk(self) -> None:
text = "Hello world."
chunks = chunk_text(text, chunk_size=100)
assert len(chunks) == 1
assert chunks[0] == "Hello world."
def test_splits_on_paragraph_boundaries(self) -> None:
text = "First paragraph.\n\nSecond paragraph.\n\nThird paragraph."
chunks = chunk_text(text, chunk_size=30, overlap=0)
# Each paragraph fits in 30 chars, so they should be split
assert len(chunks) >= 2
assert "First paragraph." in chunks[0]
def test_long_text_produces_multiple_chunks(self) -> None:
text = " ".join(["word"] * 200)
chunks = chunk_text(text, chunk_size=100, overlap=10)
assert len(chunks) > 1
for chunk in chunks:
# Chunks may slightly exceed due to sentence boundary snapping
assert len(chunk) <= 200 # generous upper bound
def test_overlap_produces_shared_content(self) -> None:
sentences = [f"Sentence number {i}." for i in range(20)]
text = " ".join(sentences)
chunks = chunk_text(text, chunk_size=100, overlap=30)
if len(chunks) >= 2:
# After word-boundary snapping, the second chunk should share
# some content with the tail of the first chunk
words_in_first_tail = chunks[0].split()[-3:] # last 3 words
overlap_text = " ".join(words_in_first_tail)
assert overlap_text in chunks[1], f"Expected overlap '{overlap_text}' in chunk[1]: '{chunks[1][:80]}...'"
# ── SqliteVecStore ──────────────────────────────────────────────────────
class TestSqliteVecStore:
"""Each test gets its own ephemeral store to avoid cross-test dimension conflicts."""
@pytest.mark.anyio
async def test_add_and_search(self) -> None:
store = SqliteVecStore.ephemeral()
docs = [
Document(id="1", text="Python is a programming language", metadata={"source": "test"}),
Document(id="2", text="Java is another programming language", metadata={"source": "test"}),
Document(id="3", text="The weather today is sunny", metadata={"source": "test"}),
]
# Simple 3-dimensional embeddings for testing
embeddings = [
[1.0, 0.0, 0.0],
[0.9, 0.1, 0.0],
[0.0, 0.0, 1.0],
]
await store.add_documents("test-col", docs, embeddings)
# Search with a query close to the programming-related docs
results = await store.search("test-col", [1.0, 0.05, 0.0], top_k=2)
assert len(results) == 2
assert isinstance(results[0], SearchResult)
# The closest should be doc "1" (exact match on first dimension)
assert results[0].document.id == "1"
assert results[0].score > 0.5
@pytest.mark.anyio
async def test_list_and_has_collection(self) -> None:
store = SqliteVecStore.ephemeral()
docs = [Document(id="1", text="test", metadata={})]
await store.add_documents("my-collection", docs, [[1.0, 0.0]])
collections = await store.list_collections()
assert "my-collection" in collections
assert await store.has_collection("my-collection") is True
assert await store.has_collection("nonexistent") is False
@pytest.mark.anyio
async def test_delete_collection(self) -> None:
store = SqliteVecStore.ephemeral()
docs = [Document(id="1", text="test", metadata={})]
await store.add_documents("to-delete", docs, [[1.0]])
assert await store.has_collection("to-delete") is True
await store.delete_collection("to-delete")
assert await store.has_collection("to-delete") is False
@pytest.mark.anyio
async def test_search_empty_collection(self) -> None:
store = SqliteVecStore.ephemeral()
docs = [Document(id="1", text="test", metadata={})]
await store.add_documents("empty-test", docs, [[1.0, 0.0]])
results = await store.search("empty-test", [1.0, 0.0], top_k=5)
assert len(results) == 1
@pytest.mark.anyio
async def test_mismatched_docs_embeddings_raises(self) -> None:
store = SqliteVecStore.ephemeral()
docs = [Document(id="1", text="test", metadata={})]
with pytest.raises(ValueError, match="documents.*embeddings"):
await store.add_documents("bad", docs, [[1.0], [2.0]])
# ── RagService (with stub embedder) ────────────────────────────────────
class StubEmbeddingService:
"""A minimal stub that returns fixed-dimension embeddings for testing."""
def __init__(self, dim: int = 8) -> None:
self._dim = dim
async def embed_query(self, text: str) -> list[float]:
# Deterministic embedding based on hash of text
h = hash(text) % 1000
return [(h + i) / 1000.0 for i in range(self._dim)]
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
return [await self.embed_query(t) for t in texts]
def chunk_and_prepare(
self,
text: str,
source: str = "",
base_metadata: dict[str, str] | None = None,
) -> list[Document]:
from stirling.rag.chunker import chunk_text
chunks = chunk_text(text, 100, 10)
docs = []
for i, chunk in enumerate(chunks):
meta = dict(base_metadata) if base_metadata else {}
meta["source"] = source
meta["chunk_index"] = str(i)
doc_id = f"{source}:chunk:{i}" if source else f"chunk:{i}"
docs.append(Document(id=doc_id, text=chunk, metadata=meta))
return docs
@pytest.fixture
def rag_service() -> RagService:
"""Each RagService test gets its own fresh ephemeral store to avoid dimension conflicts."""
store = SqliteVecStore.ephemeral()
return RagService(embedder=StubEmbeddingService(), store=store, default_top_k=3) # type: ignore[arg-type]
class TestRagService:
@pytest.mark.anyio
async def test_index_and_search(self, rag_service: RagService) -> None:
text = "Python is great for data science. It has many libraries like pandas and numpy."
count = await rag_service.index_text("docs", text, source="guide.pdf")
assert count > 0
results = await rag_service.search("Python libraries", collection="docs")
assert len(results) > 0
assert results[0].document.text # non-empty text
@pytest.mark.anyio
async def test_index_empty_text_returns_zero(self, rag_service: RagService) -> None:
count = await rag_service.index_text("docs", "", source="empty.pdf")
assert count == 0
@pytest.mark.anyio
async def test_search_nonexistent_collection_returns_empty(self, rag_service: RagService) -> None:
results = await rag_service.search("anything", collection="nonexistent")
assert results == []
@pytest.mark.anyio
async def test_search_all_collections(self, rag_service: RagService) -> None:
await rag_service.index_text("col-a", "Machine learning overview.", source="ml.pdf")
await rag_service.index_text("col-b", "Deep learning with neural networks.", source="dl.pdf")
results = await rag_service.search("neural networks")
assert len(results) > 0
@pytest.mark.anyio
async def test_delete_collection(self, rag_service: RagService) -> None:
await rag_service.index_text("temp", "Temporary data.", source="tmp.pdf")
collections = await rag_service.list_collections()
assert "temp" in collections
await rag_service.delete_collection("temp")
collections = await rag_service.list_collections()
assert "temp" not in collections
# ── RagCapability ──────────────────────────────────────────────────────
async def _invoke_search_knowledge(capability: RagCapability, query: str, max_results: int = 5) -> str:
"""Extract and call the search_knowledge tool function from a RagCapability's toolset."""
from pydantic_ai import FunctionToolset
toolset = capability.toolset
assert isinstance(toolset, FunctionToolset)
tool = toolset.tools["search_knowledge"]
return await tool.function(query=query, max_results=max_results) # type: ignore[call-arg] — pyright can't infer the generic tool function's kwargs
class TestRagCapability:
def test_instructions_static_when_collections_pinned(self, rag_service: RagService) -> None:
cap = RagCapability(rag_service, collections=["docs", "manuals"])
instructions = cap.instructions
assert isinstance(instructions, str)
assert "docs, manuals" in instructions
assert "search_knowledge" in instructions
def test_instructions_dynamic_when_no_collections(self, rag_service: RagService) -> None:
cap = RagCapability(rag_service)
instructions = cap.instructions
assert callable(instructions)
@pytest.mark.anyio
async def test_dynamic_instructions_list_available_collections(self, rag_service: RagService) -> None:
await rag_service.index_text("col-a", "Alpha content.", source="a.pdf")
await rag_service.index_text("col-b", "Beta content.", source="b.pdf")
cap = RagCapability(rag_service)
instructions_fn = cap.instructions
assert callable(instructions_fn)
text = await instructions_fn()
assert "col-a" in text
assert "col-b" in text
@pytest.mark.anyio
async def test_dynamic_instructions_when_store_empty(self, rag_service: RagService) -> None:
cap = RagCapability(rag_service)
instructions_fn = cap.instructions
assert callable(instructions_fn)
text = await instructions_fn()
assert "empty" in text.lower()
@pytest.mark.anyio
async def test_search_knowledge_returns_no_results_message_when_empty(self, rag_service: RagService) -> None:
cap = RagCapability(rag_service)
output = await _invoke_search_knowledge(cap, "anything")
assert output == "No relevant results found in the knowledge base."
@pytest.mark.anyio
async def test_search_knowledge_formats_results_with_source_and_score(self, rag_service: RagService) -> None:
await rag_service.index_text("docs", "Python is a programming language.", source="guide.pdf")
cap = RagCapability(rag_service)
output = await _invoke_search_knowledge(cap, "Python")
assert "[Result 1" in output
assert "source: guide.pdf" in output
assert "chunk:" in output
assert "relevance:" in output
@pytest.mark.anyio
async def test_search_knowledge_restricts_to_pinned_collections(self, rag_service: RagService) -> None:
await rag_service.index_text("pinned", "Pinned collection content.", source="pinned.pdf")
await rag_service.index_text("other", "Content in another collection.", source="other.pdf")
cap = RagCapability(rag_service, collections=["pinned"])
output = await _invoke_search_knowledge(cap, "content")
assert "pinned.pdf" in output
assert "other.pdf" not in output
@pytest.mark.anyio
async def test_search_knowledge_respects_max_results(self, rag_service: RagService) -> None:
paragraphs = "\n\n".join(f"Paragraph {i} about topic." for i in range(10))
await rag_service.index_text("bulk", paragraphs, source="bulk.pdf")
cap = RagCapability(rag_service)
output = await _invoke_search_knowledge(cap, "topic", max_results=2)
# Only two results requested, so only Result 1 and Result 2 should appear
assert "[Result 1" in output
assert "[Result 2" in output
assert "[Result 3" not in output
engine/tests/test_rag_routes.py
+205
View File
@@ -0,0 +1,205 @@
from __future__ import annotations
from collections.abc import Iterator
import pytest
from fastapi.testclient import TestClient
from stirling.api import app
from stirling.api.dependencies import get_rag_embedding_model, get_rag_service
from stirling.rag import Document, RagService, SqliteVecStore
TEST_EMBEDDING_MODEL = "test-embedder"
class StubEmbedder:
"""Deterministic embeddings for route tests — no network, no provider needed."""
def __init__(self, dim: int = 8) -> None:
self._dim = dim
async def embed_query(self, text: str) -> list[float]:
h = hash(text) % 1000
return [(h + i) / 1000.0 for i in range(self._dim)]
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
return [await self.embed_query(t) for t in texts]
def chunk_and_prepare(
self,
text: str,
source: str = "",
base_metadata: dict[str, str] | None = None,
) -> list[Document]:
from stirling.rag.chunker import chunk_text
chunks = chunk_text(text, 100, 10)
docs = []
for i, chunk in enumerate(chunks):
meta = dict(base_metadata) if base_metadata else {}
meta["source"] = source
meta["chunk_index"] = str(i)
doc_id = f"{source}:chunk:{i}" if source else f"chunk:{i}"
docs.append(Document(id=doc_id, text=chunk, metadata=meta))
return docs
def _build_service() -> RagService:
return RagService(
embedder=StubEmbedder(), # type: ignore[arg-type]
store=SqliteVecStore.ephemeral(),
default_top_k=3,
)
@pytest.fixture
def client() -> Iterator[TestClient]:
service = _build_service()
app.dependency_overrides[get_rag_service] = lambda: service
app.dependency_overrides[get_rag_embedding_model] = lambda: TEST_EMBEDDING_MODEL
try:
yield TestClient(app)
finally:
app.dependency_overrides.pop(get_rag_service, None)
app.dependency_overrides.pop(get_rag_embedding_model, None)
# ── /status ─────────────────────────────────────────────────────────────
def test_status_reports_embedding_model_and_collections(client: TestClient) -> None:
client.post(
"/api/v1/rag/index",
json={"collection": "my-docs", "text": "Hello world.", "source": "a.pdf"},
)
response = client.get("/api/v1/rag/status")
assert response.status_code == 200
body = response.json()
assert body["embeddingModel"] == TEST_EMBEDDING_MODEL
assert "my-docs" in body["collections"]
def test_status_when_empty(client: TestClient) -> None:
response = client.get("/api/v1/rag/status")
assert response.status_code == 200
body = response.json()
assert body == {"embeddingModel": TEST_EMBEDDING_MODEL, "collections": []}
# ── /index ──────────────────────────────────────────────────────────────
def test_index_returns_chunk_count(client: TestClient) -> None:
response = client.post(
"/api/v1/rag/index",
json={"collection": "indexed", "text": "Short text.", "source": "doc.pdf"},
)
assert response.status_code == 200
body = response.json()
assert body["collection"] == "indexed"
assert body["chunksIndexed"] >= 1
def test_index_rejects_empty_collection_name(client: TestClient) -> None:
response = client.post(
"/api/v1/rag/index",
json={"collection": "", "text": "Text.", "source": "x.pdf"},
)
assert response.status_code == 422
def test_index_rejects_oversized_text(client: TestClient) -> None:
huge = "x" * 1_000_001 # Just over the 1MB cap
response = client.post(
"/api/v1/rag/index",
json={"collection": "toobig", "text": huge},
)
assert response.status_code == 422
# ── /search ─────────────────────────────────────────────────────────────
def test_search_returns_results(client: TestClient) -> None:
client.post(
"/api/v1/rag/index",
json={"collection": "search-test", "text": "Python is fun.", "source": "guide.pdf"},
)
response = client.post(
"/api/v1/rag/search",
json={"query": "Python", "collection": "search-test", "topK": 3},
)
assert response.status_code == 200
body = response.json()
assert body["query"] == "Python"
assert len(body["results"]) >= 1
first = body["results"][0]
assert first["source"] == "guide.pdf"
assert "score" in first
def test_search_rejects_empty_collection_name(client: TestClient) -> None:
response = client.post(
"/api/v1/rag/search",
json={"query": "anything", "collection": ""},
)
assert response.status_code == 422
def test_search_without_collection_searches_all(client: TestClient) -> None:
client.post(
"/api/v1/rag/index",
json={"collection": "col-one", "text": "Alpha content.", "source": "one.pdf"},
)
client.post(
"/api/v1/rag/index",
json={"collection": "col-two", "text": "Beta content.", "source": "two.pdf"},
)
response = client.post(
"/api/v1/rag/search",
json={"query": "content"},
)
assert response.status_code == 200
body = response.json()
assert len(body["results"]) >= 1
# ── /collections ────────────────────────────────────────────────────────
def test_collections_empty_when_no_data(client: TestClient) -> None:
response = client.get("/api/v1/rag/collections")
assert response.status_code == 200
assert response.json() == {"collections": []}
def test_collections_lists_indexed(client: TestClient) -> None:
client.post(
"/api/v1/rag/index",
json={"collection": "list-me", "text": "Text.", "source": "x.pdf"},
)
response = client.get("/api/v1/rag/collections")
assert response.status_code == 200
assert "list-me" in response.json()["collections"]
# ── DELETE /collections/{name} ──────────────────────────────────────────
def test_delete_collection_removes_it(client: TestClient) -> None:
client.post(
"/api/v1/rag/index",
json={"collection": "to-delete", "text": "Text.", "source": "x.pdf"},
)
response = client.delete("/api/v1/rag/collections/to-delete")
assert response.status_code == 200
assert response.json() == {"status": "deleted", "collection": "to-delete"}
listing = client.get("/api/v1/rag/collections").json()
assert "to-delete" not in listing["collections"]
def test_delete_nonexistent_collection_is_idempotent(client: TestClient) -> None:
response = client.delete("/api/v1/rag/collections/never-existed")
assert response.status_code == 200
assert response.json() == {"status": "deleted", "collection": "never-existed"}
engine/tests/test_stirling_api.py
+29 -38
View File
@@ -1,3 +1,4 @@
from conftest import build_app_settings
from fastapi.testclient import TestClient
from stirling.api import app
@@ -11,7 +12,7 @@ from stirling.api.dependencies import (
get_pdf_question_agent,
get_user_spec_agent,
)
from stirling.config import AppSettings, load_settings
from stirling.config import load_settings
from stirling.contracts import (
AgentDraft,
AgentDraftRequest,
@@ -19,42 +20,30 @@ from stirling.contracts import (
AgentExecutionRequest,
AgentRevisionRequest,
AgentRevisionResponse,
CannotContinueExecutionAction,
EditCannotDoResponse,
OrchestratorRequest,
PdfEditRequest,
PdfQuestionNotFoundResponse,
PdfQuestionRequest,
UnsupportedCapabilityResponse,
)
from stirling.contracts.form_fill import (
AnalysedFileResult,
CannotContinueExecutionAction,
CrossFileRole,
DetectedRole,
DocumentExtractionRequest,
EditCannotDoResponse,
FileFillResult,
FormAnalysisRequest,
FormAnalysisResponse,
FormFillBatchRequest,
FormFillBatchResponse,
KnowledgeUpdateResponse,
OrchestratorRequest,
PdfEditRequest,
PdfQuestionNeedContentResponse,
PdfQuestionNotFoundResponse,
PdfQuestionRequest,
)
from stirling.models.tool_models import RotateParams
class StubSettingsProvider:
def __call__(self) -> AppSettings:
return AppSettings(
smart_model_name="test",
fast_model_name="test",
smart_model_max_tokens=8192,
fast_model_max_tokens=2048,
)
from stirling.models.tool_models import Angle, RotatePdfParams
class StubOrchestratorAgent:
async def handle(self, request: OrchestratorRequest) -> UnsupportedCapabilityResponse:
return UnsupportedCapabilityResponse(capability="pdf_edit", message=request.user_message)
async def handle(self, request: OrchestratorRequest) -> PdfQuestionNeedContentResponse:
return PdfQuestionNeedContentResponse(reason=request.user_message, files=[], max_pages=1, max_characters=1000)
class StubPdfEditAgent:
@@ -127,14 +116,7 @@ class StubExecutionPlanningAgent:
return CannotContinueExecutionAction(reason=str(request.current_step_index))
client: TestClient = TestClient(app)
def override_settings() -> AppSettings:
return StubSettingsProvider()()
app.dependency_overrides[load_settings] = override_settings
app.dependency_overrides[load_settings] = build_app_settings
app.dependency_overrides[get_orchestrator_agent] = lambda: StubOrchestratorAgent()
app.dependency_overrides[get_pdf_edit_agent] = lambda: StubPdfEditAgent()
app.dependency_overrides[get_pdf_question_agent] = lambda: StubPdfQuestionAgent()
@@ -144,6 +126,8 @@ app.dependency_overrides[get_form_analyser_agent] = lambda: StubFormAnalyserAgen
app.dependency_overrides[get_form_filler_agent] = lambda: StubFormFillerAgent()
app.dependency_overrides[get_document_extractor_agent] = lambda: StubDocumentExtractorAgent()
client: TestClient = TestClient(app)
def test_health_route() -> None:
response = client.get("/health")
@@ -153,10 +137,10 @@ def test_health_route() -> None:
def test_orchestrator_route() -> None:
response = client.post("/api/v1/orchestrator", json={"userMessage": "route this"})
response = client.post("/api/v1/orchestrator", json={"userMessage": "route this", "fileNames": ["test.pdf"]})
assert response.status_code == 200
assert response.json()["outcome"] == "unsupported_capability"
assert response.json()["outcome"] == "need_content"
def test_pdf_edit_route() -> None:
@@ -167,7 +151,14 @@ def test_pdf_edit_route() -> None:
def test_pdf_questions_route() -> None:
response = client.post("/api/v1/pdf/questions", json={"question": "what is this?"})
response = client.post(
"/api/v1/pdf/questions",
json={
"question": "what is this?",
"fileNames": ["test.pdf"],
"pageText": [{"fileName": "test.pdf", "pages": [{"pageNumber": 1, "text": "Example"}]}],
},
)
assert response.status_code == 200
assert response.json()["outcome"] == "not_found"
@@ -192,8 +183,8 @@ def test_agent_revise_route() -> None:
"steps": [
{
"kind": "tool",
"tool": "rotate",
"parameters": RotateParams(angle=90).model_dump(by_alias=True),
"tool": "/api/v1/general/rotate-pdf",
"parameters": RotatePdfParams(angle=Angle(90)).model_dump(by_alias=True),
}
],
},
@@ -252,8 +243,8 @@ def test_next_action_route() -> None:
"steps": [
{
"kind": "tool",
"tool": "rotate",
"parameters": RotateParams(angle=90).model_dump(by_alias=True),
"tool": "/api/v1/general/rotate-pdf",
"parameters": RotatePdfParams(angle=Angle(90)).model_dump(by_alias=True),
}
],
},
engine/tests/test_stirling_contracts.py
+41 -23
View File
@@ -1,36 +1,47 @@
from collections.abc import Iterator
import pytest
from stirling.config import AppSettings, load_settings
from stirling.config import AppSettings
from stirling.contracts import (
AgentExecutionRequest,
AgentSpec,
AgentSpecStep,
EditPlanResponse,
ExecutionContext,
OrchestratorRequest,
PdfQuestionAnswerResponse,
ToolOperationStep,
)
from stirling.contracts.form_fill import (
ExtractedFileText,
ExtractedTextArtifact,
KnowledgeEntry,
KnowledgeUpdateResponse,
OrchestratorRequest,
PdfQuestionAnswerResponse,
PdfTextSelection,
ToolOperationStep,
)
from stirling.models.tool_models import OperationId, RotateParams
from stirling.models.tool_models import Angle, RotatePdfParams, ToolEndpoint
def test_orchestrator_request_accepts_user_message() -> None:
request = OrchestratorRequest(user_message="Rotate the PDF")
request = OrchestratorRequest(
user_message="Rotate the PDF",
file_names=["test.pdf"],
artifacts=[
ExtractedTextArtifact(
files=[
ExtractedFileText(
file_name="test.pdf",
pages=[PdfTextSelection(page_number=1, text="Hello")],
)
]
)
],
)
assert request.user_message == "Rotate the PDF"
assert len(request.artifacts) == 1
def test_agent_execution_request_uses_typed_agent_spec() -> None:
steps: list[AgentSpecStep] = [
ToolOperationStep(
tool=OperationId.ROTATE,
parameters=RotateParams(angle=90),
tool=ToolEndpoint.ROTATE_PDF,
parameters=RotatePdfParams(angle=Angle(90)),
)
]
request = AgentExecutionRequest(
@@ -48,13 +59,13 @@ def test_agent_execution_request_uses_typed_agent_spec() -> None:
def test_edit_plan_response_has_typed_steps() -> None:
steps = [ToolOperationStep(tool=OperationId.ROTATE, parameters=RotateParams(angle=90))]
steps = [ToolOperationStep(tool=ToolEndpoint.ROTATE_PDF, parameters=RotatePdfParams(angle=Angle(90)))]
response = EditPlanResponse(
summary="Rotate the input PDF by 90 degrees.",
steps=steps,
)
assert response.steps[0].tool == OperationId.ROTATE
assert response.steps[0].tool == ToolEndpoint.ROTATE_PDF
def test_pdf_question_answer_defaults_evidence_list() -> None:
@@ -63,13 +74,6 @@ def test_pdf_question_answer_defaults_evidence_list() -> None:
assert response.evidence == []
@pytest.fixture(autouse=True)
def clear_settings_cache() -> Iterator[None]:
load_settings.cache_clear()
yield
load_settings.cache_clear()
def test_knowledge_update_response_discriminator() -> None:
update = KnowledgeUpdateResponse(
proposed_entries=[KnowledgeEntry(key="name", value="John", source="CV")],
@@ -79,11 +83,25 @@ def test_knowledge_update_response_discriminator() -> None:
def test_app_settings_accepts_model_configuration() -> None:
from pathlib import Path
from stirling.config import RagBackend
settings = AppSettings(
smart_model_name="claude-sonnet-4-5-20250929",
fast_model_name="claude-haiku-4-5-20251001",
smart_model_max_tokens=8192,
fast_model_max_tokens=2048,
rag_backend=RagBackend.SQLITE,
rag_embedding_model="voyageai:voyage-4",
rag_store_path=Path(":memory:"),
rag_pgvector_dsn="",
rag_chunk_size=512,
rag_chunk_overlap=64,
rag_default_top_k=5,
posthog_enabled=False,
posthog_api_key="",
posthog_host="https://eu.i.posthog.com",
)
assert settings.smart_model_name
engine/tests/test_user_spec_agent.py
+30 -34
View File
@@ -4,7 +4,6 @@ import pytest
from pydantic import ValidationError
from stirling.agents import UserSpecAgent
from stirling.config import AppSettings
from stirling.contracts import (
AgentDraft,
AgentDraftRequest,
@@ -14,35 +13,28 @@ from stirling.contracts import (
EditPlanResponse,
ToolOperationStep,
)
from stirling.models.tool_models import CompressParams, OperationId, RotateParams
from stirling.services import build_runtime
def build_test_settings() -> AppSettings:
return AppSettings(
smart_model_name="test",
fast_model_name="test",
smart_model_max_tokens=8192,
fast_model_max_tokens=2048,
)
from stirling.models.tool_models import Angle, FlattenParams, RotatePdfParams, ToolEndpoint
from stirling.services.runtime import AppRuntime
class StubUserSpecAgent(UserSpecAgent):
def __init__(self, draft_result: AgentDraft, revision_result: AgentDraft) -> None:
super().__init__(build_runtime(build_test_settings()))
def __init__(self, runtime: AppRuntime, draft_result: AgentDraft, revision_result: AgentDraft) -> None:
super().__init__(runtime)
self.draft_result = draft_result
self.revision_result = revision_result
self.edit_plan = EditPlanResponse(
summary="Rotate the document.",
steps=[
ToolOperationStep(
tool=OperationId.ROTATE,
parameters=RotateParams(angle=90),
tool=ToolEndpoint.ROTATE_PDF,
parameters=RotatePdfParams(angle=Angle(90)),
)
],
)
async def _build_edit_plan(self, user_message: str) -> EditPlanResponse:
async def _build_edit_plan(
self, user_message: str, conversation_history: list[ConversationMessage]
) -> EditPlanResponse:
return self.edit_plan
async def _run_draft_agent(self, request: AgentDraftRequest, edit_plan: EditPlanResponse) -> AgentDraft:
@@ -53,10 +45,12 @@ class StubUserSpecAgent(UserSpecAgent):
class ClarifyingUserSpecAgent(UserSpecAgent):
def __init__(self) -> None:
super().__init__(build_runtime(build_test_settings()))
def __init__(self, runtime: AppRuntime) -> None:
super().__init__(runtime)
async def _build_edit_plan(self, user_message: str) -> EditClarificationRequest:
async def _build_edit_plan(
self, user_message: str, conversation_history: list[ConversationMessage]
) -> EditClarificationRequest:
return EditClarificationRequest(
question="Which pages should be changed?",
reason="The request does not specify the target pages.",
@@ -64,16 +58,17 @@ class ClarifyingUserSpecAgent(UserSpecAgent):
@pytest.mark.anyio
async def test_user_spec_agent_drafts_agent_spec() -> None:
async def test_user_spec_agent_drafts_agent_spec(runtime: AppRuntime) -> None:
agent = StubUserSpecAgent(
runtime,
AgentDraft(
name="Invoice Cleanup",
description="Prepare invoices for review.",
objective="Normalize invoices before accounting review.",
steps=[
ToolOperationStep(
tool=OperationId.ROTATE,
parameters=RotateParams(angle=90),
tool=ToolEndpoint.ROTATE_PDF,
parameters=RotatePdfParams(angle=Angle(90)),
)
],
),
@@ -100,19 +95,20 @@ async def test_user_spec_agent_drafts_agent_spec() -> None:
@pytest.mark.anyio
async def test_user_spec_agent_revises_existing_draft() -> None:
async def test_user_spec_agent_revises_existing_draft(runtime: AppRuntime) -> None:
current_draft = AgentDraft(
name="Invoice Cleanup",
description="Prepare invoices for review.",
objective="Normalize invoices before accounting review.",
steps=[
ToolOperationStep(
tool=OperationId.ROTATE,
parameters=RotateParams(angle=90),
tool=ToolEndpoint.ROTATE_PDF,
parameters=RotatePdfParams(angle=Angle(90)),
)
],
)
agent = StubUserSpecAgent(
runtime,
draft_result=current_draft,
revision_result=AgentDraft(
name="Invoice Cleanup",
@@ -120,12 +116,12 @@ async def test_user_spec_agent_revises_existing_draft() -> None:
objective="Normalize invoices before accounting review.",
steps=[
ToolOperationStep(
tool=OperationId.ROTATE,
parameters=RotateParams(angle=90),
tool=ToolEndpoint.ROTATE_PDF,
parameters=RotatePdfParams(angle=Angle(90)),
),
ToolOperationStep(
tool=OperationId.COMPRESS,
parameters=CompressParams(compression_level=5),
tool=ToolEndpoint.FLATTEN,
parameters=FlattenParams(flatten_only_forms=False, render_dpi=None),
),
],
),
@@ -146,14 +142,14 @@ async def test_user_spec_agent_revises_existing_draft() -> None:
def test_tool_operation_step_rejects_mismatched_parameters() -> None:
with pytest.raises(ValidationError):
ToolOperationStep(
tool=OperationId.ROTATE,
parameters=CompressParams(compression_level=5),
tool=ToolEndpoint.ROTATE_PDF,
parameters=FlattenParams(flatten_only_forms=False, render_dpi=None),
)
@pytest.mark.anyio
async def test_user_spec_agent_propagates_edit_clarification() -> None:
agent = ClarifyingUserSpecAgent()
async def test_user_spec_agent_propagates_edit_clarification(runtime: AppRuntime) -> None:
agent = ClarifyingUserSpecAgent(runtime)
response = await agent.draft(AgentDraftRequest(user_message="Build an agent to rotate some pages."))
engine/uv.lock
Generated
+844 -103
View File
File diff suppressed because it is too large Load Diff