mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-07-20 23:21:06 +02:00
feat(runs): introduce Run and ToolOutputSpill models for logging scraper invocations
- Added `Run` model to track scraper invocations, including metadata such as status, input, and output. - Implemented `ToolOutputSpill` model to store context-editing spills separately from user-facing logs. - Updated middleware to handle spill placeholders and integrate with the new models. - Enhanced REST API to record runs and expose run history through new endpoints. - Adjusted tests to validate the new run logging functionality and ensure proper integration with existing capabilities.
This commit is contained in:
parent
ab747e7a49
commit
b6e378b070
14 changed files with 1247 additions and 130 deletions
|
|
@ -1,16 +1,26 @@
|
|||
"""Generate the agent door from the capability registry (05).
|
||||
|
||||
One LangChain tool per verb; each runs the same thin adapter as the REST door
|
||||
(``access/rest.py``): meter-gate -> executor -> charge. The tool returns the
|
||||
verb's serialized output so the model can reason over it; UI cards are the SSE
|
||||
emission handler's job, not this generator's.
|
||||
(``access/rest.py``): meter-gate -> executor -> charge. Every run is recorded to
|
||||
the ``runs`` table (best-effort). Outputs that fit under ``RUN_OUTPUT_CHAR_CAP``
|
||||
are returned inline; larger ones are stored and the model gets a char-budgeted
|
||||
preview plus a ``run_<id>`` reference it can page with ``read_run``/``search_run``.
|
||||
Those two read tools are appended to the tool list so every capability-using
|
||||
subagent can follow a truncation reference without extra wiring.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
from langchain_core.tools import BaseTool, StructuredTool
|
||||
|
||||
from app.capabilities.core.billing import charge_capability, gate_capability
|
||||
from app.capabilities.core.runs import (
|
||||
RUN_OUTPUT_CHAR_CAP,
|
||||
record_run,
|
||||
serialize_output,
|
||||
)
|
||||
from app.capabilities.core.store import all_capabilities
|
||||
from app.capabilities.core.types import Capability, CapabilityContext
|
||||
from app.db import async_session_maker
|
||||
|
|
@ -22,31 +32,129 @@ def build_capability_tools(
|
|||
workspace_id: int,
|
||||
capabilities: list[Capability] | None = None,
|
||||
) -> list[BaseTool]:
|
||||
"""Emit one tool per verb (defaults to the whole registry)."""
|
||||
"""Emit one tool per verb (defaults to the whole registry), plus the run readers."""
|
||||
caps = capabilities if capabilities is not None else all_capabilities()
|
||||
return [_capability_tool(cap, workspace_id) for cap in caps]
|
||||
tools = [_capability_tool(cap, workspace_id) for cap in caps]
|
||||
# Deferred import: the reader lives in the agents package (which imports from
|
||||
# here), so importing it lazily avoids an import-time cycle.
|
||||
from app.agents.chat.multi_agent_chat.subagents.shared.run_reader import (
|
||||
build_run_reader_tools,
|
||||
)
|
||||
|
||||
tools.extend(build_run_reader_tools(workspace_id=workspace_id))
|
||||
return tools
|
||||
|
||||
|
||||
def _current_thread_id() -> str | None:
|
||||
"""Best-effort ``configurable.thread_id`` from the active LangGraph config."""
|
||||
try:
|
||||
from langgraph.config import get_config
|
||||
|
||||
cfg = get_config()
|
||||
tid = (cfg.get("configurable") or {}).get("thread_id")
|
||||
return str(tid) if tid is not None else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _capability_tool(capability: Capability, workspace_id: int) -> BaseTool:
|
||||
input_model = capability.input_schema
|
||||
unit = capability.billing_unit
|
||||
executor = capability.executor
|
||||
name = capability.name
|
||||
|
||||
async def _run(**kwargs: object) -> dict | str:
|
||||
payload = input_model(**kwargs)
|
||||
input_dump = payload.model_dump(exclude_none=True)
|
||||
thread_id = _current_thread_id()
|
||||
|
||||
async with async_session_maker() as session:
|
||||
ctx = CapabilityContext(session=session, workspace_id=workspace_id)
|
||||
try:
|
||||
await gate_capability(payload, unit, ctx)
|
||||
except InsufficientCreditsError as exc:
|
||||
return str(exc)
|
||||
output = await executor(payload)
|
||||
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
output = await executor(payload)
|
||||
except Exception as exc:
|
||||
duration_ms = int((time.perf_counter() - started) * 1000)
|
||||
async with async_session_maker() as rec_session:
|
||||
await record_run(
|
||||
rec_session,
|
||||
workspace_id=workspace_id,
|
||||
capability=name,
|
||||
origin="agent",
|
||||
status="error",
|
||||
input=input_dump,
|
||||
error=str(exc),
|
||||
thread_id=thread_id,
|
||||
duration_ms=duration_ms,
|
||||
)
|
||||
raise
|
||||
|
||||
duration_ms = int((time.perf_counter() - started) * 1000)
|
||||
await charge_capability(output, unit, ctx)
|
||||
return output.model_dump()
|
||||
|
||||
serialized = serialize_output(output)
|
||||
async with async_session_maker() as rec_session:
|
||||
run_id = await record_run(
|
||||
rec_session,
|
||||
workspace_id=workspace_id,
|
||||
capability=name,
|
||||
origin="agent",
|
||||
status="success",
|
||||
serialized=serialized,
|
||||
input=input_dump,
|
||||
thread_id=thread_id,
|
||||
duration_ms=duration_ms,
|
||||
)
|
||||
|
||||
if serialized.char_count <= RUN_OUTPUT_CHAR_CAP:
|
||||
dump = output.model_dump(exclude_none=True)
|
||||
if run_id is not None:
|
||||
dump["run_id"] = f"run_{run_id}"
|
||||
return dump
|
||||
|
||||
return _build_preview(serialized, run_id)
|
||||
|
||||
return StructuredTool.from_function(
|
||||
coroutine=_run,
|
||||
name=capability.name.replace(".", "_"),
|
||||
name=name.replace(".", "_"),
|
||||
description=capability.description,
|
||||
args_schema=input_model,
|
||||
)
|
||||
|
||||
|
||||
def _build_preview(serialized, run_id: str | None) -> str:
|
||||
"""Char-budgeted preview: whole JSONL items until the cap is spent."""
|
||||
lines = serialized.text.split("\n")
|
||||
preview_lines: list[str] = []
|
||||
used = 0
|
||||
for line in lines:
|
||||
cost = len(line) + 1
|
||||
if used + cost > RUN_OUTPUT_CHAR_CAP:
|
||||
break
|
||||
preview_lines.append(line)
|
||||
used += cost
|
||||
|
||||
if not preview_lines and lines:
|
||||
# A single item larger than the cap: show a clipped head so the model
|
||||
# still sees the shape and can page/search for the rest.
|
||||
preview_lines = [lines[0][:RUN_OUTPUT_CHAR_CAP]]
|
||||
|
||||
shown = len(preview_lines)
|
||||
preview = "\n".join(preview_lines)
|
||||
|
||||
if run_id is None:
|
||||
return (
|
||||
f"{preview}\n\n...Showing {shown} of {serialized.item_count} items "
|
||||
f"({serialized.char_count} chars). Full output unavailable (storage error)."
|
||||
)
|
||||
return (
|
||||
f"{preview}\n\n...Showing {shown} of {serialized.item_count} items "
|
||||
f"({serialized.char_count} chars). Full run stored as run_{run_id}. Use "
|
||||
f"read_run('run_{run_id}', offset, limit) or search_run('run_{run_id}', "
|
||||
"pattern) to inspect the rest."
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,32 +1,77 @@
|
|||
"""Generate the REST door from the capability registry (05).
|
||||
|
||||
One typed ``POST`` per verb; each runs the same thin adapter:
|
||||
authn -> workspace authz -> meter-gate -> executor -> charge -> typed output.
|
||||
One typed ``POST`` per verb under ``/workspaces/{id}/scrapers/{platform}/{verb}``;
|
||||
each runs the same thin adapter: authn -> workspace authz -> meter-gate -> executor
|
||||
-> charge -> typed output. Every request is recorded to the ``runs`` table
|
||||
(best-effort) and its id returned via the ``X-Run-Id`` header. Two ``GET`` routes
|
||||
expose the run history that backs the Scraper-API logs UI.
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
import time
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Response, status
|
||||
from pydantic import BaseModel
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.auth.context import AuthContext
|
||||
from app.capabilities.core.access.rate_limit import enforce_capability_rate_limit
|
||||
from app.capabilities.core.billing import charge_capability, gate_capability
|
||||
from app.capabilities.core.runs import record_run, serialize_output
|
||||
from app.capabilities.core.store import all_capabilities
|
||||
from app.capabilities.core.types import Capability, CapabilityContext
|
||||
from app.db import get_async_session
|
||||
from app.db import Run, async_session_maker, get_async_session
|
||||
from app.exceptions import ExternalServiceError, SurfSenseError
|
||||
from app.services.web_crawl_credit_service import InsufficientCreditsError
|
||||
from app.users import get_auth_context
|
||||
from app.utils.rbac import check_workspace_access
|
||||
|
||||
|
||||
class RunSummary(BaseModel):
|
||||
"""Metadata row for the runs list (output body omitted)."""
|
||||
|
||||
id: str
|
||||
capability: str
|
||||
origin: str
|
||||
status: str
|
||||
item_count: int
|
||||
char_count: int
|
||||
duration_ms: int | None
|
||||
cost_micros: int | None
|
||||
error: str | None
|
||||
created_at: datetime
|
||||
|
||||
|
||||
class RunDetail(RunSummary):
|
||||
"""Full run including input and stored output."""
|
||||
|
||||
thread_id: str | None
|
||||
input: dict | None
|
||||
output_text: str | None
|
||||
|
||||
|
||||
def _origin_for(auth: AuthContext) -> str:
|
||||
"""Session callers are the in-app UI; PAT/system callers are the public API."""
|
||||
return "ui" if getattr(auth, "method", None) == "session" else "api"
|
||||
|
||||
|
||||
async def _record_rest_run(**kwargs) -> str | None:
|
||||
"""Record a run on a dedicated session so it survives a failed request txn."""
|
||||
async with async_session_maker() as session:
|
||||
return await record_run(session, **kwargs)
|
||||
|
||||
|
||||
def build_capabilities_router(
|
||||
capabilities: list[Capability] | None = None,
|
||||
) -> APIRouter:
|
||||
"""Emit one typed route per verb (defaults to the whole registry)."""
|
||||
router = APIRouter(tags=["capabilities"])
|
||||
"""Emit one typed route per verb (defaults to the whole registry) + run history."""
|
||||
router = APIRouter(tags=["scrapers"])
|
||||
caps = capabilities if capabilities is not None else all_capabilities()
|
||||
for capability in caps:
|
||||
_register_verb(router, capability)
|
||||
_register_run_history(router)
|
||||
return router
|
||||
|
||||
|
||||
|
|
@ -35,10 +80,13 @@ def _register_verb(router: APIRouter, capability: Capability) -> None:
|
|||
output_model = capability.output_schema
|
||||
unit = capability.billing_unit
|
||||
executor = capability.executor
|
||||
name = capability.name
|
||||
platform, _, verb = name.partition(".")
|
||||
|
||||
async def endpoint(
|
||||
workspace_id: int,
|
||||
payload: input_model,
|
||||
response: Response,
|
||||
session: AsyncSession = Depends(get_async_session),
|
||||
auth: AuthContext = Depends(get_auth_context),
|
||||
):
|
||||
|
|
@ -57,24 +105,152 @@ def _register_verb(router: APIRouter, capability: Capability) -> None:
|
|||
},
|
||||
) from exc
|
||||
|
||||
input_dump = payload.model_dump(exclude_none=True)
|
||||
user_id = getattr(auth.user, "id", None)
|
||||
origin = _origin_for(auth)
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
output = await executor(payload)
|
||||
except (SurfSenseError, HTTPException):
|
||||
except (SurfSenseError, HTTPException) as exc:
|
||||
await _record_rest_run(
|
||||
workspace_id=workspace_id,
|
||||
capability=name,
|
||||
origin=origin,
|
||||
status="error",
|
||||
input=input_dump,
|
||||
user_id=user_id,
|
||||
error=str(exc),
|
||||
duration_ms=int((time.perf_counter() - started) * 1000),
|
||||
)
|
||||
raise
|
||||
except Exception as exc:
|
||||
await _record_rest_run(
|
||||
workspace_id=workspace_id,
|
||||
capability=name,
|
||||
origin=origin,
|
||||
status="error",
|
||||
input=input_dump,
|
||||
user_id=user_id,
|
||||
error=str(exc),
|
||||
duration_ms=int((time.perf_counter() - started) * 1000),
|
||||
)
|
||||
raise ExternalServiceError(
|
||||
f"The '{capability.name}' capability failed due to an upstream error.",
|
||||
f"The '{name}' capability failed due to an upstream error.",
|
||||
code="CAPABILITY_UPSTREAM_ERROR",
|
||||
) from exc
|
||||
|
||||
duration_ms = int((time.perf_counter() - started) * 1000)
|
||||
await charge_capability(output, unit, ctx)
|
||||
|
||||
serialized = serialize_output(output)
|
||||
run_id = await _record_rest_run(
|
||||
workspace_id=workspace_id,
|
||||
capability=name,
|
||||
origin=origin,
|
||||
status="success",
|
||||
serialized=serialized,
|
||||
input=input_dump,
|
||||
user_id=user_id,
|
||||
duration_ms=duration_ms,
|
||||
)
|
||||
if run_id is not None:
|
||||
response.headers["X-Run-Id"] = f"run_{run_id}"
|
||||
return output
|
||||
|
||||
router.add_api_route(
|
||||
f"/workspaces/{{workspace_id}}/capabilities/{capability.name}",
|
||||
f"/workspaces/{{workspace_id}}/scrapers/{platform}/{verb}",
|
||||
endpoint,
|
||||
methods=["POST"],
|
||||
response_model=output_model,
|
||||
name=f"capability:{capability.name}",
|
||||
name=f"scraper:{name}",
|
||||
dependencies=[Depends(enforce_capability_rate_limit)],
|
||||
)
|
||||
|
||||
|
||||
def _register_run_history(router: APIRouter) -> None:
|
||||
"""Register the ``GET`` list + detail routes that back the logs UI."""
|
||||
|
||||
async def list_runs(
|
||||
workspace_id: int,
|
||||
session: AsyncSession = Depends(get_async_session),
|
||||
auth: AuthContext = Depends(get_auth_context),
|
||||
limit: int = Query(default=50, ge=1, le=100),
|
||||
offset: int = Query(default=0, ge=0),
|
||||
capability: str | None = Query(default=None),
|
||||
run_status: str | None = Query(default=None, alias="status"),
|
||||
) -> list[RunSummary]:
|
||||
await check_workspace_access(session, auth, workspace_id)
|
||||
stmt = select(Run).where(Run.workspace_id == workspace_id)
|
||||
if capability:
|
||||
stmt = stmt.where(Run.capability == capability)
|
||||
if run_status:
|
||||
stmt = stmt.where(Run.status == run_status)
|
||||
stmt = stmt.order_by(Run.created_at.desc()).limit(limit).offset(offset)
|
||||
rows = (await session.execute(stmt)).scalars().all()
|
||||
return [_to_summary(row) for row in rows]
|
||||
|
||||
async def get_run(
|
||||
workspace_id: int,
|
||||
run_id: str,
|
||||
session: AsyncSession = Depends(get_async_session),
|
||||
auth: AuthContext = Depends(get_auth_context),
|
||||
) -> RunDetail:
|
||||
await check_workspace_access(session, auth, workspace_id)
|
||||
raw_id = run_id[len("run_") :] if run_id.startswith("run_") else run_id
|
||||
try:
|
||||
parsed_id = uuid.UUID(raw_id)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND, detail="Run not found."
|
||||
) from exc
|
||||
row = (
|
||||
await session.execute(
|
||||
select(Run).where(
|
||||
Run.id == parsed_id, Run.workspace_id == workspace_id
|
||||
)
|
||||
)
|
||||
).scalar_one_or_none()
|
||||
if row is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND, detail="Run not found."
|
||||
)
|
||||
return _to_detail(row)
|
||||
|
||||
router.add_api_route(
|
||||
"/workspaces/{workspace_id}/scrapers/runs",
|
||||
list_runs,
|
||||
methods=["GET"],
|
||||
response_model=list[RunSummary],
|
||||
name="scraper:list_runs",
|
||||
)
|
||||
router.add_api_route(
|
||||
"/workspaces/{workspace_id}/scrapers/runs/{run_id}",
|
||||
get_run,
|
||||
methods=["GET"],
|
||||
response_model=RunDetail,
|
||||
name="scraper:get_run",
|
||||
)
|
||||
|
||||
|
||||
def _to_summary(row: Run) -> RunSummary:
|
||||
return RunSummary(
|
||||
id=f"run_{row.id}",
|
||||
capability=row.capability,
|
||||
origin=row.origin,
|
||||
status=row.status,
|
||||
item_count=row.item_count,
|
||||
char_count=row.char_count,
|
||||
duration_ms=row.duration_ms,
|
||||
cost_micros=row.cost_micros,
|
||||
error=row.error,
|
||||
created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
def _to_detail(row: Run) -> RunDetail:
|
||||
return RunDetail(
|
||||
**_to_summary(row).model_dump(),
|
||||
thread_id=row.thread_id,
|
||||
input=row.input,
|
||||
output_text=row.output_text,
|
||||
)
|
||||
|
|
|
|||
191
surfsense_backend/app/capabilities/core/runs.py
Normal file
191
surfsense_backend/app/capabilities/core/runs.py
Normal file
|
|
@ -0,0 +1,191 @@
|
|||
"""Shared run recorder: persist one ``runs`` row per capability invocation.
|
||||
|
||||
Both doors (the agent tool adapter and the REST endpoint) call :func:`record_run`
|
||||
so agent and API runs land identically. Output is serialized to JSONL (one item
|
||||
per line, ``exclude_none``) so the ``read_run``/``search_run`` tools can page and
|
||||
grep by line. Recording is best-effort: a failure here never fails an otherwise
|
||||
successful scrape — the caller degrades gracefully on a ``None`` return.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from pydantic import BaseModel
|
||||
from sqlalchemy import text
|
||||
|
||||
from app.db import Run, ToolOutputSpill
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
RUN_OUTPUT_CHAR_CAP = 40_000
|
||||
"""~10k tokens. Capability outputs at/under this are returned inline; larger ones
|
||||
are stored and previewed. Read-tool responses pass through the same cap."""
|
||||
|
||||
RUNS_RETENTION_DAYS = 30
|
||||
SPILLS_RETENTION_DAYS = 7
|
||||
_CLEANUP_BATCH = 200
|
||||
_CLEANUP_SAMPLE_RATE = 0.01
|
||||
"""ponytail: opportunistic bounded cleanup fired on ~1% of inserts. A dedicated
|
||||
cron/scheduler is the upgrade path if row volume ever outpaces this."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SerializedOutput:
|
||||
"""A capability output flattened to JSONL, computed once and reused."""
|
||||
|
||||
text: str
|
||||
item_count: int
|
||||
char_count: int
|
||||
|
||||
|
||||
def serialize_output(output: BaseModel) -> SerializedOutput:
|
||||
"""Flatten a capability output to JSONL (one item per line, ``exclude_none``).
|
||||
|
||||
Capability outputs wrap their results in an ``items`` list; each item becomes
|
||||
one JSON line so line-based paging/grep works. A non-``items`` output is dumped
|
||||
as a single line.
|
||||
"""
|
||||
items = getattr(output, "items", None)
|
||||
if isinstance(items, list):
|
||||
lines = [
|
||||
json.dumps(_dump(item), default=str, ensure_ascii=False) for item in items
|
||||
]
|
||||
else:
|
||||
lines = [
|
||||
json.dumps(
|
||||
output.model_dump(exclude_none=True), default=str, ensure_ascii=False
|
||||
)
|
||||
]
|
||||
|
||||
body = "\n".join(lines)
|
||||
return SerializedOutput(text=body, item_count=len(lines), char_count=len(body))
|
||||
|
||||
|
||||
def _dump(item: Any) -> Any:
|
||||
if isinstance(item, BaseModel):
|
||||
return item.model_dump(exclude_none=True)
|
||||
return item
|
||||
|
||||
|
||||
async def record_run(
|
||||
session: AsyncSession,
|
||||
*,
|
||||
workspace_id: int,
|
||||
capability: str,
|
||||
origin: str,
|
||||
status: str,
|
||||
serialized: SerializedOutput | None = None,
|
||||
input: dict | None = None,
|
||||
user_id: Any | None = None,
|
||||
thread_id: str | None = None,
|
||||
error: str | None = None,
|
||||
duration_ms: int | None = None,
|
||||
cost_micros: int | None = None,
|
||||
) -> str | None:
|
||||
"""Persist a run row and return its id, or ``None`` on failure (best-effort).
|
||||
|
||||
Both doors pass a dedicated session (from ``async_session_maker``), so this
|
||||
function owns the commit — recording never entangles the request transaction
|
||||
and survives an executor error that leaves the request session unusable.
|
||||
"""
|
||||
try:
|
||||
run = Run(
|
||||
workspace_id=workspace_id,
|
||||
user_id=user_id,
|
||||
thread_id=thread_id,
|
||||
capability=capability,
|
||||
origin=origin,
|
||||
status=status,
|
||||
error=error,
|
||||
input=input,
|
||||
output_text=serialized.text if serialized else None,
|
||||
item_count=serialized.item_count if serialized else 0,
|
||||
char_count=serialized.char_count if serialized else 0,
|
||||
duration_ms=duration_ms,
|
||||
cost_micros=cost_micros,
|
||||
)
|
||||
session.add(run)
|
||||
await session.flush()
|
||||
run_id = str(run.id)
|
||||
await _maybe_cleanup(session, "runs", RUNS_RETENTION_DAYS)
|
||||
await session.commit()
|
||||
return run_id
|
||||
except Exception:
|
||||
logger.exception("record_run failed for capability=%s", capability)
|
||||
try:
|
||||
await session.rollback()
|
||||
except Exception:
|
||||
logger.exception("record_run rollback failed")
|
||||
return None
|
||||
|
||||
|
||||
async def record_spill(
|
||||
session: AsyncSession,
|
||||
*,
|
||||
content: str,
|
||||
spill_id: Any | None = None,
|
||||
workspace_id: int | None = None,
|
||||
thread_id: str | None = None,
|
||||
tool_name: str | None = None,
|
||||
) -> str | None:
|
||||
"""Persist a context-editing spill row and return its id, or ``None``.
|
||||
|
||||
``spill_id`` may be supplied so the caller's placeholder can reference the id
|
||||
before the row is flushed (the context-editing middleware needs this). The
|
||||
write is idempotent on that id: context edits re-apply on every model call
|
||||
(they operate on a per-call copy of the messages), so the same spill arrives
|
||||
repeatedly — an existing row is left as-is and its id returned.
|
||||
"""
|
||||
try:
|
||||
kwargs: dict[str, Any] = {}
|
||||
if spill_id is not None:
|
||||
kwargs["id"] = spill_id
|
||||
existing = await session.get(ToolOutputSpill, spill_id)
|
||||
if existing is not None:
|
||||
return str(spill_id)
|
||||
spill = ToolOutputSpill(
|
||||
workspace_id=workspace_id,
|
||||
thread_id=thread_id,
|
||||
tool_name=tool_name,
|
||||
content=content,
|
||||
char_count=len(content),
|
||||
**kwargs,
|
||||
)
|
||||
session.add(spill)
|
||||
await session.flush()
|
||||
spill_id = str(spill.id)
|
||||
await _maybe_cleanup(session, "tool_output_spills", SPILLS_RETENTION_DAYS)
|
||||
await session.commit()
|
||||
return spill_id
|
||||
except Exception:
|
||||
logger.exception("record_spill failed")
|
||||
try:
|
||||
await session.rollback()
|
||||
except Exception:
|
||||
logger.exception("record_spill rollback failed")
|
||||
return None
|
||||
|
||||
|
||||
async def _maybe_cleanup(session: AsyncSession, table: str, retention_days: int) -> None:
|
||||
"""Delete a bounded batch of expired rows on ~1% of inserts."""
|
||||
if random.random() >= _CLEANUP_SAMPLE_RATE:
|
||||
return
|
||||
cutoff = datetime.now(UTC) - timedelta(days=retention_days)
|
||||
# ponytail: LIMIT-bounded so a long backlog never lands as one giant delete.
|
||||
# `table` is one of two hardcoded literals below — never user input.
|
||||
await session.execute(
|
||||
text(
|
||||
f"DELETE FROM {table} WHERE id IN "
|
||||
f"(SELECT id FROM {table} WHERE created_at < :cutoff LIMIT :batch)"
|
||||
),
|
||||
{"cutoff": cutoff, "batch": _CLEANUP_BATCH},
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue