mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-07-10 22:32:16 +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:
DESKTOP-RTLN3BA\$punk
parent
ab747e7a49
commit
b6e378b070
14 changed files with 1247 additions and 130 deletions
|
|
@ -3,7 +3,6 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.tools import BaseTool
|
||||
|
||||
|
|
@ -25,7 +24,7 @@ def build_context_editing_mw(
|
|||
flags: AgentFeatureFlags,
|
||||
max_input_tokens: int | None,
|
||||
tools: Sequence[BaseTool],
|
||||
backend_resolver: Any,
|
||||
workspace_id: int | None = None,
|
||||
) -> SpillingContextEditingMiddleware | None:
|
||||
if not enabled(flags, "enable_context_editing") or not max_input_tokens:
|
||||
return None
|
||||
|
|
@ -46,5 +45,5 @@ def build_context_editing_mw(
|
|||
)
|
||||
return SpillingContextEditingMiddleware(
|
||||
edits=[spill_edit, clear_edit],
|
||||
backend_resolver=backend_resolver,
|
||||
workspace_id=workspace_id,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -4,25 +4,31 @@ SpillToBackendEdit + SpillingContextEditingMiddleware.
|
|||
LangChain's :class:`ClearToolUsesEdit` discards old ``ToolMessage.content``
|
||||
when the context-editing budget triggers, replacing the body with a fixed
|
||||
placeholder. That's lossy: anything the agent might want to revisit is
|
||||
gone. The spill-to-disk pattern (originally from OpenCode's
|
||||
gone. The spill pattern (originally from OpenCode's
|
||||
``opencode/packages/opencode/src/tool/truncate.ts``) keeps the prune
|
||||
behavior but writes the full original payload to the runtime backend
|
||||
under ``/tool_outputs/{thread_id}/{message_id}.txt`` first. The
|
||||
placeholder is then upgraded to point at the spill path so the agent
|
||||
(or a subagent) can read it back on demand.
|
||||
behavior but persists the full original payload first — to the
|
||||
``tool_output_spills`` table — and upgrades the placeholder to a
|
||||
``spill_<uuid>`` reference the agent can read back with the shared
|
||||
``read_run``/``search_run`` tools on demand.
|
||||
|
||||
Why the DB and not the runtime filesystem backend: the previous version
|
||||
wrote via ``backend.aupload_files``, which is a no-op on cloud
|
||||
(``KBPostgresBackend`` raises ``NotImplementedError``) and mismatched on
|
||||
desktop (paths must be ``/{mount_id}/...``), so spills were unrecoverable
|
||||
in production. A table works on every deployment and needs no sandbox.
|
||||
|
||||
Why this is a middleware subclass instead of a plain ``ContextEdit``:
|
||||
``ContextEdit.apply`` is sync, but writing to the backend is async. We
|
||||
capture the spill payloads inside ``apply`` and flush them via
|
||||
``await backend.aupload_files(...)`` from ``awrap_model_call`` *before*
|
||||
delegating to the handler, so the explore subagent can always read what
|
||||
the placeholder advertises.
|
||||
``ContextEdit.apply`` is sync, but the DB write is async. We generate the
|
||||
spill id and capture the payload inside ``apply`` (so the placeholder can
|
||||
reference the final id immediately) and flush the rows from
|
||||
``awrap_model_call`` *before* delegating to the handler.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass, field
|
||||
|
|
@ -44,7 +50,6 @@ from langchain_core.messages.utils import count_tokens_approximately
|
|||
from langgraph.config import get_config
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from deepagents.backends.protocol import BackendProtocol
|
||||
from langchain.agents.middleware.types import (
|
||||
ModelRequest,
|
||||
ModelResponse,
|
||||
|
|
@ -52,24 +57,27 @@ if TYPE_CHECKING:
|
|||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_SPILL_PREFIX = "/tool_outputs"
|
||||
# Namespace for deterministic spill ids: the same (thread, message) always maps
|
||||
# to the same row, so re-running the edit on later model calls (edits apply to a
|
||||
# per-call copy of the messages, never to persisted state) re-references the
|
||||
# existing spill instead of inserting a duplicate every turn.
|
||||
_SPILL_NAMESPACE = uuid.UUID("6ba7b810-9dad-11d1-80b4-00c04fd430c8")
|
||||
|
||||
|
||||
def _build_spill_placeholder(spill_path: str) -> str:
|
||||
def _spill_id_for(thread_id: str | None, message_key: str) -> uuid.UUID:
|
||||
return uuid.uuid5(_SPILL_NAMESPACE, f"{thread_id or 'no_thread'}:{message_key}")
|
||||
|
||||
|
||||
def _build_spill_placeholder(spill_id: uuid.UUID) -> str:
|
||||
"""Build the user-facing placeholder text shown to the model."""
|
||||
return (
|
||||
f"[cleared — full output at {spill_path}; ask the explore subagent to read it]"
|
||||
f"[cleared — full output stored as spill_{spill_id}; "
|
||||
"use read_run/search_run to read it]"
|
||||
)
|
||||
|
||||
|
||||
def _get_thread_id_or_session() -> str:
|
||||
"""Best-effort thread_id discovery for the spill path.
|
||||
|
||||
Falls back to a process-stable string if no LangGraph config is
|
||||
available (e.g. unit tests). The exact value doesn't matter as long
|
||||
as it's stable within one stream so the placeholder paths line up
|
||||
with the actual upload path.
|
||||
"""
|
||||
def _get_thread_id() -> str | None:
|
||||
"""Best-effort ``configurable.thread_id`` for the spill row (``None`` if absent)."""
|
||||
try:
|
||||
config = get_config()
|
||||
thread_id = config.get("configurable", {}).get("thread_id")
|
||||
|
|
@ -77,17 +85,18 @@ def _get_thread_id_or_session() -> str:
|
|||
return str(thread_id)
|
||||
except RuntimeError:
|
||||
pass
|
||||
return "no_thread"
|
||||
return None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class SpillToBackendEdit(ContextEdit):
|
||||
"""Capture-and-replace context edit that spills full tool output to the backend.
|
||||
"""Capture-and-replace context edit that spills full tool output to the DB.
|
||||
|
||||
Behaves like :class:`ClearToolUsesEdit` (same trigger / keep / exclude
|
||||
semantics) **and** records the original ``ToolMessage.content`` in
|
||||
:attr:`pending_spills` so the wrapping middleware can flush them
|
||||
before the model call.
|
||||
:attr:`pending_spills` so the wrapping middleware can flush the rows to
|
||||
``tool_output_spills`` before the model call. The spill id is generated up
|
||||
front so the placeholder can reference it immediately.
|
||||
|
||||
Args:
|
||||
trigger: Token threshold above which the edit fires.
|
||||
|
|
@ -97,8 +106,6 @@ class SpillToBackendEdit(ContextEdit):
|
|||
exclude_tools: Names of tools whose output is NOT spilled.
|
||||
clear_tool_inputs: Also clear the originating ``AIMessage.tool_calls``
|
||||
args when their pair is cleared.
|
||||
path_prefix: Path under the backend where spills are written.
|
||||
Default ``"/tool_outputs"``.
|
||||
"""
|
||||
|
||||
trigger: int = 100_000
|
||||
|
|
@ -106,12 +113,16 @@ class SpillToBackendEdit(ContextEdit):
|
|||
keep: int = 3
|
||||
clear_tool_inputs: bool = False
|
||||
exclude_tools: Sequence[str] = ()
|
||||
path_prefix: str = DEFAULT_SPILL_PREFIX
|
||||
|
||||
pending_spills: list[tuple[str, bytes]] = field(default_factory=list)
|
||||
# (spill_id, content_bytes, tool_name, thread_id)
|
||||
pending_spills: list[tuple[uuid.UUID, bytes, str | None, str | None]] = field(
|
||||
default_factory=list
|
||||
)
|
||||
_lock: threading.Lock = field(default_factory=threading.Lock)
|
||||
|
||||
def drain_pending(self) -> list[tuple[str, bytes]]:
|
||||
def drain_pending(
|
||||
self,
|
||||
) -> list[tuple[uuid.UUID, bytes, str | None, str | None]]:
|
||||
"""Return and clear the pending-spill list atomically."""
|
||||
with self._lock:
|
||||
out = list(self.pending_spills)
|
||||
|
|
@ -139,7 +150,7 @@ class SpillToBackendEdit(ContextEdit):
|
|||
if self.keep:
|
||||
candidates = candidates[: -self.keep]
|
||||
|
||||
thread_id = _get_thread_id_or_session()
|
||||
thread_id = _get_thread_id()
|
||||
excluded_tools = set(self.exclude_tools)
|
||||
|
||||
for idx, tool_message in candidates:
|
||||
|
|
@ -168,24 +179,23 @@ class SpillToBackendEdit(ContextEdit):
|
|||
if tool_name in excluded_tools:
|
||||
continue
|
||||
|
||||
message_id = tool_message.id or tool_message.tool_call_id or "unknown"
|
||||
spill_path = f"{self.path_prefix}/{thread_id}/{message_id}.txt"
|
||||
|
||||
message_key = tool_message.id or tool_message.tool_call_id or "unknown"
|
||||
spill_id = _spill_id_for(thread_id, message_key)
|
||||
original = tool_message.content
|
||||
payload = self._encode_payload(original)
|
||||
with self._lock:
|
||||
self.pending_spills.append((spill_path, payload))
|
||||
self.pending_spills.append((spill_id, payload, tool_name, thread_id))
|
||||
|
||||
messages[idx] = tool_message.model_copy(
|
||||
update={
|
||||
"artifact": None,
|
||||
"content": _build_spill_placeholder(spill_path),
|
||||
"content": _build_spill_placeholder(spill_id),
|
||||
"response_metadata": {
|
||||
**tool_message.response_metadata,
|
||||
"context_editing": {
|
||||
"cleared": True,
|
||||
"strategy": "spill_to_backend",
|
||||
"spill_path": spill_path,
|
||||
"strategy": "spill_to_db",
|
||||
"spill_id": str(spill_id),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
|
@ -243,52 +253,30 @@ class SpillToBackendEdit(ContextEdit):
|
|||
)
|
||||
|
||||
|
||||
BackendResolver = "Callable[[Any], BackendProtocol] | BackendProtocol"
|
||||
|
||||
|
||||
class SpillingContextEditingMiddleware(ContextEditingMiddleware):
|
||||
""":class:`ContextEditingMiddleware` that flushes :class:`SpillToBackendEdit` writes.
|
||||
|
||||
Runs the configured edits as the parent does, then flushes any
|
||||
pending spills via the supplied backend resolver before delegating
|
||||
to the model handler. Spill failures are logged but never abort the
|
||||
model call — the placeholder text is already in the message, so the
|
||||
worst case is the agent gets a placeholder it cannot follow up on.
|
||||
Runs the configured edits as the parent does, then persists any pending
|
||||
spills to ``tool_output_spills`` before delegating to the model handler.
|
||||
Spill failures are logged but never abort the model call — the placeholder
|
||||
text is already in the message, so the worst case is the agent gets a
|
||||
placeholder it cannot follow up on.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
edits: Sequence[ContextEdit],
|
||||
backend_resolver: BackendResolver | None = None,
|
||||
workspace_id: int | None = None,
|
||||
token_count_method: str = "approximate",
|
||||
) -> None:
|
||||
super().__init__(edits=list(edits), token_count_method=token_count_method) # type: ignore[arg-type]
|
||||
self._backend_resolver = backend_resolver
|
||||
self._workspace_id = workspace_id
|
||||
|
||||
def _resolve_backend(self, request: ModelRequest) -> BackendProtocol | None:
|
||||
if self._backend_resolver is None:
|
||||
return None
|
||||
if callable(self._backend_resolver):
|
||||
try:
|
||||
from langchain.tools import ToolRuntime
|
||||
|
||||
tool_runtime = ToolRuntime(
|
||||
state=getattr(request, "state", {}),
|
||||
context=getattr(request.runtime, "context", None),
|
||||
stream_writer=getattr(request.runtime, "stream_writer", None),
|
||||
store=getattr(request.runtime, "store", None),
|
||||
config=getattr(request.runtime, "config", None) or {},
|
||||
tool_call_id=None,
|
||||
)
|
||||
return self._backend_resolver(tool_runtime)
|
||||
except Exception:
|
||||
logger.exception("Failed to resolve spill backend")
|
||||
return None
|
||||
return self._backend_resolver # type: ignore[return-value]
|
||||
|
||||
def _collect_pending(self) -> list[tuple[str, bytes]]:
|
||||
out: list[tuple[str, bytes]] = []
|
||||
def _collect_pending(
|
||||
self,
|
||||
) -> list[tuple[uuid.UUID, bytes, str | None, str | None]]:
|
||||
out: list[tuple[uuid.UUID, bytes, str | None, str | None]] = []
|
||||
for edit in self.edits:
|
||||
if isinstance(edit, SpillToBackendEdit):
|
||||
out.extend(edit.drain_pending())
|
||||
|
|
@ -321,28 +309,37 @@ class SpillingContextEditingMiddleware(ContextEditingMiddleware):
|
|||
|
||||
pending = self._collect_pending()
|
||||
if pending:
|
||||
backend = self._resolve_backend(request)
|
||||
if backend is not None:
|
||||
try:
|
||||
await backend.aupload_files(pending)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Spill-to-backend upload failed (%d files); placeholders "
|
||||
"remain in messages but content is unrecoverable",
|
||||
len(pending),
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"SpillToBackendEdit produced %d pending spills but no backend "
|
||||
"resolver was configured; content is unrecoverable",
|
||||
len(pending),
|
||||
)
|
||||
await self._flush_spills(pending)
|
||||
|
||||
return await handler(request.override(messages=edited_messages))
|
||||
|
||||
async def _flush_spills(
|
||||
self, pending: list[tuple[uuid.UUID, bytes, str | None, str | None]]
|
||||
) -> None:
|
||||
"""Persist spilled tool outputs to the DB (best-effort)."""
|
||||
from app.capabilities.core.runs import record_spill
|
||||
from app.db import async_session_maker
|
||||
|
||||
try:
|
||||
async with async_session_maker() as session:
|
||||
for spill_id, payload, tool_name, thread_id in pending:
|
||||
await record_spill(
|
||||
session,
|
||||
spill_id=spill_id,
|
||||
content=payload.decode("utf-8", errors="replace"),
|
||||
workspace_id=self._workspace_id,
|
||||
thread_id=thread_id,
|
||||
tool_name=tool_name,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Spill-to-DB flush failed (%d rows); placeholders remain in "
|
||||
"messages but content is unrecoverable",
|
||||
len(pending),
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_SPILL_PREFIX",
|
||||
"ClearToolUsesEdit",
|
||||
"SpillToBackendEdit",
|
||||
"SpillingContextEditingMiddleware",
|
||||
|
|
|
|||
|
|
@ -7,6 +7,11 @@ retrieval are owned by the ``knowledge_base`` subagent and reached via the
|
|||
the main agent only sees the specialist's grounded summary. The stack here
|
||||
computes the workspace tree, commits any subagent-side staged writes at end of
|
||||
turn (cloud mode), and wires the supporting middleware.
|
||||
|
||||
One deliberate, read-only exception to the pure-router stance: the main agent
|
||||
also carries the ``read_run``/``search_run`` tools (added in ``runtime/factory``)
|
||||
so it can follow the context-editing spill placeholder — evicted tool output is
|
||||
persisted to ``tool_output_spills`` and the placeholder advertises those tools.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -260,7 +265,7 @@ def build_main_agent_deepagent_middleware(
|
|||
flags=flags,
|
||||
max_input_tokens=max_input_tokens,
|
||||
tools=tools,
|
||||
backend_resolver=backend_resolver,
|
||||
workspace_id=workspace_id,
|
||||
),
|
||||
build_compaction_mw(llm),
|
||||
build_noop_injection_mw(flags),
|
||||
|
|
|
|||
|
|
@ -224,6 +224,15 @@ async def create_multi_agent_chat_deep_agent(
|
|||
additional_tools=list(additional_tools) if additional_tools else None,
|
||||
)
|
||||
|
||||
# Read-only exception to the "main agent is a pure router" stance: the
|
||||
# context-editing spill placeholder points at read_run/search_run, so the
|
||||
# main agent needs those tools to follow it. See middleware/stack.py.
|
||||
from app.agents.chat.multi_agent_chat.subagents.shared.run_reader import (
|
||||
build_run_reader_tools,
|
||||
)
|
||||
|
||||
tools = [*list(tools), *build_run_reader_tools(workspace_id=workspace_id)]
|
||||
|
||||
_flags: AgentFeatureFlags = get_flags()
|
||||
if _flags.enable_tool_call_repair and INVALID_TOOL_NAME not in {
|
||||
t.name for t in tools
|
||||
|
|
|
|||
|
|
@ -0,0 +1,175 @@
|
|||
"""``read_run`` / ``search_run``: page and grep a stored run or spill by line.
|
||||
|
||||
Scraper capability outputs and evicted context spills are stored full in Postgres
|
||||
(``runs`` / ``tool_output_spills``); the model only ever sees a capped preview plus
|
||||
a reference like ``run_<uuid>`` or ``spill_<uuid>``. These two tools let the agent
|
||||
retrieve the rest on demand — line-based paging and pattern search — without ever
|
||||
loading the whole payload into context. Every lookup is scoped to the caller's
|
||||
workspace (the trust boundary).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Annotated
|
||||
from uuid import UUID
|
||||
|
||||
from langchain_core.tools import BaseTool, StructuredTool
|
||||
from sqlalchemy import select
|
||||
|
||||
from app.capabilities.core.runs import RUN_OUTPUT_CHAR_CAP
|
||||
from app.db import Run, ToolOutputSpill, shielded_async_session
|
||||
|
||||
_MAX_LIMIT = 100
|
||||
_MAX_PATTERN_LEN = 200
|
||||
"""ReDoS guard: reject/simplify absurdly long model-supplied patterns."""
|
||||
|
||||
|
||||
def _parse_ref(ref: str) -> tuple[str, UUID] | None:
|
||||
"""Split ``run_<uuid>`` / ``spill_<uuid>`` into ``(kind, uuid)``; ``None`` if malformed."""
|
||||
ref = (ref or "").strip()
|
||||
for prefix, kind in (("run_", "run"), ("spill_", "spill")):
|
||||
if ref.startswith(prefix):
|
||||
try:
|
||||
return kind, UUID(ref[len(prefix) :])
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
async def _load_body(ref: str, workspace_id: int) -> tuple[str, str] | str:
|
||||
"""Return ``(body_text, kind)`` for a ref, or an error string.
|
||||
|
||||
Workspace-scoped: a ref belonging to another workspace reads as not found.
|
||||
"""
|
||||
parsed = _parse_ref(ref)
|
||||
if parsed is None:
|
||||
return (
|
||||
f"Error: '{ref}' is not a valid run reference. "
|
||||
"Expected 'run_<uuid>' or 'spill_<uuid>'."
|
||||
)
|
||||
kind, ref_id = parsed
|
||||
async with shielded_async_session() as session:
|
||||
if kind == "run":
|
||||
row = (
|
||||
await session.execute(
|
||||
select(Run.output_text).where(
|
||||
Run.id == ref_id, Run.workspace_id == workspace_id
|
||||
)
|
||||
)
|
||||
).scalar_one_or_none()
|
||||
else:
|
||||
row = (
|
||||
await session.execute(
|
||||
select(ToolOutputSpill.content).where(
|
||||
ToolOutputSpill.id == ref_id,
|
||||
ToolOutputSpill.workspace_id == workspace_id,
|
||||
)
|
||||
)
|
||||
).scalar_one_or_none()
|
||||
if row is None:
|
||||
return f"Error: {ref} not found in this workspace."
|
||||
return (row or ""), kind
|
||||
|
||||
|
||||
def _cap(body: str) -> str:
|
||||
"""Clip a response to the shared char cap with an explicit truncation note."""
|
||||
if len(body) <= RUN_OUTPUT_CHAR_CAP:
|
||||
return body
|
||||
return (
|
||||
body[:RUN_OUTPUT_CHAR_CAP]
|
||||
+ f"\n\n...[response truncated at {RUN_OUTPUT_CHAR_CAP} chars; "
|
||||
"narrow with a larger offset or use search_run]..."
|
||||
)
|
||||
|
||||
|
||||
def build_run_reader_tools(*, workspace_id: int) -> list[BaseTool]:
|
||||
"""Build the ``read_run`` / ``search_run`` tools bound to one workspace."""
|
||||
|
||||
async def _read_run(
|
||||
ref: Annotated[
|
||||
str, "The run reference to read, e.g. 'run_<uuid>' or 'spill_<uuid>'."
|
||||
],
|
||||
offset: Annotated[int, "0-based line (item) index to start from."] = 0,
|
||||
limit: Annotated[int, "Max lines (items) to return (default 20)."] = 20,
|
||||
) -> str:
|
||||
loaded = await _load_body(ref, workspace_id)
|
||||
if isinstance(loaded, str):
|
||||
return loaded
|
||||
body, _kind = loaded
|
||||
lines = body.split("\n")
|
||||
start = max(0, offset)
|
||||
count = min(max(1, limit), _MAX_LIMIT)
|
||||
window = lines[start : start + count]
|
||||
if not window:
|
||||
return (
|
||||
f"No lines at offset {start} (total {len(lines)} lines in {ref})."
|
||||
)
|
||||
header = (
|
||||
f"Showing lines {start}-{start + len(window) - 1} of {len(lines)} in {ref}:\n"
|
||||
)
|
||||
return _cap(header + "\n".join(window))
|
||||
|
||||
async def _search_run(
|
||||
ref: Annotated[str, "The run reference to search, e.g. 'run_<uuid>'."],
|
||||
pattern: Annotated[str, "Substring or regular expression to match per line."],
|
||||
max_matches: Annotated[int, "Max matching lines to return (default 20)."] = 20,
|
||||
) -> str:
|
||||
loaded = await _load_body(ref, workspace_id)
|
||||
if isinstance(loaded, str):
|
||||
return loaded
|
||||
body, _kind = loaded
|
||||
pattern = (pattern or "").strip()
|
||||
if not pattern:
|
||||
return "Error: provide a non-empty search pattern."
|
||||
|
||||
matcher = _build_matcher(pattern)
|
||||
limit = min(max(1, max_matches), _MAX_LIMIT)
|
||||
matches: list[str] = []
|
||||
total = 0
|
||||
for idx, line in enumerate(body.split("\n")):
|
||||
if matcher(line):
|
||||
total += 1
|
||||
if len(matches) < limit:
|
||||
matches.append(f"[{idx}] {line}")
|
||||
if not matches:
|
||||
return f"No lines in {ref} matched {pattern!r}."
|
||||
header = (
|
||||
f"Found {total} matching line(s) in {ref}"
|
||||
f"{f' (showing first {limit})' if total > limit else ''}:\n"
|
||||
)
|
||||
return _cap(header + "\n".join(matches))
|
||||
|
||||
return [
|
||||
StructuredTool.from_function(
|
||||
name="read_run",
|
||||
description=(
|
||||
"Read a stored scraper run or spilled tool output by line, in pages. "
|
||||
"Use the reference from a truncated tool result (e.g. 'run_<uuid>'). "
|
||||
"Each line is one result item (JSON). Page with offset/limit; prefer "
|
||||
"search_run when hunting for something specific."
|
||||
),
|
||||
coroutine=_read_run,
|
||||
),
|
||||
StructuredTool.from_function(
|
||||
name="search_run",
|
||||
description=(
|
||||
"Search a stored scraper run or spilled tool output for lines matching "
|
||||
"a substring or regular expression. Returns matching items with their "
|
||||
"line index. Cheaper than reading the whole run when you know what you "
|
||||
"are looking for."
|
||||
),
|
||||
coroutine=_search_run,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _build_matcher(pattern: str):
|
||||
"""Compile a line matcher; fall back to substring on bad/oversized regex (ReDoS guard)."""
|
||||
if len(pattern) > _MAX_PATTERN_LEN:
|
||||
return lambda line: pattern in line
|
||||
try:
|
||||
compiled = re.compile(pattern, re.IGNORECASE)
|
||||
except re.error:
|
||||
return lambda line: pattern.lower() in line.lower()
|
||||
return lambda line: compiled.search(line) is not None
|
||||
|
|
@ -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},
|
||||
)
|
||||
|
|
@ -2779,6 +2779,81 @@ class PersonalAccessToken(BaseModel, TimestampMixin):
|
|||
return not self.is_expired
|
||||
|
||||
|
||||
class Run(Base, TimestampMixin):
|
||||
"""One row per scraper-capability invocation, from either the agent door
|
||||
or the REST/API-key door.
|
||||
|
||||
Backs the user-facing Scraper-API logs and the agent's tool-boundary
|
||||
truncation: the full output lives here while the model sees only a capped
|
||||
preview plus this row's id. ``output_text`` is stored as JSONL (one item
|
||||
per line, ``exclude_none``) so ``read_run``/``search_run`` can page and grep
|
||||
by line without parsing the whole payload. Retained ~30 days via opportunistic
|
||||
bounded cleanup on insert.
|
||||
|
||||
``cost_micros`` ships nullable and unpopulated in this pass; the planned
|
||||
per-verb pricing rework will fill it. ``thread_id`` is a free-form string
|
||||
(subagent ids look like ``"2099::task:call_x"``), so it is intentionally not
|
||||
a foreign key.
|
||||
"""
|
||||
|
||||
__tablename__ = "runs"
|
||||
__allow_unmapped__ = True
|
||||
|
||||
__table_args__ = (
|
||||
Index("ix_runs_workspace_created", "workspace_id", "created_at"),
|
||||
)
|
||||
|
||||
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
||||
workspace_id = Column(
|
||||
Integer,
|
||||
ForeignKey("workspaces.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
index=True,
|
||||
)
|
||||
user_id = Column(
|
||||
UUID(as_uuid=True),
|
||||
ForeignKey("user.id", ondelete="CASCADE"),
|
||||
nullable=True,
|
||||
index=True,
|
||||
)
|
||||
thread_id = Column(String(255), nullable=True)
|
||||
capability = Column(String(100), nullable=False, index=True)
|
||||
origin = Column(String(16), nullable=False)
|
||||
status = Column(String(16), nullable=False)
|
||||
error = Column(Text, nullable=True)
|
||||
input = Column(JSONB, nullable=True)
|
||||
output_text = Column(Text, nullable=True)
|
||||
item_count = Column(Integer, nullable=False, default=0)
|
||||
char_count = Column(Integer, nullable=False, default=0)
|
||||
duration_ms = Column(Integer, nullable=True)
|
||||
cost_micros = Column(BigInteger, nullable=True)
|
||||
|
||||
|
||||
class ToolOutputSpill(Base, TimestampMixin):
|
||||
"""Internal scratch store for main-agent context-editing spills.
|
||||
|
||||
Kept separate from ``runs`` so customer-facing scraper logs stay clean.
|
||||
The full ``ToolMessage`` content that context editing evicts is written here
|
||||
and the message body is replaced with a ``spill_{id}`` placeholder the agent
|
||||
can read back via ``read_run``/``search_run``. Retained ~7 days.
|
||||
"""
|
||||
|
||||
__tablename__ = "tool_output_spills"
|
||||
__allow_unmapped__ = True
|
||||
|
||||
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
||||
workspace_id = Column(
|
||||
Integer,
|
||||
ForeignKey("workspaces.id", ondelete="CASCADE"),
|
||||
nullable=True,
|
||||
index=True,
|
||||
)
|
||||
thread_id = Column(String(255), nullable=True)
|
||||
tool_name = Column(String(255), nullable=True)
|
||||
content = Column(Text, nullable=False)
|
||||
char_count = Column(Integer, nullable=False, default=0)
|
||||
|
||||
|
||||
# Register model packages that live outside this file so their classes
|
||||
# are present in Base.metadata before configure_mappers() resolves any
|
||||
# string-based relationship() references.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue