refactor(api-slim): write each duplicated policy once (#4083)

* refactor(api-slim): replace bare tuple contracts with named result types

The project standard says "never use multi-item tuple return values -- not even
for internal/private functions", but a diff-based review can only enforce it on
new code. This sweeps the pre-existing violations where a positional contract
was actually load-bearing, and adds a structural guard so they cannot come back.

Six named types, chosen by how a mistake would surface rather than by count:

  RetainBatchResult      retain pipeline (memory_ids, usage, processed_tokens)
  ExtractionResult       fact extraction (facts, chunks, usage)
  TagClause              tag SQL (sql, params, next_param_offset)
  GraphRetrieval         graph retrieval (results, timings)
  FittedDeltaPrompt      delta prompt fitting (doc, candidate, facts, truncated)
  ValidatedOperations    delta ops (valid, skipped)
  PreparedFactEntities   entity prep (fact_texts, fact_dates, entities_per_fact)

Three defects this surfaced:

* `_streaming_retain_batch` was annotated `tuple[list[list[str]], TokenUsage]`
  while returning three values, and `retain_batch` handed that straight back as
  its own 3-tuple. Nothing caught it: tuple arity is checked nowhere and `ty`
  has `invalid-return-type` disabled. It was harmless only because every caller
  happened to unpack three.

* `_fit_structured_delta_prompt_parts` returns three `str` whose meaning is
  caller-dependent -- the refresh prompt puts the synthesis in slot 2 and new
  facts in slot 3, the retraction prompt puts still-supported facts in slot 2
  and retracted ones in slot 3. Transposing those two type-checks and builds a
  grammatical prompt telling the model to strip content resting on still-valid
  facts and keep content resting on retracted ones. Now asserted directly by
  test_retraction_prompt_does_not_transpose_surviving_and_retracted, since no
  type can express it.

* `_merge_processed_content_tokens` (the "None is contagious, bill full
  content" rule) was re-derived inline twice in memory_engine.py without the
  docstring explaining why `None + int` is a semantic and not a bug. Moved
  beside the field it governs; its docstring also cited
  `RetainResult.processed_content_tokens`, a field neither RetainResult has.

Also deduplicates the memory_units link-expansion projection, which was spelled
out ~20 times across both DB backends -- once per UNION ALL arm, again in the PG
semantic arm's GROUP BY, again in each dialect's outer re-projection. Adding a
column was a 20-edit change where a dropped column breaks union arity and a
reordered one corrupts results silently, since the arms are read positionally.

Behaviour-preserving, and checked rather than asserted:

* 18/18 generated link-expansion SQL strings byte-identical (2 dialects x 3
  builders x 3 window variants).
* 930 tag-builder invocations and 380 Python-side filter evaluations identical
  to main, across every match mode, tag set, alias, offset and nesting shape.
* 288 delta/retraction prompts identical to main, including at caps that force
  truncation.
* Guard tests enumerate each family from source and assert no member returns a
  bare tuple; all are mutation-checked by reintroducing the original defect.

Deliberately left alone: MemoryEngineInterface and mcp_tools dict returns
(extension and MCP wire contracts), memories/base.graph_links_and_entities (an
overridable store method), and tuples that are genuinely tuple-shaped -- a dict
cache key and a sort key.

Claude-Session: https://claude.ai/code/session_01UBcDzagMhuXsYZDp63i7aB

* refactor(api-slim): write each duplicated policy once

Four duplications found by auditing api-slim for logic that exists in more than
one place. One of them was already producing malformed output; the other three
are the conditions that let it happen unnoticed.

Fixes a real bug: 35 MCP error payloads were not valid JSON

The bank-id variants of the MCP tools declare `-> str` and return JSON text, and
built their error branch by interpolation:

    return f'{{"error": "{e}"}}'

which emits invalid JSON as soon as the message contains a double quote, a
backslash or a newline. PostgreSQL quotes identifiers with double quotes, so an
ordinary `relation "memory_units" does not exist` was already enough to hand the
caller something it could not parse. 35 of the 38 sites were in the bank-id half
of a duplicated tool; the session half returned a dict and was always correct --
which is exactly why no test caught it. All now go through `_error_json`, which
serializes with `json.dumps`.

Each MCP tool is written once instead of twice

Every tool is registered twice, once taking an explicit `bank_id` and returning
JSON text, once resolving the bank from the session and returning a dict.
FastMCP builds each tool's schema from the literal signature and docstring, so
those declarations genuinely have to exist twice -- but the logic did not, and
all 36 pairs had diverging bodies. The shared flow (resolve the bank, reject a
missing one, serialize, map OperationValidationError and everything else onto an
error payload) is now `_run_tool`, and each registrar writes its one engine call
once. 32 of 36 tools; the four exceptions are named in the guard test with their
reasons -- recall/reflect use different pydantic serializers in the two copies
(model_dump_json vs model_dump, which do not render datetimes identically), and
retain/sync_retain resolve a bank per content item.

`ValueError` handling is now explicit rather than incidental: tools treated it
three different ways -- a fault, a rejected input logged beside
OperationValidationError, or a normal negative answer with no log at all -- so
the wrapper takes which one applies instead of averaging them.

api/http.py: one error policy instead of 20 spellings

89 route handlers carried 20 distinct spellings of the same mapping: 72
byte-identical copies of the four-line catch-all, plus seven that logged the
traceback without the message. All 79 now call `_internal_error`, which logs the
route, the message and the traceback and returns a 500 carrying only the
message. Two handlers keep their own: recall logs its handler duration, and
retain maps MemoryDefenseAllBlockedError to a 422 and logs an input summary.
Both are named in the guard test rather than pattern-matched.

Providers: a capability that was half-declared

openai_responses and github_copilot declared `cached_prefix` and
`cached_prefix_message_count` on call/call_with_tools but never read them, and
override none of the prompt-caching methods -- so `get_or_create_cached_prefix`
returns None for them and callers can never pass one. Dead surface that read
like a supported feature; removed, which also makes them consistent with the ten
other providers that never declared it.

Verification

Behaviour is checked, not assumed. A harness registers every tool in both bank
modes against a mock engine across three scenarios (success, not-found, and an
exception whose message carries every JSON metacharacter) and captures each
tool's parameter schema, output schema, description and returned payload:
468/468 entries identical before and after, including the four that caught a
real regression mid-refactor -- routing recall and reflect through the shared
wrapper changed their serializer, so they were reverted.

New guard tests assert over each family rather than any one member, since the
member that forgets is by construction the one without a test: that no tool
builds error JSON by interpolation, that both copies of a tool declare the same
parameters and share one implementation, that a provider declares a capability's
parameters only if it implements the capability, and that no handler builds a
500 inline. Every one is mutation-checked by reintroducing the defect it
describes.
This commit is contained in:
Nicolò Boschi
2026-09-03 17:43:41 +02:00
committed by GitHub
parent 9697c69008
commit 6e206b191b
48 changed files with 2698 additions and 2102 deletions
+103 -402
View File
@@ -10,6 +10,7 @@ import json
import logging
import os
import re
import traceback
import uuid
from collections.abc import Awaitable
from contextlib import asynccontextmanager
@@ -209,6 +210,28 @@ from hindsight_api.models import RequestContext
logger = logging.getLogger(__name__)
def _internal_error(exc: Exception, where: str) -> HTTPException:
"""Log an unhandled handler exception with its traceback, and map it to a 500.
Every route's catch-all used to do this inline — 72 byte-identical copies of
the same four lines, plus seven near-copies that logged the traceback without
the message. The duplication is why the policy had already drifted: what gets
logged, and the fact that the client sees ``detail=str(exc)`` rather than a
traceback, was re-decided per route instead of once.
``traceback.format_exc()`` reads the *currently handled* exception, so this
must be called from inside an ``except`` block which is where every call
site is.
Returns the exception rather than raising it, so call sites read
``raise _internal_error(e, ...)`` and keep the ``raise`` visible at the
handler instead of hidden behind a call.
"""
logger.error(f"Error in {where}: {exc}\n\nTraceback:\n{traceback.format_exc()}")
return HTTPException(status_code=500, detail=str(exc))
# 499 is the de facto reverse-proxy status for "client closed request".
_CLIENT_CLOSED_REQUEST_STATUS_CODE = 499
@@ -4573,11 +4596,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/graph: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/graph")
@app.get(
"/v1/default/banks/{bank_id}/memories/list",
@@ -4647,11 +4666,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/memories/list: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/memories/list")
async def _require_dry_run_enabled() -> None:
"""Feature-flag gate for dry-run extraction.
@@ -4720,11 +4735,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/memories/dry-run-extract: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/memories/dry-run-extract")
@app.get(
"/v1/default/banks/{bank_id}/memories/{memory_id}",
@@ -4755,11 +4766,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/memories/{memory_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/memories/{memory_id}")
@app.patch(
"/v1/default/banks/{bank_id}/memories/{memory_id}",
@@ -4814,11 +4821,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in PATCH /v1/default/banks/{bank_id}/memories/{memory_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"PATCH /v1/default/banks/{bank_id}/memories/{memory_id}")
@app.get(
"/v1/default/banks/{bank_id}/memories/{memory_id}/history",
@@ -4849,11 +4852,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/memories/{memory_id}/history: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/memories/{memory_id}/history")
@app.post(
"/v1/default/banks/{bank_id}/memories/recall",
@@ -5221,11 +5220,7 @@ def _register_routes(app: FastAPI):
detail=str(e) or "Reflect operation timed out. Consider reducing the budget or simplifying the query.",
)
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/reflect: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/reflect")
@app.get(
"/v1/default/banks",
@@ -5253,11 +5248,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, "/v1/default/banks")
@app.get(
"/v1/default/banks/{bank_id}/stats",
@@ -5314,11 +5305,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/stats: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/stats")
@app.post(
"/v1/default/banks/{bank_id}/health/llm",
@@ -5353,11 +5340,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in POST /v1/default/banks/{bank_id}/health/llm: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"POST /v1/default/banks/{bank_id}/health/llm")
@app.get(
"/v1/default/banks/{bank_id}/stats/memories-timeseries",
@@ -5391,11 +5374,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/stats/memories-timeseries: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/stats/memories-timeseries")
@app.get(
"/v1/default/banks/{bank_id}/entities",
@@ -5427,11 +5406,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/entities: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/entities")
@app.get(
"/v1/default/banks/{bank_id}/entities/graph",
@@ -5457,11 +5432,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/entities/graph: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/entities/graph")
@app.get(
"/v1/default/banks/{bank_id}/entities/{entity_id}",
@@ -5500,11 +5471,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/entities/{entity_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/entities/{entity_id}")
@app.post(
"/v1/default/banks/{bank_id}/entities/{entity_id}/regenerate",
@@ -5574,11 +5541,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in GET /v1/default/banks/{bank_id}/mental-models: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/mental-models")
@app.get(
"/v1/default/banks/{bank_id}/mental-models/{mental_model_id}",
@@ -5616,11 +5579,7 @@ def _register_routes(app: FastAPI):
except ValueError as e:
raise HTTPException(status_code=400, detail=str(e))
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in GET /v1/default/banks/{bank_id}/mental-models/{mental_model_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/mental-models/{mental_model_id}")
@app.get(
"/v1/default/banks/{bank_id}/mental-models/{mental_model_id}/history",
@@ -5651,13 +5610,7 @@ def _register_routes(app: FastAPI):
except ValueError as e:
raise HTTPException(status_code=400, detail=str(e))
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(
f"Error in GET /v1/default/banks/{bank_id}/mental-models/{mental_model_id}/history: {error_detail}"
)
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/mental-models/{mental_model_id}/history")
@app.post(
"/v1/default/banks/{bank_id}/mental-models",
@@ -5704,11 +5657,7 @@ def _register_routes(app: FastAPI):
except OperationValidationError as e:
raise HTTPException(status_code=e.status_code, detail=e.reason)
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in POST /v1/default/banks/{bank_id}/mental-models: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"POST /v1/default/banks/{bank_id}/mental-models")
@app.post(
"/v1/default/banks/{bank_id}/mental-models/{mental_model_id}/refresh",
@@ -5742,13 +5691,7 @@ def _register_routes(app: FastAPI):
except OperationValidationError as e:
raise HTTPException(status_code=e.status_code, detail=e.reason)
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(
f"Error in POST /v1/default/banks/{bank_id}/mental-models/{mental_model_id}/refresh: {error_detail}"
)
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"POST /v1/default/banks/{bank_id}/mental-models/{mental_model_id}/refresh")
@app.post(
"/v1/default/banks/{bank_id}/mental-models/{mental_model_id}/dry-run-refresh",
@@ -5796,14 +5739,9 @@ def _register_routes(app: FastAPI):
except OperationValidationError as e:
raise HTTPException(status_code=e.status_code, detail=e.reason)
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(
f"Error in POST /v1/default/banks/{bank_id}/mental-models/{mental_model_id}/"
f"dry-run-refresh: {error_detail}"
raise _internal_error(
e, f"POST /v1/default/banks/{bank_id}/mental-models/{mental_model_id}/dry-run-refresh"
)
raise HTTPException(status_code=500, detail=str(e))
@app.post(
"/v1/default/banks/{bank_id}/mental-models/{mental_model_id}/clear",
@@ -5839,13 +5777,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(
f"Error in POST /v1/default/banks/{bank_id}/mental-models/{mental_model_id}/clear: {error_detail}"
)
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"POST /v1/default/banks/{bank_id}/mental-models/{mental_model_id}/clear")
@app.patch(
"/v1/default/banks/{bank_id}/mental-models/{mental_model_id}",
@@ -5888,11 +5820,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in PATCH /v1/default/banks/{bank_id}/mental-models/{mental_model_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"PATCH /v1/default/banks/{bank_id}/mental-models/{mental_model_id}")
@app.delete(
"/v1/default/banks/{bank_id}/mental-models/{mental_model_id}",
@@ -5924,11 +5852,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in DELETE /v1/default/banks/{bank_id}/mental-models/{mental_model_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"DELETE /v1/default/banks/{bank_id}/mental-models/{mental_model_id}")
# =========================================================================
# KNOWLEDGE BASE ENDPOINTS (folders + pages, markdown)
@@ -5961,11 +5885,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in GET /v1/default/banks/{bank_id}/knowledge-base/tree: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/knowledge-base/tree")
@app.post(
"/v1/default/banks/{bank_id}/knowledge-base/folders",
@@ -5997,11 +5917,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in POST /v1/default/banks/{bank_id}/knowledge-base/folders: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"POST /v1/default/banks/{bank_id}/knowledge-base/folders")
@app.post(
"/v1/default/banks/{bank_id}/knowledge-base/pages",
@@ -6054,11 +5970,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in POST /v1/default/banks/{bank_id}/knowledge-base/pages: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"POST /v1/default/banks/{bank_id}/knowledge-base/pages")
@app.get(
"/v1/default/banks/{bank_id}/knowledge-base/export",
@@ -6100,11 +6012,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in GET /v1/default/banks/{bank_id}/knowledge-base/export: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/knowledge-base/export")
@app.get(
"/v1/default/banks/{bank_id}/knowledge-base/search",
@@ -6137,11 +6045,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in GET /v1/default/banks/{bank_id}/knowledge-base/search: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/knowledge-base/search")
@app.get(
"/v1/default/banks/{bank_id}/knowledge-base/pages/{page_id}",
@@ -6169,11 +6073,7 @@ def _register_routes(app: FastAPI):
except OperationValidationError as e:
raise HTTPException(status_code=e.status_code, detail=e.reason)
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in GET /v1/default/banks/{bank_id}/knowledge-base/pages/{page_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/knowledge-base/pages/{page_id}")
@app.patch(
"/v1/default/banks/{bank_id}/knowledge-base/nodes/{node_id}",
@@ -6242,11 +6142,7 @@ def _register_routes(app: FastAPI):
except OperationValidationError as e:
raise HTTPException(status_code=e.status_code, detail=e.reason)
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in PATCH /v1/default/banks/{bank_id}/knowledge-base/nodes/{node_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"PATCH /v1/default/banks/{bank_id}/knowledge-base/nodes/{node_id}")
@app.delete(
"/v1/default/banks/{bank_id}/knowledge-base/nodes/{node_id}",
@@ -6273,11 +6169,7 @@ def _register_routes(app: FastAPI):
except OperationValidationError as e:
raise HTTPException(status_code=e.status_code, detail=e.reason)
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in DELETE /v1/default/banks/{bank_id}/knowledge-base/nodes/{node_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"DELETE /v1/default/banks/{bank_id}/knowledge-base/nodes/{node_id}")
# =========================================================================
# DIRECTIVES ENDPOINTS
@@ -6329,11 +6221,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in GET /v1/default/banks/{bank_id}/directives: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/directives")
@app.get(
"/v1/default/banks/{bank_id}/directives/{directive_id}",
@@ -6363,11 +6251,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in GET /v1/default/banks/{bank_id}/directives/{directive_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/directives/{directive_id}")
@app.post(
"/v1/default/banks/{bank_id}/directives",
@@ -6402,11 +6286,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in POST /v1/default/banks/{bank_id}/directives: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"POST /v1/default/banks/{bank_id}/directives")
@app.patch(
"/v1/default/banks/{bank_id}/directives/{directive_id}",
@@ -6443,11 +6323,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in PATCH /v1/default/banks/{bank_id}/directives/{directive_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"PATCH /v1/default/banks/{bank_id}/directives/{directive_id}")
@app.delete(
"/v1/default/banks/{bank_id}/directives/{directive_id}",
@@ -6477,11 +6353,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in DELETE /v1/default/banks/{bank_id}/directives/{directive_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"DELETE /v1/default/banks/{bank_id}/directives/{directive_id}")
@app.get(
"/v1/default/banks/{bank_id}/documents",
@@ -6531,11 +6403,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/documents: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/documents")
@app.get(
"/v1/default/banks/{bank_id}/documents/{document_id:path}/chunks",
@@ -6577,11 +6445,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/documents/{document_id}/chunks: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/documents/{document_id}/chunks")
@app.post(
"/v1/default/banks/{bank_id}/documents/{document_id:path}/reprocess",
@@ -6624,11 +6488,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/documents/{document_id}/reprocess: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/documents/{document_id}/reprocess")
@app.get(
"/v1/default/banks/{bank_id}/documents/{document_id:path}",
@@ -6658,11 +6518,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/documents/{document_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/documents/{document_id}")
@app.get(
"/v1/default/banks/{bank_id}/tags",
@@ -6728,11 +6584,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/tags: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/tags")
@app.get(
"/v1/default/chunks/{chunk_id:path}",
@@ -6759,11 +6611,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/chunks/{chunk_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/chunks/{chunk_id}")
@app.patch(
"/v1/default/banks/{bank_id}/documents/{document_id:path}",
@@ -6809,11 +6657,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in PATCH /v1/default/banks/{bank_id}/documents/{document_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"PATCH /v1/default/banks/{bank_id}/documents/{document_id}")
@app.delete(
"/v1/default/banks/{bank_id}/documents/{document_id:path}",
@@ -6856,11 +6700,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/documents/{document_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/documents/{document_id}")
@app.get(
"/v1/default/banks/{bank_id}/operations",
@@ -6907,11 +6747,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/operations: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/operations")
@app.get(
"/v1/default/banks/{bank_id}/operations/{operation_id}",
@@ -6949,11 +6785,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in GET /v1/default/banks/{bank_id}/operations/{operation_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/operations/{operation_id}")
@app.delete(
"/v1/default/banks/{bank_id}/operations/{operation_id}",
@@ -6984,11 +6816,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/operations/{operation_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/operations/{operation_id}")
@app.post(
"/v1/default/banks/{bank_id}/operations/{operation_id}/retry",
@@ -7018,11 +6846,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in POST /v1/default/banks/{bank_id}/operations/{operation_id}/retry: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"POST /v1/default/banks/{bank_id}/operations/{operation_id}/retry")
@app.delete(
"/v1/default/banks/{bank_id}/operations/{operation_id}/delete",
@@ -7052,13 +6876,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(
f"Error in DELETE /v1/default/banks/{bank_id}/operations/{operation_id}/delete: {error_detail}"
)
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"DELETE /v1/default/banks/{bank_id}/operations/{operation_id}/delete")
@app.get(
"/v1/default/banks/{bank_id}/profile",
@@ -7100,11 +6918,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/profile: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/profile")
@app.put(
"/v1/default/banks/{bank_id}/profile",
@@ -7145,11 +6959,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/profile: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/profile")
@app.post(
"/v1/default/banks/{bank_id}/background",
@@ -7175,11 +6985,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/background: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/background")
@app.put(
"/v1/default/banks/{bank_id}",
@@ -7224,11 +7030,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}")
@app.patch(
"/v1/default/banks/{bank_id}",
@@ -7274,11 +7076,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in PATCH /v1/default/banks/{bank_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"PATCH /v1/default/banks/{bank_id}")
@app.delete(
"/v1/default/banks/{bank_id}",
@@ -7306,11 +7104,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in DELETE /v1/default/banks/{bank_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"DELETE /v1/default/banks/{bank_id}")
# =====================================================================
# Bank Template Import / Export
@@ -7384,11 +7178,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in POST /v1/default/banks/{bank_id}/import: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"POST /v1/default/banks/{bank_id}/import")
@app.get(
"/v1/default/banks/{bank_id}/export",
@@ -7476,11 +7266,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in GET /v1/default/banks/{bank_id}/export: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/export")
# =====================================================================
# Document Transfer (Export / Import between banks — no LLM re-extraction)
@@ -7576,12 +7362,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
logger.error(
f"Error in POST /v1/default/banks/{bank_id}/document-transfer/export: {traceback.format_exc()}"
)
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"POST /v1/default/banks/{bank_id}/document-transfer/export")
@app.post(
"/v1/default/banks/{bank_id}/document-transfer",
@@ -7629,10 +7410,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
logger.error(f"Error in POST /v1/default/banks/{bank_id}/document-transfer: {traceback.format_exc()}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"POST /v1/default/banks/{bank_id}/document-transfer")
@app.get(
"/v1/default/files/download/{key:path}",
@@ -7683,10 +7461,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
logger.error(f"Error in GET /v1/default/files/download/{key}: {traceback.format_exc()}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/files/download/{key}")
@app.get(
"/v1/bank-template-schema",
@@ -7723,11 +7498,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in DELETE /v1/default/banks/{bank_id}/observations: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"DELETE /v1/default/banks/{bank_id}/observations")
@app.get(
"/v1/default/banks/{bank_id}/observations/scopes",
@@ -7752,11 +7523,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in GET /v1/default/banks/{bank_id}/observations/scopes: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/observations/scopes")
@app.post(
"/v1/default/banks/{bank_id}/consolidation/recover",
@@ -7781,11 +7548,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in POST /v1/default/banks/{bank_id}/consolidation/recover: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"POST /v1/default/banks/{bank_id}/consolidation/recover")
@app.delete(
"/v1/default/banks/{bank_id}/memories/{memory_id}/observations",
@@ -7816,13 +7579,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(
f"Error in DELETE /v1/default/banks/{bank_id}/memories/{memory_id}/observations: {error_detail}"
)
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"DELETE /v1/default/banks/{bank_id}/memories/{memory_id}/observations")
@app.get(
"/v1/default/banks/{bank_id}/config",
@@ -7848,11 +7605,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in GET /v1/default/banks/{bank_id}/config: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/config")
@app.patch(
"/v1/default/banks/{bank_id}/config",
@@ -7889,11 +7642,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in PATCH /v1/default/banks/{bank_id}/config: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"PATCH /v1/default/banks/{bank_id}/config")
@app.delete(
"/v1/default/banks/{bank_id}/config",
@@ -7920,11 +7669,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in DELETE /v1/default/banks/{bank_id}/config: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"DELETE /v1/default/banks/{bank_id}/config")
@app.post(
"/v1/default/banks/{bank_id}/consolidate",
@@ -7957,11 +7702,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in POST /v1/default/banks/{bank_id}/consolidate: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"POST /v1/default/banks/{bank_id}/consolidate")
# =========================================================================
# Webhook Endpoints
@@ -8051,11 +7792,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in POST /v1/default/banks/{bank_id}/webhooks: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"POST /v1/default/banks/{bank_id}/webhooks")
@app.get(
"/v1/default/banks/{bank_id}/webhooks",
@@ -8107,11 +7844,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in GET /v1/default/banks/{bank_id}/webhooks: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/webhooks")
@app.delete(
"/v1/default/banks/{bank_id}/webhooks/{webhook_id}",
@@ -8142,11 +7875,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in DELETE /v1/default/banks/{bank_id}/webhooks/{webhook_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"DELETE /v1/default/banks/{bank_id}/webhooks/{webhook_id}")
@app.patch(
"/v1/default/banks/{bank_id}/webhooks/{webhook_id}",
@@ -8229,11 +7958,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in PATCH /v1/default/banks/{bank_id}/webhooks/{webhook_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"PATCH /v1/default/banks/{bank_id}/webhooks/{webhook_id}")
@app.get(
"/v1/default/banks/{bank_id}/webhooks/{webhook_id}/deliveries",
@@ -8279,11 +8004,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in GET /v1/default/banks/{bank_id}/webhooks/{webhook_id}/deliveries: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/webhooks/{webhook_id}/deliveries")
@app.post(
"/v1/default/banks/{bank_id}/memories",
@@ -8645,11 +8366,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/files/retain: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/files/retain")
@app.delete(
"/v1/default/banks/{bank_id}/memories",
@@ -8677,11 +8394,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in /v1/default/banks/{bank_id}/memories: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"/v1/default/banks/{bank_id}/memories")
# ---- Audit Logs ----
# Response models live in engine/audit.py so the MemoryEngine read methods
@@ -8730,10 +8443,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
logger.error(f"Error listing audit logs: {traceback.format_exc()}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/audit-logs")
@app.get(
"/v1/default/banks/{bank_id}/audit-logs/stats",
@@ -8765,10 +8475,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
logger.error(f"Error getting audit log stats: {traceback.format_exc()}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/audit-logs/stats")
@app.get(
"/v1/default/banks/{bank_id}/llm-requests",
@@ -8825,10 +8532,7 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
logger.error(f"Error listing LLM requests: {traceback.format_exc()}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/llm-requests")
@app.get(
"/v1/default/banks/{bank_id}/llm-requests/stats",
@@ -8860,7 +8564,4 @@ def _register_routes(app: FastAPI):
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
logger.error(f"Error getting LLM request stats: {traceback.format_exc()}")
raise HTTPException(status_code=500, detail=str(e))
raise _internal_error(e, f"GET /v1/default/banks/{bank_id}/llm-requests/stats")
@@ -24,6 +24,55 @@ from typing import Any
from .base import DatabaseConnection
from .result import ResultRow
#: The ``memory_units`` columns every link-expansion arm projects, in the order the
#: arms are ``UNION ALL``-ed together. Order is part of the contract, not a style
#: choice: the arms are combined positionally, so two arms listing the same columns
#: in different orders would union cleanly and silently mis-assign every value.
MEMORY_UNIT_COLUMNS: tuple[str, ...] = (
"id",
"text",
"context",
"event_date",
"occurred_start",
"occurred_end",
"mentioned_at",
"fact_type",
"document_id",
"chunk_id",
"tags",
"proof_count",
)
def memory_unit_columns(alias: str = "", *, indent: int = 0) -> str:
"""The shared ``memory_units`` projection for link expansion, optionally alias-qualified.
Single source of truth for a list that link expansion repeats ~20 times across
both backends — once per ``UNION ALL`` arm, again in the PostgreSQL semantic
arm's ``GROUP BY``, and again in each dialect's outer re-projection. Spelling
it out per site made adding a column a 20-edit change where *every* miss is a
failure: a dropped column breaks the union arity, a reordered one corrupts the
results silently (see ``MEMORY_UNIT_COLUMNS``), and a ``GROUP BY`` left behind
is a runtime SQL error on a path only Oracle or a live recall exercises.
The score/weight/source expressions that follow the projection stay at the call
sites — they are what actually differs between the arms, and hiding them here
would trade a real distinction for a false one.
Args:
alias: Correlation name to qualify each column with (e.g. ``"mu"``). Pass
``""`` for bare columns, as subquery re-projections and ``GROUP BY``
lists need.
indent: Spaces to indent continuation lines by, so the generated SQL stays
readable in logs and ``EXPLAIN`` output.
"""
prefix = f"{alias}." if alias else ""
columns = [f"{prefix}{column}" for column in MEMORY_UNIT_COLUMNS]
# Four per line keeps the longest alias-qualified row well inside the 120-column
# limit the rest of the file is formatted to.
lines = [", ".join(columns[i : i + 4]) for i in range(0, len(columns), 4)]
return (",\n" + " " * indent).join(lines)
def document_serialization_sql(table: str, alias: str) -> str:
"""SQL predicate keeping one document to a single in-flight retain.
@@ -18,6 +18,7 @@ from .ops import (
UpdatedWindow,
bank_serialization_sql,
document_serialization_sql,
memory_unit_columns,
)
from .result import DictResultRow as ResultRow
@@ -604,9 +605,7 @@ class OracleOps(DataAccessOps):
GROUP BY t.unit_id
),
entity_expanded AS (
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
SELECT {memory_unit_columns("mu", indent=23)},
es.score, 'entity' AS source
FROM entity_scores es
JOIN {mu_table} mu ON mu.id = es.unit_id
@@ -647,9 +646,7 @@ class OracleOps(DataAccessOps):
GROUP BY id
),
semantic_expanded AS (
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
SELECT {memory_unit_columns("mu", indent=23)},
ss.score, 'semantic' AS source
FROM sem_scores ss
JOIN {mu_table} mu ON mu.id = ss.id
@@ -658,9 +655,7 @@ class OracleOps(DataAccessOps):
),
causal_ranked AS (
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
{memory_unit_columns("mu", indent=20)},
ml.weight AS score,
'causal' AS source,
ROW_NUMBER() OVER (PARTITION BY mu.id ORDER BY ml.weight DESC) AS rn_
@@ -672,8 +667,8 @@ class OracleOps(DataAccessOps):
{window.clause("mu")}
),
causal_expanded AS (
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at,
fact_type, document_id, chunk_id, tags, proof_count, score, source
SELECT {memory_unit_columns(indent=23)},
score, source
FROM causal_ranked WHERE rn_ = 1
ORDER BY score DESC
FETCH FIRST $3 ROWS ONLY
@@ -751,9 +746,7 @@ class OracleOps(DataAccessOps):
),
observation_entity_expanded AS (
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
{memory_unit_columns("mu", indent=20)},
(SELECT COUNT(*)
FROM {obs_sources_table} os2
WHERE os2.observation_id = mu.id
@@ -792,9 +785,7 @@ class OracleOps(DataAccessOps):
GROUP BY id
),
semantic_expanded AS (
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
SELECT {memory_unit_columns("mu", indent=23)},
ss.score, 'semantic' AS source
FROM sem_scores ss
JOIN {mu_table} mu ON mu.id = ss.id
@@ -803,9 +794,8 @@ class OracleOps(DataAccessOps):
),
causal_ranked AS (
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
mu.chunk_id, mu.tags, mu.proof_count, ml.weight AS score,
{memory_unit_columns("mu", indent=20)},
ml.weight AS score,
'causal' AS source,
ROW_NUMBER() OVER (PARTITION BY mu.id ORDER BY ml.weight DESC) AS rn_
FROM {ml_table} ml
@@ -816,8 +806,8 @@ class OracleOps(DataAccessOps):
{window.clause("mu")}
),
causal_expanded AS (
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at,
fact_type, document_id, chunk_id, tags, proof_count, score, source
SELECT {memory_unit_columns(indent=23)},
score, source
FROM causal_ranked WHERE rn_ = 1
ORDER BY score DESC
FETCH FIRST $2 ROWS ONLY
@@ -16,6 +16,7 @@ from .ops import (
UpdatedWindow,
bank_serialization_sql,
document_serialization_sql,
memory_unit_columns,
)
from .result import ResultRow
@@ -792,9 +793,7 @@ class PostgreSQLOps(DataAccessOps):
WHERE ue.unit_id = ANY($1::uuid[])
),
entity_expanded AS (
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
SELECT {memory_unit_columns("mu", indent=23)},
COUNT(DISTINCT se.entity_id)::float AS score,
'entity'::text AS source
FROM seed_entities se
@@ -833,16 +832,12 @@ class PostgreSQLOps(DataAccessOps):
return f"""
semantic_expanded AS (
SELECT
id, text, context, event_date, occurred_start,
occurred_end, mentioned_at,
fact_type, document_id, chunk_id, tags, proof_count,
{memory_unit_columns(indent=20)},
MAX(weight) AS score,
'semantic'::text AS source
FROM (
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
{memory_unit_columns("mu", indent=24)},
ml.weight
FROM {ml_table} ml
JOIN {mu_table} mu ON mu.id = ml.to_unit_id
@@ -853,9 +848,7 @@ class PostgreSQLOps(DataAccessOps):
{window.clause("mu")}
UNION ALL
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
{memory_unit_columns("mu", indent=24)},
ml.weight
FROM {ml_table} ml
JOIN {mu_table} mu ON mu.id = ml.from_unit_id
@@ -865,17 +858,13 @@ class PostgreSQLOps(DataAccessOps):
AND mu.id != ALL($1::uuid[])
{window.clause("mu")}
) sem_raw
GROUP BY id, text, context, event_date, occurred_start,
occurred_end, mentioned_at,
fact_type, document_id, chunk_id, tags, proof_count
GROUP BY {memory_unit_columns(indent=25)}
ORDER BY score DESC
LIMIT $3
),
causal_expanded AS (
SELECT DISTINCT ON (mu.id)
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
{memory_unit_columns("mu", indent=20)},
ml.weight AS score,
'causal'::text AS source
FROM {ml_table} ml
@@ -979,9 +968,7 @@ class PostgreSQLOps(DataAccessOps):
),
candidates AS (
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
{memory_unit_columns("mu", indent=20)},
mu.source_memory_ids
FROM {mu_table} mu, connected_array ca
WHERE mu.fact_type = 'observation'
@@ -999,9 +986,7 @@ class PostgreSQLOps(DataAccessOps):
),
observation_entity_expanded AS (
SELECT
c.id, c.text, c.context, c.event_date, c.occurred_start,
c.occurred_end, c.mentioned_at,
c.fact_type, c.document_id, c.chunk_id, c.tags, c.proof_count,
{memory_unit_columns("c", indent=20)},
sc.score,
'entity'::text AS source
FROM candidates c
@@ -1013,39 +998,33 @@ class PostgreSQLOps(DataAccessOps):
-- DISTINCT ON for causal, hardcoded to fact_type='observation'.
semantic_expanded AS (
SELECT
id, text, context, event_date, occurred_start,
occurred_end, mentioned_at,
fact_type, document_id, chunk_id, tags, proof_count,
{memory_unit_columns(indent=20)},
MAX(weight) AS score,
'semantic'::text AS source
FROM (
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
mu.chunk_id, mu.tags, mu.proof_count, ml.weight
SELECT {memory_unit_columns("mu", indent=27)},
ml.weight
FROM {ml_table} ml JOIN {mu_table} mu ON mu.id = ml.to_unit_id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
AND mu.id != ALL($1::uuid[])
{window.clause("mu")}
UNION ALL
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
mu.chunk_id, mu.tags, mu.proof_count, ml.weight
SELECT {memory_unit_columns("mu", indent=27)},
ml.weight
FROM {ml_table} ml JOIN {mu_table} mu ON mu.id = ml.from_unit_id
WHERE ml.to_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
AND mu.id != ALL($1::uuid[])
{window.clause("mu")}
) sem_raw
GROUP BY id, text, context, event_date, occurred_start, occurred_end,
mentioned_at, fact_type, document_id, chunk_id, tags, proof_count
GROUP BY {memory_unit_columns(indent=25)}
ORDER BY score DESC LIMIT $2
),
causal_expanded AS (
SELECT DISTINCT ON (mu.id)
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
mu.chunk_id, mu.tags, mu.proof_count, ml.weight AS score, 'causal'::text AS source
{memory_unit_columns("mu", indent=20)},
ml.weight AS score, 'causal'::text AS source
FROM {ml_table} ml JOIN {mu_table} mu ON ml.to_unit_id = mu.id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
@@ -202,7 +202,10 @@ async def list_memory_units(
)
if tags:
tags_clause, tags_params, next_param = build_tags_where_clause(tags, param_count + 1, "", tags_match)
built = build_tags_where_clause(tags, param_count + 1, "", tags_match)
tags_clause = built.sql
tags_params = built.params
next_param = built.next_param_offset
if tags_clause:
query_conditions.append(tags_clause.removeprefix("AND "))
query_params.extend(tags_params)
@@ -238,7 +238,9 @@ async def scan_memories(
# Compound tag groups (AND/OR/NOT trees), AND-ed on. Also owns its `AND` prefix and appends
# one bind param per leaf; empty/absent groups yield no clause and no params.
groups_clause, group_params, _ = build_tag_groups_where_clause(tag_groups, param_offset=len(params) + 1)
built = build_tag_groups_where_clause(tag_groups, param_offset=len(params) + 1)
groups_clause = built.sql
group_params = built.params
params.extend(group_params)
offset = _decode_page_token(page_token) + max(int(skip or 0), 0)
@@ -507,12 +509,17 @@ async def any_memory_updated_since(
params: list[Any] = [bank_id, since]
where = ["bank_id = $1", "updated_at > $2"]
tag_clause, tag_params, next_param = build_tags_where_clause(tags, param_offset=len(params) + 1, match=tags_match)
built = build_tags_where_clause(tags, param_offset=len(params) + 1, match=tags_match)
tag_clause = built.sql
tag_params = built.params
next_param = built.next_param_offset
if tag_clause:
where.append(tag_clause.removeprefix("AND "))
params.extend(tag_params)
group_clause, group_params, _ = build_tag_groups_where_clause(tag_groups, param_offset=next_param)
built = build_tag_groups_where_clause(tag_groups, param_offset=next_param)
group_clause = built.sql
group_params = built.params
if group_clause:
where.append(group_clause.removeprefix("AND "))
params.extend(group_params)
@@ -608,9 +615,8 @@ async def any_memory_updated_since_batch(
# in the group whatever its actual tags. The `mu.` alias is required, not
# cosmetic: both sides of the join expose a `tags` column, and an unqualified
# one would resolve by scoping rules rather than by intent.
tag_clause, _, _ = build_tags_where_clause(
scope.tags, table_alias="mu.", match=scope.tags_match, value_expr="s.tags"
)
built = build_tags_where_clause(scope.tags, table_alias="mu.", match=scope.tags_match, value_expr="s.tags")
tag_clause = built.sql
by_clause.setdefault(tag_clause, []).append(scope)
for tag_clause, group in by_clause.items():
@@ -177,7 +177,7 @@ class PostgresMemories(MemoriesExtension):
assert retriever is not None # only resolved when the arm is on
async def _run_graph(ft: str) -> list:
results, _timing = await retriever.retrieve(
retrieved = await retriever.retrieve(
pool=pool,
query_embedding_str=query_embedding,
bank_id=bank_id,
@@ -191,7 +191,8 @@ class PostgresMemories(MemoriesExtension):
created_before=created_before,
preselected_semantic_seeds=semantic_bm25[ft].graph_seeds,
)
return results
# Timings are diagnostics for the perf harness; this path drops them.
return retrieved.results
# gather preserves input order, so zip back onto fact_types positionally.
graph_lists = await asyncio.gather(*[_run_graph(ft) for ft in fact_types])
@@ -563,7 +563,7 @@ from .response_models import (
from .response_models import RecallResult as RecallResultModel
from .retain import bank_utils, embedding_utils
from .retain.fold import FoldMemberRef
from .retain.types import RetainContentDict
from .retain.types import RetainBatchResult, RetainContentDict, merge_processed_content_tokens
from .search.reranking import CrossEncoderReranker, apply_combined_scoring
from .search.tag_resolution import MAX_VOCABULARY, TagResolutionError, needs_resolution, resolve_tag_groups
from .search.tags import TagGroup, TagsMatch, build_tag_groups_where_clause, build_tags_where_clause
@@ -5237,7 +5237,7 @@ class MemoryEngine(MemoryEngineInterface):
else:
group_outbox_callback = outbox_callback if is_last_group else None
group_result, group_usage, group_processed = await self._retain_batch_async_internal(
group_outcome = await self._retain_batch_async_internal(
bank_id=bank_id,
contents=group.contents,
request_context=request_context,
@@ -5250,13 +5250,12 @@ class MemoryEngine(MemoryEngineInterface):
outbox_callback=group_outbox_callback,
)
for local_idx, origin_idx in enumerate(group.origins):
if local_idx < len(group_result):
result[origin_idx] = group_result[local_idx]
total_usage = total_usage + group_usage
if total_processed_content_tokens is None or group_processed is None:
total_processed_content_tokens = None
else:
total_processed_content_tokens = total_processed_content_tokens + group_processed
if local_idx < len(group_outcome.memory_ids):
result[origin_idx] = group_outcome.memory_ids[local_idx]
total_usage = total_usage + group_outcome.usage
total_processed_content_tokens = merge_processed_content_tokens(
total_processed_content_tokens, group_outcome.processed_content_tokens
)
# A cancelled run (bank deleted mid-flight) skips the completion side
# effects, mirroring the pre-grouping early return from the sub-batch loop.
@@ -5658,7 +5657,7 @@ class MemoryEngine(MemoryEngineInterface):
collected: list[_SubBatchOutcome] = []
async def _run_sub(idx: int, contents_, origins_, offset_, is_last_, body_, body_hash_):
r, u, pr = await self._retain_batch_async_internal(
sub_outcome = await self._retain_batch_async_internal(
bank_id=bank_id,
contents=contents_,
request_context=request_context,
@@ -5678,7 +5677,13 @@ class MemoryEngine(MemoryEngineInterface):
body_accum=body_accum,
retain_session=retain_session,
)
return _SubBatchOutcome(index=idx, origins=origins_, results=r, usage=u, processed=pr)
return _SubBatchOutcome(
index=idx,
origins=origins_,
results=sub_outcome.memory_ids,
usage=sub_outcome.usage,
processed=sub_outcome.processed_content_tokens,
)
try:
for sub in sub_batch_stream:
@@ -5812,10 +5817,9 @@ class MemoryEngine(MemoryEngineInterface):
if sub_idx < len(sub_results):
per_input_results[origin_idx].extend(sub_results[sub_idx])
total_usage = total_usage + sub_usage
if total_processed_content_tokens is None or sub_processed is None:
total_processed_content_tokens = None
else:
total_processed_content_tokens = total_processed_content_tokens + sub_processed
total_processed_content_tokens = merge_processed_content_tokens(
total_processed_content_tokens, sub_processed
)
# Per-sub-batch progress is intentionally not written here: the streaming
# retain pipeline emits finer-grained "storing N/total chunks" snapshots
# via progress_callback as each sub-batch's chunks commit.
@@ -5832,7 +5836,7 @@ class MemoryEngine(MemoryEngineInterface):
# In a try/finally for the same reason the split path's commit is: a retain that fails
# part-way must not discard what its earlier parts already produced.
try:
result, total_usage, total_processed_content_tokens = await self._retain_batch_async_internal(
sub_batch_outcome = await self._retain_batch_async_internal(
bank_id=bank_id,
contents=contents,
request_context=request_context,
@@ -5850,6 +5854,9 @@ class MemoryEngine(MemoryEngineInterface):
if retain_session is not None:
async with _retain_timing_mod.timed("store.commit"):
await retain_session.commit()
result = sub_batch_outcome.memory_ids
total_usage = sub_batch_outcome.usage
total_processed_content_tokens = sub_batch_outcome.processed_content_tokens
# Progress for this path is emitted by the streaming pipeline as
# "storing N/total chunks" via progress_callback (see _retain_batch_async_internal).
@@ -5898,7 +5905,7 @@ class MemoryEngine(MemoryEngineInterface):
chunk_index_offset: int = 0,
body_accum: "dict[str, DocumentBodyAccumulator] | None" = None,
retain_session=None,
) -> tuple[list[list[str]], "TokenUsage", int | None]:
) -> "RetainBatchResult":
"""
Internal method for batch processing without chunking logic.
@@ -5972,9 +5979,8 @@ class MemoryEngine(MemoryEngineInterface):
audit_logger=self._audit_logger,
)
# Map the created facts onto this retain's trace so the trace view can
# show which memories the ingestion produced. result[0] is the
# per-content-item list of created unit ids (see retain_batch).
created_ids = [uid for group in result[0] for uid in group]
# show which memories the ingestion produced.
created_ids = [uid for group in result.memory_ids for uid in group]
# Fire-and-forget: the mapping is patched on a background task so it
# never adds latency to the retain response.
self._llm_recorder.attach_memory_ids(trace_context_of(retain_llm), created=created_ids)
@@ -5987,7 +5993,7 @@ class MemoryEngine(MemoryEngineInterface):
# the operation-level retry re-runs the whole submission anyway.
_APPEND_CONFLICT_ATTEMPTS = 3
async def _retain_batch_with_append_retry(self, **kwargs) -> tuple[list[list[str]], "TokenUsage", int | None]:
async def _retain_batch_with_append_retry(self, **kwargs) -> "RetainBatchResult":
"""Run ``orchestrator.retain_batch``, redoing an append that lost its race.
``update_mode="append"`` reads the stored document, concatenates onto it
@@ -11054,9 +11060,10 @@ class MemoryEngine(MemoryEngineInterface):
query_conditions.append(f"id ILIKE ${param_count}")
query_params.append(f"%{search_query}%")
tags_clause, tags_params, next_param = build_tags_where_clause(
tags, param_offset=param_count + 1, match=tags_match
)
built = build_tags_where_clause(tags, param_offset=param_count + 1, match=tags_match)
tags_clause = built.sql
tags_params = built.params
next_param = built.next_param_offset
query_params.extend(tags_params)
param_count = next_param - 1 # next_param is next available; convert to last used
@@ -16725,19 +16732,24 @@ class MemoryEngine(MemoryEngineInterface):
params: list[Any] = [bank_id]
where = ["bank_id = $1"]
tag_clause, tag_params, next_param = build_tags_where_clause(
built = build_tags_where_clause(
tag_filtering.tags,
param_offset=len(params) + 1,
match=tag_filtering.tags_match,
)
tag_clause = built.sql
tag_params = built.params
next_param = built.next_param_offset
if tag_clause:
where.append(tag_clause.removeprefix("AND "))
params.extend(tag_params)
group_clause, group_params, _ = build_tag_groups_where_clause(
built = build_tag_groups_where_clause(
tag_filtering.tag_groups,
param_offset=next_param,
)
group_clause = built.sql
group_params = built.params
if group_clause:
where.append(group_clause.removeprefix("AND "))
params.extend(group_params)
@@ -18037,18 +18049,20 @@ class MemoryEngine(MemoryEngineInterface):
# both filters apply independently — each wrapped in the untagged-OR rule —
# so the directive set is the intersection of what either filter would admit.
if tags:
tags_clause, tags_params, param_idx = build_tags_where_clause(
tags=tags, param_offset=param_idx, table_alias="", match=tags_match
)
built = build_tags_where_clause(tags=tags, param_offset=param_idx, table_alias="", match=tags_match)
tags_clause = built.sql
tags_params = built.params
param_idx = built.next_param_offset
if tags_clause:
# Always include untagged directives; tagged ones must match the reflect tags
scoped_clause = tags_clause.replace("AND ", "", 1)
filters.append(f"((tags IS NULL OR tags = '{{}}') OR ({scoped_clause}))")
params.extend(tags_params)
if tag_groups:
groups_clause, groups_params, param_idx = build_tag_groups_where_clause(
tag_groups, param_offset=param_idx
)
built = build_tag_groups_where_clause(tag_groups, param_offset=param_idx)
groups_clause = built.sql
groups_params = built.params
param_idx = built.next_param_offset
if groups_clause:
# Same untagged-OR rule as the flat-tags branch above.
scoped_clause = groups_clause.replace("AND ", "", 1)
@@ -594,7 +594,6 @@ class GitHubCopilotLLM(LLMInterface):
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
cached_prefix: str | None = None,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> Any:
start_time = time.time()
@@ -702,8 +701,6 @@ class GitHubCopilotLLM(LLMInterface):
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
cached_prefix: str | None = None,
cached_prefix_message_count: int = 0,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> LLMToolCallResult:
start_time = time.time()
@@ -439,7 +439,6 @@ class OpenAIResponsesLLM(LLMInterface):
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
cached_prefix: str | None = None,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> Any:
"""Make a Responses API call with retry logic (see ``LLMInterface.call``)."""
@@ -533,8 +532,6 @@ class OpenAIResponsesLLM(LLMInterface):
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
cached_prefix: str | None = None,
cached_prefix_message_count: int = 0,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> LLMToolCallResult:
"""Make a Responses API call with tools (see ``LLMInterface.call_with_tools``).
@@ -38,6 +38,7 @@ from __future__ import annotations
import json
import logging
import re
from dataclasses import dataclass
from typing import Annotated, Any, Literal, Union
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError, field_validator
@@ -202,7 +203,22 @@ _OPERATION_ADAPTER: TypeAdapter[Operation] = TypeAdapter(Operation)
_BODY_FIELDS = ("text", "blocks")
def _validate_operations_list(raw_ops: Any) -> tuple[list[Operation], list[dict[str, Any]]]:
@dataclass(frozen=True)
class ValidatedOperations:
"""The surviving and rejected halves of one LLM operation list.
Both halves are lists and both are non-empty in the interesting cases, so as a
bare tuple they were transposable without a type error and the caller feeds
them to ``_finalize_operations``, which raises ``DeltaAllOpsInvalidError`` off
the *valid* half being empty while the *skipped* half is not. Swapped, a clean
refresh would abort and a fully-malformed one would be applied.
"""
valid: list[Operation]
skipped: list[dict[str, Any]]
def _validate_operations_list(raw_ops: Any) -> ValidatedOperations:
"""Validate each operation independently; drop invalid ops instead of failing the batch."""
if not isinstance(raw_ops, list):
raise TypeError(f"operations must be a list, got {type(raw_ops)!r}")
@@ -218,7 +234,7 @@ def _validate_operations_list(raw_ops: Any) -> tuple[list[Operation], list[dict[
i,
exc.errors(include_url=False),
)
return valid, skipped
return ValidatedOperations(valid, skipped)
class DeltaOperationList(BaseModel):
@@ -289,7 +305,8 @@ def parse_delta_operation_list(raw: Any) -> DeltaOperationList:
return raw
if isinstance(raw, dict):
ops_raw = raw.get("operations", [])
valid, skipped = _validate_operations_list(ops_raw)
validated = _validate_operations_list(ops_raw)
valid, skipped = validated.valid, validated.skipped
if skipped:
logger.info(
"[STRUCTURED_DELTA] parsed %s op(s), skipped %s invalid op(s) from dict payload",
@@ -320,7 +337,8 @@ def parse_delta_operation_list(raw: Any) -> DeltaOperationList:
last_error = ValueError("delta payload must be an object with an operations array")
continue
try:
valid, skipped = _validate_operations_list(payload["operations"])
validated = _validate_operations_list(payload["operations"])
valid, skipped = validated.valid, validated.skipped
except TypeError as exc:
last_error = exc
continue
@@ -8,6 +8,7 @@ The reflect agent uses hierarchical retrieval:
"""
import json
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any
@@ -1065,6 +1066,29 @@ def _truncate_prompt_text(text: str, max_tokens: int) -> str:
return truncate_to_tokens(text, max_tokens).text
@dataclass(frozen=True)
class FittedDeltaPrompt:
"""The three oversized prompt sections after budget-fitting, and whether any was cut.
The sections are named for their *slots*, not their contents, because both
callers reuse this fitter with different material in them: the refresh prompt
puts the synthesis in ``candidate`` and the new facts in ``facts``, while the
retraction prompt puts the still-supported facts in ``candidate`` and the
retracted ones in ``facts``.
That reuse is why these must not be a bare tuple. All three are ``str``, so
transposing two positions type-checks perfectly and for the retraction
caller, swapping ``candidate`` and ``facts`` builds a prompt instructing the
model to strip content resting on still-valid facts while keeping content
resting on retracted ones. Nothing downstream could detect it.
"""
document_json: str
candidate: str
facts: str
truncated: bool
def _fit_structured_delta_prompt_parts(
*,
source_query: str,
@@ -1074,7 +1098,7 @@ def _fit_structured_delta_prompt_parts(
budget_hint: str,
task_footer: str,
max_input_tokens: int,
) -> tuple[str, str, str, bool]:
) -> FittedDeltaPrompt:
"""Shrink large prompt sections to fit within max_input_tokens (tokenizer estimate)."""
from .tokenization import count_prompt_tokens
@@ -1098,7 +1122,7 @@ def _fit_structured_delta_prompt_parts(
candidate = _truncate_prompt_text(candidate_markdown, cand_budget)
facts_body = _truncate_prompt_text(facts_block, facts_budget)
truncated = doc_json != current_document_json or candidate != candidate_markdown or facts_body != facts_block
return doc_json, candidate, facts_body, truncated
return FittedDeltaPrompt(doc_json, candidate, facts_body, truncated)
def build_structured_delta_prompt(
@@ -1176,7 +1200,7 @@ def build_structured_delta_prompt(
"as needed. Preserve unchanged sections and blocks by not mentioning them."
)
input_cap = max_input_tokens if max_input_tokens is not None else _STRUCTURED_DELTA_DEFAULT_MAX_INPUT_TOKENS
doc_json, candidate, facts_body, input_truncated = _fit_structured_delta_prompt_parts(
fitted = _fit_structured_delta_prompt_parts(
source_query=source_query,
current_document_json=current_document_json,
candidate_markdown=candidate_markdown,
@@ -1186,7 +1210,7 @@ def build_structured_delta_prompt(
max_input_tokens=input_cap,
)
truncation_note = ""
if input_truncated:
if fitted.truncated:
truncation_note = (
"\n\n*Note: Document, synthesis, or facts were truncated to fit the model "
"context window. Prefer minimal, high-leverage operations.*"
@@ -1195,10 +1219,10 @@ def build_structured_delta_prompt(
return (
f"## Topic\n{source_query}\n\n"
f"## CURRENT DOCUMENT (apply ops to this; copy section and block ids from it verbatim)\n"
f"```json\n{doc_json}\n```\n\n"
f"```json\n{fitted.document_json}\n```\n\n"
f"## NEW INFORMATION SYNTHESIS (context for how new facts relate to the topic)\n"
f"```markdown\n{candidate}\n```\n\n"
f"## SUPPORTING FACTS (new since last refresh — integrate these)\n{facts_body}"
f"```markdown\n{fitted.candidate}\n```\n\n"
f"## SUPPORTING FACTS (new since last refresh — integrate these)\n{fitted.facts}"
f"{document_hint}{budget_hint}{truncation_note}\n\n"
f"{task_footer}"
)
@@ -1326,7 +1350,7 @@ def build_structured_retraction_prompt(
# lists standing in for synthesis + facts), same budget split, so a large
# document cannot push the retracted list out of the window.
input_cap = max_input_tokens if max_input_tokens is not None else _STRUCTURED_DELTA_DEFAULT_MAX_INPUT_TOKENS
doc_json, surviving_body, retracted_body, input_truncated = _fit_structured_delta_prompt_parts(
fitted = _fit_structured_delta_prompt_parts(
source_query=source_query,
current_document_json=current_document_json,
candidate_markdown=surviving_block,
@@ -1336,7 +1360,7 @@ def build_structured_retraction_prompt(
max_input_tokens=input_cap,
)
truncation_note = ""
if input_truncated:
if fitted.truncated:
truncation_note = (
"\n\n*Note: Document or fact lists were truncated to fit the model context "
"window. Prefer minimal, high-leverage operations, and keep anything you "
@@ -1346,11 +1370,11 @@ def build_structured_retraction_prompt(
return (
f"## Topic\n{source_query}\n\n"
f"## CURRENT DOCUMENT (apply ops to this; reference section ids as listed)\n"
f"```json\n{doc_json}\n```\n\n"
f"```json\n{fitted.document_json}\n```\n\n"
f"## STILL-SUPPORTED FACTS (these remain valid — do not remove content resting on them)\n"
f"{surviving_body}\n\n"
f"{fitted.candidate}\n\n"
f"## RETRACTED FACTS (no longer in the memory bank — remove content resting on these)\n"
f"{retracted_body}"
f"{fitted.facts}"
f"{budget_hint}{truncation_note}\n\n"
f"{task_footer}"
)
@@ -131,12 +131,18 @@ async def tool_search_mental_models(
# skip the filter, or mental models would see every scope while the other
# reflect retrieval tools correctly see only untagged data.
if tags or tags_match == "exact":
tag_clause, tag_params, next_param = build_tags_where_clause(tags, param_offset=next_param, match=tags_match)
built = build_tags_where_clause(tags, param_offset=next_param, match=tags_match)
tag_clause = built.sql
tag_params = built.params
next_param = built.next_param_offset
filters += f" {tag_clause}"
params.extend(tag_params)
if tag_groups:
groups_clause, groups_params, next_param = build_tag_groups_where_clause(tag_groups, next_param)
built = build_tag_groups_where_clause(tag_groups, next_param)
groups_clause = built.sql
groups_params = built.params
next_param = built.next_param_offset
filters += f" {groups_clause}"
params.extend(groups_params)
@@ -5,6 +5,7 @@ Handles entity extraction and resolution for stored facts.
"""
import logging
from dataclasses import dataclass
from . import link_utils
from .types import EntityResolutionResult, ProcessedFact, UserEntities
@@ -12,10 +13,27 @@ from .types import EntityResolutionResult, ProcessedFact, UserEntities
logger = logging.getLogger(__name__)
@dataclass(frozen=True)
class PreparedFactEntities:
"""Per-fact inputs to entity resolution, all three aligned to ``facts`` by index.
The alignment is the whole contract: resolution zips these together, so a
caller that reordered or filtered one list without the others would silently
attach entities to the wrong fact. Three same-shaped lists in a bare tuple
made that easy to do and impossible to see and most callers wanted only
``entities_per_fact``, so they wrote ``_texts, _dates, entities_per_fact``
and depended on the other two staying exactly where they were.
"""
fact_texts: list[str]
fact_dates: list
entities_per_fact: list[list[dict]]
def _prepare_facts_for_entity_processing(
facts: list[ProcessedFact],
user_entities_per_content: dict[int, UserEntities] | None = None,
) -> tuple[list[str], list, list[list[dict]]]:
) -> PreparedFactEntities:
"""
Extract fact texts, dates, and merged entity lists from ProcessedFact objects.
@@ -25,7 +43,7 @@ def _prepare_facts_for_entity_processing(
without turning off resolution for the extractor's (#3479).
Returns:
Tuple of (fact_texts, fact_dates, entities_per_fact)
A PreparedFactEntities whose three lists are index-aligned to ``facts``.
"""
user_entities_per_content = user_entities_per_content or {}
@@ -60,7 +78,7 @@ def _prepare_facts_for_entity_processing(
entities_per_fact.append(llm_entities)
return fact_texts, fact_dates, entities_per_fact
return PreparedFactEntities(fact_texts, fact_dates, entities_per_fact)
async def resolve_entities(
@@ -99,7 +117,10 @@ async def resolve_entities(
if len(unit_ids) != len(facts):
raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and facts ({len(facts)})")
fact_texts, fact_dates, entities_per_fact = _prepare_facts_for_entity_processing(facts, user_entities_per_content)
prepared = _prepare_facts_for_entity_processing(facts, user_entities_per_content)
fact_texts = prepared.fact_texts
fact_dates = prepared.fact_dates
entities_per_fact = prepared.entities_per_fact
return await link_utils.resolve_entities_only(
entity_resolver,
@@ -2210,7 +2210,7 @@ async def extract_facts_from_text(
# Import types for the orchestration layer (note: ExtractedFact here is different from the Pydantic model above)
from .types import CausalRelation as CausalRelationType
from .types import ChunkMetadata, RetainContent
from .types import ChunkMetadata, ExtractionResult, RetainContent
from .types import ExtractedFact as ExtractedFactType
logger = logging.getLogger(__name__)
@@ -2259,7 +2259,7 @@ async def extract_facts_from_contents_batch_api(
pool=None,
operation_id: str | None = None,
schema: str | None = None,
) -> tuple[list[ExtractedFactType], list[ChunkMetadata], TokenUsage]:
) -> ExtractionResult:
"""
Extract facts using LLM Batch API (OpenAI/Groq).
@@ -2275,10 +2275,10 @@ async def extract_facts_from_contents_batch_api(
schema: Database schema (for multi-tenant support)
Returns:
Tuple of (extracted_facts, chunks_metadata, usage)
An ExtractionResult carrying the facts, their chunk metadata, and token usage.
"""
if not contents:
return [], [], TokenUsage()
return ExtractionResult([], [], TokenUsage())
logger.info(f"Using Batch API for fact extraction ({len(contents)} contents)")
@@ -2396,7 +2396,7 @@ async def extract_facts_from_contents_batch_api(
)
if not batch_requests and not batch_id: # No requests and not resuming
return [], [], TokenUsage()
return ExtractionResult([], [], TokenUsage())
# Step 2: Submit batch (skip if resuming)
if not batch_id:
@@ -2799,13 +2799,13 @@ async def extract_facts_from_contents_batch_api(
logger.info(f"Batch API extracted {len(extracted_facts)} facts from {len(all_chunks_info)} chunks")
return extracted_facts, chunks_metadata, total_usage
return ExtractionResult(extracted_facts, chunks_metadata, total_usage)
def _extract_facts_chunks(
contents: list[RetainContent],
config,
) -> tuple[list[ExtractedFactType], list[ChunkMetadata], TokenUsage]:
) -> ExtractionResult:
"""
chunks mode: no LLM call, no entity extraction.
@@ -2849,7 +2849,7 @@ def _extract_facts_chunks(
global_chunk_idx += 1
_add_temporal_offsets(extracted_facts, contents)
return extracted_facts, chunks_metadata, TokenUsage()
return ExtractionResult(extracted_facts, chunks_metadata, TokenUsage())
async def extract_facts_from_contents(
@@ -2859,7 +2859,7 @@ async def extract_facts_from_contents(
pool=None,
operation_id: str | None = None,
schema: str | None = None,
) -> tuple[list[ExtractedFactType], list[ChunkMetadata], TokenUsage]:
) -> ExtractionResult:
"""
Extract facts from multiple content items in parallel.
@@ -2880,10 +2880,10 @@ async def extract_facts_from_contents(
schema: Database schema (passed to batch API for multi-tenant support)
Returns:
Tuple of (extracted_facts, chunks_metadata, usage)
An ExtractionResult carrying the facts, their chunk metadata, and token usage.
"""
if not contents:
return [], [], TokenUsage()
return ExtractionResult([], [], TokenUsage())
# chunks mode: skip LLM entirely, store each chunk as-is
# Must come before the batch-API check so no LLM queue/locks are acquired
@@ -2996,7 +2996,7 @@ async def extract_facts_from_contents(
# Step 6: Auto-tag facts from label groups with tag=True
_inject_label_tags(extracted_facts, config)
return extracted_facts, chunks_metadata, total_usage
return ExtractionResult(extracted_facts, chunks_metadata, total_usage)
def _collapse_to_verbatim(facts: list[ExtractedFactType], chunks: list[ChunkMetadata]) -> list[ExtractedFactType]:
@@ -212,21 +212,6 @@ async def _audit_memory_defense(
audit_logger.log_fire_and_forget(entry)
def _merge_processed_content_tokens(a: int | None, b: int | None) -> int | None:
"""Combine the processed-content-tokens signal across sub-results.
Semantics (see RetainResult.processed_content_tokens):
* None means "this part of the retain did not go through chunk-level
dedup" — i.e. the entire submitted payload was processed. If any
sub-result is None, the aggregate is None so callers conservatively
bill the full content.
* Otherwise, accumulate the int values.
"""
if a is None or b is None:
return None
return a + b
def _count_delta_content_tokens(delta_contents: list["RetainContent"]) -> int:
"""Sum content + context tokens across the chunk items that were
actually fed into the extraction pipeline on a partial-delta retain.
@@ -295,9 +280,11 @@ from .types import (
Phase1Result,
ProcessedFact,
ResolvedEntity,
RetainBatchResult,
RetainContent,
RetainContentDict,
UserEntities,
merge_processed_content_tokens,
)
logger = logging.getLogger(__name__)
@@ -320,6 +307,26 @@ class _ProcessedFactBatch:
retained_index_by_original: list[int | None]
@dataclass(frozen=True)
class _EmbeddedExtraction:
"""Facts extracted *and* embedded — what the retain paths share before storage.
``extracted_facts`` and ``processed_facts`` are both here on purpose and are
not interchangeable: ``ProcessedFact.from_extracted_fact`` drops degenerate
facts, so ``processed_facts`` can be shorter, and the two are re-aligned via
``_ProcessedFactBatch.retained_index_by_original``. Callers that need the
original chunk positions (causal-relation remapping) must read
``extracted_facts``; callers that write rows must read ``processed_facts``.
Returning them as a bare 4-tuple made picking the wrong one a silent
positional mistake.
"""
extracted_facts: list[ExtractedFact]
processed_facts: list[ProcessedFact]
chunks: list[ChunkMetadata]
usage: TokenUsage
async def _record_retain_document_outcome(pool: Any, bank_id: str, document_id: str, units_created: int) -> None:
"""Emit the per-document retain outcome metric.
@@ -682,9 +689,8 @@ async def _streaming_session_retain(
for idx, content in enumerate(batch_contents)
if getattr(content, "entities", None)
}
_texts, _dates, entities_per_fact = entity_processing._prepare_facts_for_entity_processing(
batch_processed, user_entities_per_content
)
prepared = entity_processing._prepare_facts_for_entity_processing(batch_processed, user_entities_per_content)
entities_per_fact = prepared.entities_per_fact
names = {
(unit_ids or [])[i]: [e["text"] for e in entities_per_fact[i]]
for i in range(min(len(unit_ids or []), len(entities_per_fact)))
@@ -800,9 +806,8 @@ async def _streaming_store_owned_retain(
for idx, content in enumerate(batch_contents)
if getattr(content, "entities", None)
}
_texts, _dates, entities_per_fact = entity_processing._prepare_facts_for_entity_processing(
batch_processed, user_entities_per_content
)
prepared = entity_processing._prepare_facts_for_entity_processing(batch_processed, user_entities_per_content)
entities_per_fact = prepared.entities_per_fact
unit_entity_names = {
unit_ids[i]: [e["text"] for e in entities_per_fact[i]]
for i in range(min(len(unit_ids), len(entities_per_fact)))
@@ -984,9 +989,10 @@ async def _delta_store_owned_write(
for idx, content in enumerate(delta_contents)
if getattr(content, "entities", None)
}
_t, _d, entities_per_fact = entity_processing._prepare_facts_for_entity_processing(
prepared = entity_processing._prepare_facts_for_entity_processing(
processed_facts, user_entities_per_content
)
entities_per_fact = prepared.entities_per_fact
unit_entity_names = {
unit_ids[i]: [e["text"] for e in entities_per_fact[i]]
for i in range(min(len(unit_ids), len(entities_per_fact)))
@@ -1038,13 +1044,8 @@ async def _extract_and_embed(
pool: Any = None,
operation_id: str | None = None,
schema: str | None = None,
) -> tuple[list, list[ProcessedFact], list[ChunkMetadata], TokenUsage]:
"""
Shared pipeline: extract facts from contents and generate embeddings.
Returns:
Tuple of (extracted_facts, processed_facts, chunks_metadata, usage)
"""
) -> _EmbeddedExtraction:
"""Shared pipeline: extract facts from contents and generate embeddings."""
set_stage("retain.extract_and_embed")
step_start = time.time()
# No narrator: extraction takes none from this path at all. A "Narrator: {name}" line is
@@ -1055,16 +1056,17 @@ async def _extract_and_embed(
# that never mentioned them (#3962). A caller that genuinely wants to name the speaker says
# so in the item's `context`, which extraction already reads and which the dry-run
# `agent_name` override is deprecated in favour of.
extracted_facts, chunks, usage = await fact_extraction.extract_facts_from_contents(
extraction = await fact_extraction.extract_facts_from_contents(
contents, llm_config, config, pool, operation_id, schema
)
extracted_facts, chunks, usage = extraction.facts, extraction.chunks, extraction.usage
log_buffer.append(
f" Extract facts: {len(extracted_facts)} facts, {len(chunks)} chunks "
f"from {len(contents)} contents in {time.time() - step_start:.3f}s"
)
if not extracted_facts:
return extracted_facts, [], chunks, usage
return _EmbeddedExtraction(extracted_facts, [], chunks, usage)
if fact_type_override:
for fact in extracted_facts:
@@ -1078,7 +1080,7 @@ async def _extract_and_embed(
fact_batch = _process_extracted_facts(extracted_facts, embeddings)
return fact_batch.extracted_facts, fact_batch.processed_facts, chunks, usage
return _EmbeddedExtraction(fact_batch.extracted_facts, fact_batch.processed_facts, chunks, usage)
def _remap_causal_relations(
@@ -1190,7 +1192,7 @@ async def retain_batch(
webhook_manager: Any = None,
memory_defense_extension: "MemoryDefenseExtension | None" = None,
audit_logger: Any = None,
) -> tuple[list[list[str]], TokenUsage, int | None]:
) -> RetainBatchResult:
"""
Process a batch of content through the retain pipeline.
@@ -1300,7 +1302,7 @@ async def retain_batch(
outbox_callback_factory(group_dicts) if outbox_callback_factory is not None else outbox_callback
)
group_ids, group_usage, group_processed = await retain_batch(
group_result = await retain_batch(
pool=pool,
embeddings_model=embeddings_model,
llm_config=llm_config,
@@ -1338,16 +1340,18 @@ async def retain_batch(
# usage totals are not safe to accumulate from several tasks at once. The driver
# below merges them in one place, in group order, so the result does not depend on
# which group happened to finish first.
return doc_key, group_ids, group_usage, group_processed
return doc_key, group_result
group_results = await asyncio.gather(*(_run_group(k, gd, gc) for k, (gd, gc) in groups.items()))
for doc_key, group_ids, group_usage, group_processed in group_results:
for doc_key, group_result in group_results:
for group_idx, orig_idx in enumerate(original_indices[doc_key]):
if group_idx < len(group_ids):
result_unit_ids[orig_idx] = group_ids[group_idx]
total_usage = total_usage + group_usage
total_processed_tokens = _merge_processed_content_tokens(total_processed_tokens, group_processed)
return result_unit_ids, total_usage, total_processed_tokens
if group_idx < len(group_result.memory_ids):
result_unit_ids[orig_idx] = group_result.memory_ids[group_idx]
total_usage = total_usage + group_result.usage
total_processed_tokens = merge_processed_content_tokens(
total_processed_tokens, group_result.processed_content_tokens
)
return RetainBatchResult(result_unit_ids, total_usage, total_processed_tokens)
# --- Memory Defense pre-extraction screening ---
# Delegate to the loaded extension. `config` is a resolved HindsightConfig
@@ -1421,7 +1425,7 @@ async def retain_batch(
contents_dicts = [contents_dicts[i] for i in _surviving]
# If nothing survives, return empty results immediately.
if not contents:
return [[] for _ in contents_dicts], TokenUsage(), 0
return RetainBatchResult([[] for _ in contents_dicts], TokenUsage(), 0)
# Resolve effective document_id early so both delta and streaming paths
# can find existing chunks from a prior attempt. On retry, a generated
@@ -1641,7 +1645,7 @@ async def retain_batch(
logger.info("\n" + "\n".join(log_buffer) + "\n")
# No new content was processed — report 0 so callers can skip
# billing cleanly instead of falling back to full-content billing.
return [[] for _ in contents], TokenUsage(), 0
return RetainBatchResult([[] for _ in contents], TokenUsage(), 0)
# --- Delta retain: check if we can skip unchanged chunks ---
#
@@ -2208,7 +2212,7 @@ async def _streaming_retain_batch(
append_base_hash: str | None = None,
append_base_watermark: int | None = None,
force_reextract: bool = False,
) -> tuple[list[list[str]], TokenUsage]:
) -> RetainBatchResult:
"""
Process a large document in streaming mini-batches to bound memory usage.
@@ -2501,7 +2505,7 @@ async def _streaming_retain_batch(
meta_token = set_call_metadata({"document_id": effective_doc_id})
try:
extracted, processed, chunk_meta, usage = await _extract_and_embed(
embedded = await _extract_and_embed(
[content],
llm_config,
config,
@@ -2515,6 +2519,10 @@ async def _streaming_retain_batch(
)
finally:
reset_call_metadata(meta_token)
extracted = embedded.extracted_facts
processed = embedded.processed_facts
chunk_meta = embedded.chunks
usage = embedded.usage
# Reserve before queueing, so a producer running ahead of a slow write path
# waits here instead of piling extracted facts up behind the queue. Extraction
# for chunks already in flight continues; only the handover is throttled.
@@ -3365,7 +3373,7 @@ async def _streaming_retain_batch(
# The streaming path doesn't compute per-chunk content-hash dedup in
# a way that lets us report a partial-processed tokens count — signal
# ``None`` so callers bill against the full submitted payload.
return result_unit_ids, total_usage, None
return RetainBatchResult(result_unit_ids, total_usage, None)
# ---------------------------------------------------------------------------
@@ -3428,7 +3436,7 @@ async def _try_delta_retain(
# `document_body_override`, which an append fills with only the new tail.
delta_full_body: str | None = None,
append_base_hash: str | None = None,
) -> tuple[list[list[str]], TokenUsage, int | None] | None:
) -> RetainBatchResult | None:
"""
Attempt delta retain for a document upsert. Returns result tuple if delta
was performed, or None to fall back to full retain.
@@ -3732,7 +3740,7 @@ async def _try_delta_retain(
meta_token = set_call_metadata({"document_id": effective_doc_id})
try:
extracted_facts, processed_facts, new_chunk_metadata, usage = await _extract_and_embed(
embedded = await _extract_and_embed(
delta_contents,
llm_config,
config,
@@ -3746,6 +3754,10 @@ async def _try_delta_retain(
)
finally:
reset_call_metadata(meta_token)
extracted_facts = embedded.extracted_facts
processed_facts = embedded.processed_facts
new_chunk_metadata = embedded.chunks
usage = embedded.usage
# Database transaction
result_unit_ids: list[list[str]] = []
@@ -3984,7 +3996,7 @@ async def _try_delta_retain(
# changed/new chunks (see ``_build_delta_contents``) — i.e. exactly what
# the LLM pipeline saw this call. Unchanged chunks contribute zero.
processed_tokens = _count_delta_content_tokens(delta_contents)
return result_unit_ids, usage, processed_tokens
return RetainBatchResult(result_unit_ids, usage, processed_tokens)
async def _delta_metadata_only(
@@ -4001,7 +4013,7 @@ async def _delta_metadata_only(
document_body_override: str | None = None,
config: Any = None,
expected_content_hash: str | None = None,
) -> tuple[list[list[str]], TokenUsage, int] | None:
) -> RetainBatchResult | None:
"""Handle the case where no chunks changed — just update document metadata and tags."""
from ..memories import get_memories as _get_memories_meta
@@ -4062,7 +4074,7 @@ async def _delta_metadata_only(
total_time = time.time() - start_time
log_buffer.append(f"DELTA RETAIN (no changes): metadata updated in {total_time:.3f}s")
logger.info("\n" + "\n".join(log_buffer) + "\n")
return [[] for _ in contents], TokenUsage(), 0
return RetainBatchResult([[] for _ in contents], TokenUsage(), 0)
async with acquire_with_retry(pool) as conn:
async with conn.transaction():
@@ -4112,7 +4124,7 @@ async def _delta_metadata_only(
# content tokens so callers can bill accordingly (a caller that's been
# told ``0`` knows the retain was a pure metadata update and should
# charge nothing for content).
return [[] for _ in contents], TokenUsage(), 0
return RetainBatchResult([[] for _ in contents], TokenUsage(), 0)
# ---------------------------------------------------------------------------
@@ -17,6 +17,7 @@ from uuid import UUID
import numpy as np
from ..metadata_utils import drop_null_values
from ..response_models import TokenUsage
logger = logging.getLogger(__name__)
@@ -207,6 +208,83 @@ class ChunkMetadata:
chunk_index: int # Global chunk index across all contents
@dataclass(frozen=True)
class RetainBatchResult:
"""What one pass of the retain pipeline produced.
Every entry point into the pipeline returns this ``retain_batch`` and the
three paths it delegates to (``_streaming_retain_batch``, ``_try_delta_retain``,
``_delta_metadata_only``), plus the engine wrappers around them. They used to
return a bare 3-tuple each, and the arity had already drifted:
``_streaming_retain_batch`` was annotated ``tuple[list[list[str]], TokenUsage]``
while returning three values, and ``retain_batch`` handed that straight back as
its own 3-tuple. Nothing caught it a tuple's shape is checked nowhere, and
``ty`` has ``invalid-return-type`` disabled so the declared contract and the
real one simply disagreed until someone unpacked two names and got a
``ValueError`` at runtime. Naming the fields is what makes that mismatch
impossible rather than merely unlikely.
"""
memory_ids: list[list[str]]
"""Created memory-unit ids, one inner list per submitted content item, in order."""
usage: TokenUsage
"""LLM tokens consumed by this pass. Merged with ``+`` across concurrent groups."""
processed_content_tokens: int | None
"""Content+context tokens that actually reached extraction.
``0`` when nothing was re-extracted (a delta whose chunks all matched), and
``None`` when the path does not account for it (streaming, which spans many
sub-batches). ``None`` and ``0`` are therefore *not* interchangeable: the
former means "unknown", the latter "known to be nothing".
"""
@dataclass(frozen=True)
class ExtractionResult:
"""What one fact-extraction pass produced, whatever route it took.
The three extraction entry points the LLM fan-out
(``extract_facts_from_contents``), the Batch API route
(``extract_facts_from_contents_batch_api``), and chunks mode
(``_extract_facts_chunks``, no LLM at all) are interchangeable by design:
``extract_facts_from_contents`` dispatches to the other two and returns their
result unchanged. They therefore have to agree on their output exactly, which
is precisely what three separately-maintained 3-tuples could not guarantee.
``facts`` and ``chunks`` are positionally related: each fact's
``chunk_index`` indexes into ``chunks``, so the two lists must come from the
same pass and cannot be sourced independently.
"""
facts: list["ExtractedFact"]
"""Extracted facts, in chunk order, carrying their ``content_index``/``chunk_index``."""
chunks: list["ChunkMetadata"]
"""One entry per chunk the pass saw, including chunks that yielded no facts."""
usage: TokenUsage
"""LLM tokens consumed. Zero for chunks mode, which makes no model call."""
def merge_processed_content_tokens(a: int | None, b: int | None) -> int | None:
"""Combine ``RetainBatchResult.processed_content_tokens`` across sub-results.
``None`` is contagious: it means "this part of the retain did not go through
chunk-level dedup", so the aggregate is unknown and callers must
conservatively bill the full content. Only when *both* sides are known does
the total mean anything, and then it is their sum.
Lives beside the field it governs because the rule is not obvious from the
types ``None + int`` looks like a bug to fix rather than a semantic to
preserve, and it was previously re-derived inline at each merge site.
"""
if a is None or b is None:
return None
return a + b
@dataclass
class EntityRef:
"""
@@ -11,7 +11,7 @@ from abc import ABC, abstractmethod
from datetime import datetime
from .tags import TagGroup, TagsMatch
from .types import GraphRetrievalTimings, RetrievalResult
from .types import GraphRetrieval, RetrievalResult
logger = logging.getLogger(__name__)
@@ -47,7 +47,7 @@ class GraphRetriever(ABC):
created_after: datetime | None = None, # Only include memory_units created after this time
created_before: datetime | None = None, # Only include memory_units created before this time
preselected_semantic_seeds: list[RetrievalResult] | None = None,
) -> tuple[list[RetrievalResult], GraphRetrievalTimings | None]:
) -> GraphRetrieval:
"""
Retrieve relevant facts via graph traversal.
@@ -39,7 +39,7 @@ from ..db_utils import acquire_with_retry
from ..memory_engine import fq_table
from .graph_retrieval import GraphRetriever
from .tags import TagGroup, TagsMatch, filter_results_by_tag_groups, filter_results_by_tags
from .types import GraphRetrievalTimings, RetrievalResult
from .types import GraphRetrieval, GraphRetrievalTimings, RetrievalResult
logger = logging.getLogger(__name__)
@@ -64,7 +64,9 @@ async def _find_semantic_seeds(
tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match)
tag_groups_param_start = 6 + (1 if tags else 0)
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
built = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
groups_clause = built.sql
groups_params = built.params
# created_after/created_before filter `updated_at`, matching the other recall arms
# (see retrieval.py) so a window narrows every arm the same way.
@@ -141,7 +143,7 @@ class LinkExpansionRetriever(GraphRetriever):
created_after: "datetime | None" = None,
created_before: "datetime | None" = None,
preselected_semantic_seeds: list[RetrievalResult] | None = None,
) -> tuple[list[RetrievalResult], GraphRetrievalTimings | None]:
) -> GraphRetrieval:
"""
Retrieve facts by expanding links from seeds.
@@ -194,7 +196,7 @@ class LinkExpansionRetriever(GraphRetriever):
)
if not all_seeds:
return [], timings
return GraphRetrieval([], timings)
seed_ids = list({s.id for s in all_seeds})
timings.pattern_count = len(seed_ids)
@@ -290,7 +292,7 @@ class LinkExpansionRetriever(GraphRetriever):
f"in {timings.traverse * 1000:.1f}ms (query: {timings.edge_load_time * 1000:.1f}ms)"
)
return results, timings
return GraphRetrieval(results, timings)
async def _expand_combined(
self,
@@ -236,7 +236,9 @@ async def retrieve_semantic_bm25_combined_sql(
# tag_groups params start immediately after the tags param slot
tag_groups_param_start = tags_param_idx + (1 if tags else 0)
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
built = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
groups_clause = built.sql
groups_params = built.params
# --- created_after/created_before time range filter (appended after tags/groups) ---
# The bounds are named for creation but filter `updated_at` — "memories that changed
@@ -335,7 +337,8 @@ async def retrieve_semantic_bm25_combined_sql(
fb_tags_idx = 3
fb_tags_clause = build_tags_where_clause_simple(tags, fb_tags_idx, match=tags_match)
fb_groups_start = fb_tags_idx + (1 if tags else 0)
fb_groups_clause, _, _ = build_tag_groups_where_clause(tag_groups, fb_groups_start)
built = build_tag_groups_where_clause(tag_groups, fb_groups_start)
fb_groups_clause = built.sql
fb_next_idx = fb_groups_start + len(groups_params)
fb_updated_clause = ""
if created_after is not None:
@@ -504,7 +507,9 @@ async def retrieve_temporal_combined_sql(
# the backend on execute, but `unnest` is not). Mirrors retrieve_semantic_bm25_combined_sql.
tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match)
tag_groups_param_start = 6 + (1 if tags else 0)
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
built = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
groups_clause = built.sql
groups_params = built.params
# created_after/created_before time range filter (after tags/groups) — filters
# `updated_at`, as above.
@@ -665,9 +670,9 @@ async def retrieve_temporal_combined_sql(
# Build tags clause for spreading (use param 7 since 1-6 are used)
spreading_tags_clause = build_tags_where_clause_simple(tags, 7, table_alias="mu.", match=tags_match)
spreading_groups_param_start = 7 + (1 if tags else 0)
spreading_groups_clause, spreading_groups_params, _ = build_tag_groups_where_clause(
tag_groups, spreading_groups_param_start, table_alias="mu."
)
built = build_tag_groups_where_clause(tag_groups, spreading_groups_param_start, table_alias="mu.")
spreading_groups_clause = built.sql
spreading_groups_params = built.params
# The window has to be repeated here, not just on the entry-point query
# above: spreading walks temporal/causal links outward, so an in-window
# entry point would otherwise pull out-of-window neighbours into results.
@@ -24,6 +24,7 @@ EXACT matching: Memory matches only if its tag set EQUALS the request tag set (o
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Annotated, Literal
from pydantic import BaseModel, ConfigDict, Field
@@ -36,30 +37,56 @@ TagsMatch = Literal["any", "all", "any_strict", "all_strict", "exact"]
TagResolution = Literal["exact", "fuzzy"]
def _parse_tags_match(match: TagsMatch) -> tuple[str, bool]:
"""
Parse TagsMatch into operator and include_untagged flag.
@dataclass(frozen=True)
class TagMatchSemantics:
"""How one ``TagsMatch`` mode compares a tags column against a value."""
Returns:
Tuple of (operator, include_untagged)
- operator: "&&" for any/any_strict, "@>" for all/all_strict
- include_untagged: True for any/all, False for any_strict/all_strict
operator: str
"""``&&`` (overlap) for any/any_strict, ``@>`` (contains) for all/all_strict."""
include_untagged: bool
"""True for any/all — untagged rows are visible; False for the ``_strict`` modes."""
@dataclass(frozen=True)
class TagClause:
"""A tag filter rendered to SQL, with the bind bookkeeping it implies.
The three parts are one unit: ``sql`` embeds ``$n`` placeholders that only
line up if ``params`` is appended to the caller's list in this order *and*
``next_param_offset`` is threaded into whatever builds the next clause.
Returned as a bare 3-tuple, that protocol was invisible a caller that
dropped the offset (several use ``_``) or bound the params out of order got
no error, just a query silently reading the wrong placeholder.
"""
sql: str
"""The clause, starting with ``AND`` for the top-level builders, or empty for no filter."""
params: list = field(default_factory=list)
"""Values to bind, in placeholder order. Empty when the clause needs no binds."""
next_param_offset: int = 1
"""First unused placeholder number. Unchanged when nothing was bound."""
def _parse_tags_match(match: TagsMatch) -> TagMatchSemantics:
"""Parse TagsMatch into the operator and untagged-visibility it implies."""
if match == "any":
return "&&", True
return TagMatchSemantics("&&", include_untagged=True)
elif match == "all":
return "@>", True
return TagMatchSemantics("@>", include_untagged=True)
elif match == "any_strict":
return "&&", False
return TagMatchSemantics("&&", include_untagged=False)
elif match == "all_strict":
return "@>", False
return TagMatchSemantics("@>", include_untagged=False)
elif match == "exact":
# Set equality is handled by the callers via `@> AND <@`; the operator
# here is unused. Untagged rows never equal a non-empty scope.
return "@>", False
return TagMatchSemantics("@>", include_untagged=False)
else:
# Default to "any" behavior
return "&&", True
return TagMatchSemantics("&&", include_untagged=True)
def build_tags_where_clause(
@@ -68,7 +95,7 @@ def build_tags_where_clause(
table_alias: str = "",
match: TagsMatch = "any",
value_expr: str | None = None,
) -> tuple[str, list, int]:
) -> TagClause:
"""
Build a SQL WHERE clause for filtering by tags.
@@ -91,14 +118,11 @@ def build_tags_where_clause(
it must be a literal the caller controls, never user input.
Returns:
Tuple of (sql_clause, params, next_param_offset):
- sql_clause: SQL WHERE clause string
- params: List of parameter values to bind
- next_param_offset: Next available parameter number
A TagClause carrying the SQL, its bind values, and the next free placeholder.
Example:
>>> clause, params, next_offset = build_tags_where_clause(['user_a'], 3, 'mu.', 'any_strict')
>>> print(clause) # "AND mu.tags IS NOT NULL AND mu.tags != '{}' AND mu.tags && $3"
>>> built = build_tags_where_clause(['user_a'], 3, 'mu.', 'any_strict')
>>> print(built.sql) # "AND mu.tags IS NOT NULL AND mu.tags != '{}' AND mu.tags && $3"
"""
column = f"{table_alias}tags" if table_alias else "tags"
# Every branch below matches the column against one right-hand side. Naming it
@@ -113,27 +137,27 @@ def build_tags_where_clause(
if match == "exact" and not tags:
# Empty/absent scope = global/untagged: match only untagged rows. No bind param
# needed (callers gate the param on truthy `tags`, so none is appended).
return f"AND ({column} IS NULL OR {column} = '{{}}')", [], param_offset
return TagClause(f"AND ({column} IS NULL OR {column} = '{{}}')", [], param_offset)
if not tags:
return "", [], param_offset
return TagClause("", [], param_offset)
if match == "exact":
# Set equality (order-independent): superset AND subset. Untagged rows
# (empty array) never satisfy `@>` of a non-empty scope, so they're excluded.
clause = f"AND ({column} @> {value} AND {column} <@ {value})"
return clause, bound, next_offset
return TagClause(clause, bound, next_offset)
operator, include_untagged = _parse_tags_match(match)
semantics = _parse_tags_match(match)
if include_untagged:
if semantics.include_untagged:
# Include untagged memories (NULL or empty array) OR matching tags
clause = f"AND ({column} IS NULL OR {column} = '{{}}' OR {column} {operator} {value})"
clause = f"AND ({column} IS NULL OR {column} = '{{}}' OR {column} {semantics.operator} {value})"
else:
# Strict: only memories with matching tags (exclude NULL and empty)
clause = f"AND {column} IS NOT NULL AND {column} != '{{}}' AND {column} {operator} {value}"
clause = f"AND {column} IS NOT NULL AND {column} != '{{}}' AND {column} {semantics.operator} {value}"
return clause, bound, next_offset
return TagClause(clause, bound, next_offset)
def build_tags_where_clause_simple(
@@ -172,14 +196,14 @@ def build_tags_where_clause_simple(
# (empty array) never satisfy `@>` of a non-empty scope, so they're excluded.
return f"AND ({column} @> ${param_num} AND {column} <@ ${param_num})"
operator, include_untagged = _parse_tags_match(match)
semantics = _parse_tags_match(match)
if include_untagged:
if semantics.include_untagged:
# Include untagged memories (NULL or empty array) OR matching tags
return f"AND ({column} IS NULL OR {column} = '{{}}' OR {column} {operator} ${param_num})"
return f"AND ({column} IS NULL OR {column} = '{{}}' OR {column} {semantics.operator} ${param_num})"
else:
# Strict: only memories with matching tags (exclude NULL and empty)
return f"AND {column} IS NOT NULL AND {column} != '{{}}' AND {column} {operator} ${param_num}"
return f"AND {column} IS NOT NULL AND {column} != '{{}}' AND {column} {semantics.operator} ${param_num}"
def filter_results_by_tags(
@@ -207,7 +231,7 @@ def filter_results_by_tags(
if not tags:
return results
_, include_untagged = _parse_tags_match(match)
include_untagged = _parse_tags_match(match).include_untagged
is_any_match = match in ("any", "any_strict")
tags_set = set(tags)
@@ -302,64 +326,64 @@ def _build_group_clause(
group: TagGroup,
param_offset: int,
table_alias: str,
) -> tuple[str, list, int]:
"""
Recursively build an inner SQL clause (no leading AND/OR) for a single TagGroup.
) -> TagClause:
"""Recursively build an inner SQL clause (no leading AND/OR) for a single TagGroup.
Returns:
(inner_clause, params, next_param_offset)
``TagClause.sql`` here carries no leading ``AND`` the caller joins the parts.
"""
if isinstance(group, TagGroupLeaf):
column = f"{table_alias}tags" if table_alias else "tags"
if group.match == "exact":
if len(group.tags) == 0:
# Empty scope = global/untagged: match only untagged rows (no bind param).
return f"({column} IS NULL OR {column} = '{{}}')", [], param_offset
return TagClause(f"({column} IS NULL OR {column} = '{{}}')", [], param_offset)
clause = f"({column} @> ${param_offset} AND {column} <@ ${param_offset})"
return clause, [group.tags], param_offset + 1
operator, include_untagged = _parse_tags_match(group.match)
if include_untagged:
clause = f"({column} IS NULL OR {column} = '{{}}' OR {column} {operator} ${param_offset})"
return TagClause(clause, [group.tags], param_offset + 1)
semantics = _parse_tags_match(group.match)
if semantics.include_untagged:
clause = f"({column} IS NULL OR {column} = '{{}}' OR {column} {semantics.operator} ${param_offset})"
else:
clause = f"({column} IS NOT NULL AND {column} != '{{}}' AND {column} {operator} ${param_offset})"
return clause, [group.tags], param_offset + 1
clause = f"({column} IS NOT NULL AND {column} != '{{}}' AND {column} {semantics.operator} ${param_offset})"
return TagClause(clause, [group.tags], param_offset + 1)
elif isinstance(group, TagGroupAnd):
parts = []
params: list = []
offset = param_offset
for child in group.filters:
child_clause, child_params, offset = _build_group_clause(child, offset, table_alias)
parts.append(child_clause)
params.extend(child_params)
built = _build_group_clause(child, offset, table_alias)
offset = built.next_param_offset
parts.append(built.sql)
params.extend(built.params)
inner = " AND ".join(parts)
return f"({inner})", params, offset
return TagClause(f"({inner})", params, offset)
elif isinstance(group, TagGroupOr):
parts = []
params = []
offset = param_offset
for child in group.filters:
child_clause, child_params, offset = _build_group_clause(child, offset, table_alias)
parts.append(child_clause)
params.extend(child_params)
built = _build_group_clause(child, offset, table_alias)
offset = built.next_param_offset
parts.append(built.sql)
params.extend(built.params)
inner = " OR ".join(parts)
return f"({inner})", params, offset
return TagClause(f"({inner})", params, offset)
elif isinstance(group, TagGroupNot):
child_clause, child_params, next_offset = _build_group_clause(group.filter, param_offset, table_alias)
return f"NOT {child_clause}", child_params, next_offset
built = _build_group_clause(group.filter, param_offset, table_alias)
return TagClause(f"NOT {built.sql}", built.params, built.next_param_offset)
else:
# Should never happen with proper Pydantic validation
return "", [], param_offset
return TagClause("", [], param_offset)
def build_tag_groups_where_clause(
tag_groups: list[TagGroup] | None,
param_offset: int,
table_alias: str = "",
) -> tuple[str, list, int]:
) -> TagClause:
"""
Build a SQL WHERE clause for compound tag group filtering.
@@ -372,30 +396,29 @@ def build_tag_groups_where_clause(
table_alias: Optional table alias prefix (e.g., "mu." for "memory_units mu").
Returns:
Tuple of (sql_clause, params, next_param_offset):
- sql_clause: SQL WHERE clause string starting with "AND" (or empty string)
- params: List of parameter values to bind (one per leaf node)
- next_param_offset: Next available parameter number
A TagClause whose ``sql`` starts with "AND" (or is empty), carrying one bind
value per leaf node.
Example:
>>> groups = [TagGroupLeaf(tags=["user:alice"], match="all_strict")]
>>> clause, params, next_offset = build_tag_groups_where_clause(groups, 3)
>>> print(clause) # "AND (tags IS NOT NULL AND tags != '{}' AND tags @> $3)"
>>> built = build_tag_groups_where_clause(groups, 3)
>>> print(built.sql) # "AND (tags IS NOT NULL AND tags != '{}' AND tags @> $3)"
"""
if not tag_groups:
return "", [], param_offset
return TagClause("", [], param_offset)
all_params: list = []
all_clauses: list[str] = []
offset = param_offset
for group in tag_groups:
inner_clause, group_params, offset = _build_group_clause(group, offset, table_alias)
all_clauses.append(inner_clause)
all_params.extend(group_params)
built = _build_group_clause(group, offset, table_alias)
offset = built.next_param_offset
all_clauses.append(built.sql)
all_params.extend(built.params)
combined = " AND ".join(all_clauses)
return f"AND {combined}", all_params, offset
return TagClause(f"AND {combined}", all_params, offset)
# =============================================================================
@@ -420,7 +443,7 @@ def _match_group(result: object, group: TagGroup) -> bool:
if group.match == "exact" and len(group.tags) == 0:
# Empty scope = global/untagged: match only untagged results.
return is_untagged
_, include_untagged = _parse_tags_match(group.match)
include_untagged = _parse_tags_match(group.match).include_untagged
is_any_match = group.match in ("any", "any_strict")
tags_set = set(group.tags)
@@ -28,6 +28,21 @@ class GraphRetrievalTimings:
hop_details: list[dict] = field(default_factory=list)
@dataclass(frozen=True)
class GraphRetrieval:
"""What one graph-retrieval strategy returned, plus how long it took getting there.
``timings`` is diagnostics only every caller but the perf path discards it,
and it is ``None`` when instrumentation is off. Pairing it with the results in
a bare tuple meant the interesting half was always the one you had to remember
came first; a strategy that returned them the other way round would type-check
identically against ``tuple[list, X | None]`` at every implementation.
"""
results: list["RetrievalResult"]
timings: "GraphRetrievalTimings | None" = None
@dataclass
class RetrievalResult:
"""
File diff suppressed because it is too large Load Diff
@@ -598,11 +598,10 @@ async def test_retain_outcome_metadata_records_zero_counts(memory, request_conte
"""Completed retain operations expose explicit zero outcome counters."""
from hindsight_api.engine.response_models import TokenUsage
from hindsight_api.engine.retain import fact_extraction
from hindsight_api.engine.retain.types import ExtractionResult
async def empty_extract_facts_from_contents(
*args: object, **kwargs: object
) -> tuple[list[object], list[object], TokenUsage]:
return [], [], TokenUsage()
async def empty_extract_facts_from_contents(*args: object, **kwargs: object) -> ExtractionResult:
return ExtractionResult([], [], TokenUsage())
monkeypatch.setattr(fact_extraction, "extract_facts_from_contents", empty_extract_facts_from_contents)
@@ -634,10 +633,10 @@ async def test_all_degenerate_facts_still_persist_document_chunks(memory, reques
"""Filtering every extracted fact must not turn an extracted chunk into the zero-extraction fast path."""
from hindsight_api.engine.response_models import TokenUsage
from hindsight_api.engine.retain import fact_extraction
from hindsight_api.engine.retain.types import ChunkMetadata, ExtractedFact
from hindsight_api.engine.retain.types import ChunkMetadata, ExtractedFact, ExtractionResult
async def degenerate_extract_facts_from_contents(contents, *_args, **_kwargs):
return (
return ExtractionResult(
[ExtractedFact(fact_text="...", fact_type="world", content_index=0, chunk_index=0)],
[ChunkMetadata(chunk_text=contents[0].content, fact_count=1, content_index=0, chunk_index=0)],
TokenUsage(),
@@ -669,7 +668,7 @@ async def test_streaming_offsets_chunk_local_causal_fact_indices(memory, request
"""Causal targets from independently extracted chunks must stay within their source chunk."""
from hindsight_api.engine.response_models import TokenUsage
from hindsight_api.engine.retain import fact_extraction
from hindsight_api.engine.retain.types import CausalRelation, ChunkMetadata, ExtractedFact
from hindsight_api.engine.retain.types import CausalRelation, ChunkMetadata, ExtractedFact, ExtractionResult
chunks = ["first-streaming-chunk", "second-streaming-chunk"]
# Patch the generator, not `chunk_text`: retain streams its chunks since #3756, and
@@ -678,7 +677,7 @@ async def test_streaming_offsets_chunk_local_causal_fact_indices(memory, request
async def extract_chunk_facts(contents, *_args, **_kwargs):
chunk_text = contents[0].content
return (
return ExtractionResult(
[
ExtractedFact(fact_text=f"{chunk_text} cause", fact_type="world", chunk_index=0),
ExtractedFact(
@@ -738,7 +737,7 @@ async def test_degenerate_fact_preserves_later_chunk_provenance(memory, request_
"""
from hindsight_api.engine.response_models import TokenUsage
from hindsight_api.engine.retain import fact_extraction
from hindsight_api.engine.retain.types import ChunkMetadata, ExtractedFact
from hindsight_api.engine.retain.types import ChunkMetadata, ExtractedFact, ExtractionResult
chunks = ["chunk-zero-source", "chunk-one-source"]
# Patch the generator, not `chunk_text`: retain streams its chunks since #3756, and
@@ -753,7 +752,7 @@ async def test_degenerate_fact_preserves_later_chunk_provenance(memory, request_
ExtractedFact(fact_text=real_fact_by_chunk[chunk_text], fact_type="world", chunk_index=0),
ExtractedFact(fact_text="...", fact_type="world", chunk_index=0),
]
return (
return ExtractionResult(
facts,
[ChunkMetadata(chunk_text=chunk_text, fact_count=len(facts), content_index=0, chunk_index=0)],
TokenUsage(),
+37 -10
View File
@@ -178,7 +178,7 @@ async def test_batch_api_normal_flow(mock_llm_config, test_contents, hindsight_c
mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(return_value=mock_results)
# Call batch API extraction
facts, chunks, usage = await extract_facts_from_contents_batch_api(
extraction = await extract_facts_from_contents_batch_api(
contents=test_contents,
llm_config=mock_llm_config,
config=hindsight_config,
@@ -186,6 +186,9 @@ async def test_batch_api_normal_flow(mock_llm_config, test_contents, hindsight_c
operation_id=None,
schema=None,
)
facts = extraction.facts
chunks = extraction.chunks
usage = extraction.usage
# Verify results
assert len(facts) == 2, "Should extract 2 facts (one per chunk)"
@@ -265,7 +268,7 @@ async def test_batch_api_accepts_top_level_fact_list(mock_llm_config, test_conte
]
)
facts, chunks, usage = await extract_facts_from_contents_batch_api(
extraction = await extract_facts_from_contents_batch_api(
contents=[test_contents[0]],
llm_config=mock_llm_config,
config=hindsight_config,
@@ -273,6 +276,9 @@ async def test_batch_api_accepts_top_level_fact_list(mock_llm_config, test_conte
operation_id=None,
schema=None,
)
facts = extraction.facts
chunks = extraction.chunks
usage = extraction.usage
assert len(facts) == 1
assert "Alice" in facts[0].fact_text
@@ -304,7 +310,7 @@ async def test_batch_api_rejects_top_level_non_fact_list(mock_llm_config, test_c
]
)
facts, chunks, usage = await extract_facts_from_contents_batch_api(
extraction = await extract_facts_from_contents_batch_api(
contents=[test_contents[0]],
llm_config=mock_llm_config,
config=hindsight_config,
@@ -312,6 +318,9 @@ async def test_batch_api_rejects_top_level_non_fact_list(mock_llm_config, test_c
operation_id=None,
schema=None,
)
facts = extraction.facts
chunks = extraction.chunks
usage = extraction.usage
assert facts == []
assert len(chunks) == 1
@@ -372,7 +381,7 @@ async def test_batch_api_records_schema_drifted_facts_as_extraction_errors(
"hindsight_api.engine.retain.fact_extraction._write_batch_extraction_errors",
side_effect=_capture,
):
facts, chunks, _usage = await extract_facts_from_contents_batch_api(
extraction = await extract_facts_from_contents_batch_api(
contents=[test_contents[0]],
llm_config=mock_llm_config,
config=hindsight_config,
@@ -380,6 +389,8 @@ async def test_batch_api_records_schema_drifted_facts_as_extraction_errors(
operation_id=None,
schema=None,
)
facts = extraction.facts
chunks = extraction.chunks
assert facts == []
assert chunks[0].fact_count == 0
@@ -439,7 +450,7 @@ async def test_batch_api_recovers_fenced_and_control_char_json(mock_llm_config,
]
)
facts, chunks, usage = await extract_facts_from_contents_batch_api(
extraction = await extract_facts_from_contents_batch_api(
contents=[test_contents[0]],
llm_config=mock_llm_config,
config=hindsight_config,
@@ -447,6 +458,9 @@ async def test_batch_api_recovers_fenced_and_control_char_json(mock_llm_config,
operation_id=None,
schema=None,
)
facts = extraction.facts
chunks = extraction.chunks
usage = extraction.usage
# The facts are recovered rather than lost.
assert len(facts) == 1
@@ -487,7 +501,7 @@ async def test_batch_api_unparseable_json_still_records_error(mock_llm_config, t
]
)
facts, chunks, usage = await extract_facts_from_contents_batch_api(
extraction = await extract_facts_from_contents_batch_api(
contents=[test_contents[0]],
llm_config=mock_llm_config,
config=hindsight_config,
@@ -495,6 +509,9 @@ async def test_batch_api_unparseable_json_still_records_error(mock_llm_config, t
operation_id=None,
schema=None,
)
facts = extraction.facts
chunks = extraction.chunks
usage = extraction.usage
assert facts == []
assert len(chunks) == 1
@@ -613,7 +630,7 @@ async def test_batch_api_crash_recovery(mock_llm_config, test_contents, hindsigh
mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(return_value=mock_results)
# Call batch API extraction with operation_id (crash recovery scenario)
facts, chunks, usage = await extract_facts_from_contents_batch_api(
extraction = await extract_facts_from_contents_batch_api(
contents=test_contents,
llm_config=mock_llm_config,
config=hindsight_config,
@@ -621,6 +638,9 @@ async def test_batch_api_crash_recovery(mock_llm_config, test_contents, hindsigh
operation_id=operation_id, # Provides crash recovery context
schema=schema,
)
facts = extraction.facts
chunks = extraction.chunks
usage = extraction.usage
# Verify results
assert len(facts) == 2, "Should extract 2 facts after recovery"
@@ -713,7 +733,7 @@ async def test_batch_api_records_non_fatal_extraction_errors(
]
)
facts, chunks, usage = await extract_facts_from_contents_batch_api(
extraction = await extract_facts_from_contents_batch_api(
contents=test_contents,
llm_config=mock_llm_config,
config=hindsight_config,
@@ -721,6 +741,9 @@ async def test_batch_api_records_non_fatal_extraction_errors(
operation_id=operation_id,
schema=schema,
)
facts = extraction.facts
chunks = extraction.chunks
usage = extraction.usage
assert len(facts) == 1
assert len(chunks) == 2
@@ -885,7 +908,7 @@ async def test_batch_api_via_extract_facts_from_contents(
)
# Call main extract_facts_from_contents (should route to batch API)
facts, chunks, usage = await extract_facts_from_contents(
extraction = await extract_facts_from_contents(
contents=test_contents,
llm_config=mock_llm_config,
config=hindsight_config,
@@ -893,6 +916,9 @@ async def test_batch_api_via_extract_facts_from_contents(
operation_id=None,
schema=None,
)
facts = extraction.facts
chunks = extraction.chunks
usage = extraction.usage
# Verify batch API was called
mock_llm_config._provider_impl.submit_batch.assert_called_once()
@@ -952,7 +978,7 @@ async def test_batch_api_sanitizes_model_authored_text(mock_llm_config, hindsigh
]
)
facts, _chunks, _usage = await extract_facts_from_contents_batch_api(
extraction = await extract_facts_from_contents_batch_api(
contents=[RetainContent(content="Alex laughed at the joke.")],
llm_config=mock_llm_config,
config=hindsight_config,
@@ -960,6 +986,7 @@ async def test_batch_api_sanitizes_model_authored_text(mock_llm_config, hindsigh
operation_id=None,
schema=None,
)
facts = extraction.facts
assert len(facts) == 1
assert facts[0].fact_text.encode("utf-8") # raised UnicodeEncodeError before the fix
@@ -160,7 +160,7 @@ async def test_real_openai_batch_api(real_llm_config, test_contents_real, integr
# Call REAL batch API extraction
logger.info("\n📤 Submitting batch to OpenAI...")
facts, chunks, usage = await extract_facts_from_contents_batch_api(
extraction = await extract_facts_from_contents_batch_api(
contents=test_contents_real,
llm_config=real_llm_config,
config=integration_config,
@@ -168,6 +168,9 @@ async def test_real_openai_batch_api(real_llm_config, test_contents_real, integr
operation_id=None, # No crash recovery for this test
schema=schema,
)
facts = extraction.facts
chunks = extraction.chunks
usage = extraction.usage
test_end_time = time.time()
total_duration = test_end_time - test_start_time
@@ -33,11 +33,12 @@ async def test_causal_targets_are_offset_from_extraction_group_start(monkeypatch
monkeypatch.setattr(fact_extraction, "_inject_label_tags", lambda *_args: None)
config = dataclasses.replace(_get_raw_config(), retain_extraction_mode="normal", retain_batch_enabled=False)
facts, _, _ = await fact_extraction.extract_facts_from_contents(
extraction = await fact_extraction.extract_facts_from_contents(
[RetainContent(content="preceding"), RetainContent(content="causal group")],
llm_config=None,
config=config,
)
facts = extraction.facts
assert facts[5].causal_relations[0].target_fact_index == 4
@@ -66,11 +67,12 @@ async def test_each_chunk_uses_its_own_causal_index_base(monkeypatch):
monkeypatch.setattr(fact_extraction, "_inject_label_tags", lambda *_args: None)
config = dataclasses.replace(_get_raw_config(), retain_extraction_mode="normal", retain_batch_enabled=False)
facts, _, _ = await fact_extraction.extract_facts_from_contents(
extraction = await fact_extraction.extract_facts_from_contents(
[RetainContent(content="two chunks")],
llm_config=None,
config=config,
)
facts = extraction.facts
assert facts[1].causal_relations[0].target_fact_index == 0
assert facts[3].causal_relations[0].target_fact_index == 2
@@ -165,11 +167,12 @@ async def test_batch_causal_targets_use_each_chunk_start(monkeypatch):
return batch_impl
llm_config = SimpleNamespace(batch_provider_impl=_batch_provider_impl, provider="test")
facts, _, _ = await fact_extraction.extract_facts_from_contents_batch_api(
extraction = await fact_extraction.extract_facts_from_contents_batch_api(
[RetainContent(content="preceding"), RetainContent(content="first|second")],
llm_config=llm_config,
config=config,
)
facts = extraction.facts
assert facts[5].causal_relations[0].target_fact_index == 4
assert facts[7].causal_relations[0].target_fact_index == 6
@@ -986,17 +986,15 @@ async def test_delta_retain_recall_with_chunks(memory, request_context):
def test_merge_processed_content_tokens_helper():
"""Unit check on the None-propagating aggregator used by the engine."""
from hindsight_api.engine.retain.orchestrator import (
_merge_processed_content_tokens,
)
from hindsight_api.engine.retain.types import merge_processed_content_tokens
assert _merge_processed_content_tokens(0, 0) == 0
assert _merge_processed_content_tokens(5, 7) == 12
assert merge_processed_content_tokens(0, 0) == 0
assert merge_processed_content_tokens(5, 7) == 12
# None "wins" in either slot — once any sub-result bypassed dedup, the
# aggregate is None so callers bill full content.
assert _merge_processed_content_tokens(None, 10) is None
assert _merge_processed_content_tokens(10, None) is None
assert _merge_processed_content_tokens(None, None) is None
assert merge_processed_content_tokens(None, 10) is None
assert merge_processed_content_tokens(10, None) is None
assert merge_processed_content_tokens(None, None) is None
@pytest.mark.asyncio
@@ -45,11 +45,12 @@ async def test_extract_facts_preserves_experience_type():
"hindsight_api.engine.retain.fact_extraction.extract_facts_from_text",
new=AsyncMock(return_value=([extracted_fact], [(contents[0].content, 1)], TokenUsage())),
):
facts, _chunks, _usage = await extract_facts_from_contents(
extraction = await extract_facts_from_contents(
contents=contents,
llm_config=None,
config=_get_raw_config(),
)
facts = extraction.facts
assert len(facts) == 1
assert facts[0].fact_type == "experience", (
@@ -103,7 +103,7 @@ async def test_real_fireworks_batch_end_to_end(fireworks_env):
]
logger.info("Submitting a real Fireworks batch (this can take several minutes)...")
facts, chunks, usage = await extract_facts_from_contents_batch_api(
extraction = await extract_facts_from_contents_batch_api(
contents=contents,
llm_config=llm_config,
config=config,
@@ -111,6 +111,9 @@ async def test_real_fireworks_batch_end_to_end(fireworks_env):
operation_id=None,
schema=None,
)
facts = extraction.facts
chunks = extraction.chunks
usage = extraction.usage
# The end-to-end proof: if the real output shape doesn't match the normalizer,
# the consumer extracts nothing and this is empty.
@@ -105,7 +105,7 @@ async def test_real_gemini_batch_end_to_end(gemini_env):
]
logger.info("Submitting a real Gemini batch (this can take several minutes)...")
facts, chunks, usage = await extract_facts_from_contents_batch_api(
extraction = await extract_facts_from_contents_batch_api(
contents=contents,
llm_config=llm_config,
config=config,
@@ -113,6 +113,9 @@ async def test_real_gemini_batch_end_to_end(gemini_env):
operation_id=None,
schema=None,
)
facts = extraction.facts
chunks = extraction.chunks
usage = extraction.usage
# The end-to-end proof: if the real output shape doesn't match the normalizer,
# the consumer extracts nothing and this is empty.
@@ -0,0 +1,101 @@
"""`_internal_error` is the single place a handler's unhandled exception is mapped.
79 route handlers used to spell this out inline: 72 byte-identical copies of
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(f"Error in <route>: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
plus 7 near-copies that logged the traceback without the message. Re-typing the
policy per route is how it drifted; these tests pin the behaviour the copies
shared so the shared version cannot quietly change it.
"""
import logging
import pytest
from fastapi import HTTPException
from hindsight_api.api.http import _internal_error
def _raise_through(exc: Exception, where: str):
"""Call the helper from inside a real `except` block, as every call site does."""
try:
raise exc
except Exception as e:
return _internal_error(e, where)
def test_maps_to_500_with_the_message_as_detail():
"""The client gets the message, never the traceback."""
result = _raise_through(RuntimeError("boom"), "GET /v1/x")
assert isinstance(result, HTTPException)
assert result.status_code == 500
assert result.detail == "boom"
assert "Traceback" not in str(result.detail)
def test_logs_the_route_the_message_and_the_traceback(caplog):
with caplog.at_level(logging.ERROR, logger="hindsight_api.api.http"):
_raise_through(ValueError("bad input"), "GET /v1/default/banks/b1/graph")
assert len(caplog.records) == 1
msg = caplog.records[0].getMessage()
# Exactly the format the 72 inline copies produced.
assert msg.startswith("Error in GET /v1/default/banks/b1/graph: bad input\n\nTraceback:\n")
assert "ValueError: bad input" in msg, "the traceback itself must be in the log"
def test_returns_rather_than_raises():
"""Call sites read `raise _internal_error(...)`; the helper must not raise itself."""
result = _raise_through(RuntimeError("x"), "GET /v1/x")
assert isinstance(result, HTTPException)
@pytest.mark.parametrize(
"exc",
[
RuntimeError('relation "memory_units" does not exist'),
ValueError(""),
KeyError("missing"),
Exception("multi\nline\nmessage"),
],
)
def test_detail_is_always_the_stringified_exception(exc):
assert _raise_through(exc, "GET /v1/x").detail == str(exc)
def test_no_handler_still_builds_the_500_inline():
"""Structural guard over the whole module, not any one route.
A reintroduced inline copy is invisible per-route it behaves identically
until someone changes the policy in the helper and one route ignores it.
"""
import ast
import pathlib
source = pathlib.Path(__file__).resolve().parent.parent / "hindsight_api" / "api" / "http.py"
tree = ast.parse(source.read_text())
# Two handlers legitimately keep their own catch-all because they add
# diagnostics the shared helper cannot know about. Named, not silently
# skipped, so adding a third is a deliberate act:
# recall — logs handler_duration in its own [RECALL ERROR] format
# retain — maps MemoryDefenseAllBlockedError to 422 and logs an input summary
EXEMPT_MARKERS = ("[RECALL ERROR]", "MemoryDefenseAllBlockedError")
offenders = []
for node in ast.walk(tree):
if not isinstance(node, ast.Try):
continue
for handler in node.handlers:
if handler.type is None or ast.unparse(handler.type) != "Exception":
continue
body = "\n".join(ast.unparse(s) for s in handler.body)
if "HTTPException(status_code=500" not in body or "_internal_error" in body:
continue
if any(marker in body for marker in EXEMPT_MARKERS):
continue
offenders.append(handler.lineno)
assert not offenders, f"http.py builds a 500 inline instead of via _internal_error at lines {offenders}"
@@ -21,7 +21,7 @@ import pytest
from hindsight_api.config import clear_config_cache
from tests.sub_batch_helpers import collect_sub_batches
from hindsight_api.engine.response_models import TokenUsage
from hindsight_api.engine.retain.types import ChunkMetadata, ExtractedFact, RetainContent
from hindsight_api.engine.retain.types import ChunkMetadata, ExtractedFact, ExtractionResult, RetainContent
def _ts() -> float:
@@ -199,7 +199,7 @@ async def test_append_after_zero_fact_header_slice_skips_unchanged_history(
contents: list[RetainContent],
*args: Any,
**kwargs: Any,
) -> tuple[list[ExtractedFact], list[ChunkMetadata], TokenUsage]:
) -> ExtractionResult:
extracted_contents.extend(item.content for item in contents)
facts: list[ExtractedFact] = []
chunks: list[ChunkMetadata] = []
@@ -224,7 +224,7 @@ async def test_append_after_zero_fact_header_slice_skips_unchanged_history(
tags=item.tags,
)
)
return facts, chunks, TokenUsage()
return ExtractionResult(facts, chunks, TokenUsage())
monkeypatch.setattr(
fact_extraction,
@@ -0,0 +1,80 @@
"""The link-expansion arms must all project the same ``memory_units`` columns.
Every arm of link expansion (entity / semantic / causal, and their observation
variants) is combined with ``UNION ALL`` and read back positionally, so the arms
have to agree on both the set *and* the order of the leading columns. Before
``memory_unit_columns`` the list was spelled out ~20 times across the two
backends; these tests are what makes that single source of truth enforceable,
and they fail loudly if a future arm hand-rolls the projection again.
"""
import re
import pytest
from hindsight_api.engine.db.ops import MEMORY_UNIT_COLUMNS, UpdatedWindow, memory_unit_columns
from hindsight_api.engine.db.ops_oracle import OracleOps
from hindsight_api.engine.db.ops_postgresql import PostgreSQLOps
WINDOWS = [
UpdatedWindow(after=None, before=None, first_param_index=4),
UpdatedWindow(after=None, before=None, first_param_index=1),
]
def test_memory_unit_columns_qualifies_every_column():
rendered = memory_unit_columns("mu")
assert [c.strip() for c in rendered.replace("\n", " ").split(",")] == [
f"mu.{column}" for column in MEMORY_UNIT_COLUMNS
]
def test_memory_unit_columns_unqualified_by_default():
rendered = memory_unit_columns()
assert [c.strip() for c in rendered.replace("\n", " ").split(",")] == list(MEMORY_UNIT_COLUMNS)
def test_memory_unit_columns_indents_continuation_lines_only():
rendered = memory_unit_columns("mu", indent=4)
lines = rendered.split("\n")
assert len(lines) > 1, "expected the projection to wrap"
assert not lines[0].startswith(" "), "the first line sits after SELECT and must not be indented"
assert all(line.startswith(" ") for line in lines[1:])
def _arm_projections(sql: str) -> dict[str, list[str]]:
"""Leading column list of each ``<name> AS ( SELECT ... )`` arm in a CTE body."""
arms: dict[str, list[str]] = {}
for match in re.finditer(r"(\w+) AS \(\s*SELECT(?: DISTINCT ON \([^)]*\))?\s*([\s\S]*?)\n\s*FROM ", sql):
name, projection = match.group(1), match.group(2)
columns = [c.strip().split(".")[-1] for c in projection.split(",")]
arms[name] = columns[: len(MEMORY_UNIT_COLUMNS)]
return arms
@pytest.mark.parametrize("ops", [PostgreSQLOps(), OracleOps()], ids=["pg", "oracle"])
@pytest.mark.parametrize("window", WINDOWS, ids=["window_at_4", "window_at_1"])
def test_expansion_arms_share_the_memory_unit_projection(ops, window):
sql = ops.build_entity_expansion_cte("mu_t", "ue_t", 5, window) + ",\n"
sql += ops.build_semantic_causal_cte("ml_t", "mu_t", window)
arms = _arm_projections(sql)
expanded = {name: cols for name, cols in arms.items() if name.endswith("_expanded")}
assert expanded, f"no expansion arms parsed out of:\n{sql}"
for name, columns in expanded.items():
assert columns == list(MEMORY_UNIT_COLUMNS), f"arm {name!r} projects {columns}"
def test_postgres_semantic_arm_groups_by_every_projected_column():
"""A projected column missing from GROUP BY is a runtime error, not a test failure.
PostgreSQL only: Oracle cannot GROUP BY its CLOB columns, so its semantic arm
deliberately groups by ``id`` alone in a scores subquery and joins back for the
full projection.
"""
sql = PostgreSQLOps().build_semantic_causal_cte("ml_t", "mu_t", WINDOWS[0])
group_by = re.search(r"GROUP BY ([\s\S]*?)\n\s*ORDER BY", sql)
assert group_by is not None, f"expected a grouped semantic arm in:\n{sql}"
grouped = {c.strip().split(".")[-1] for c in group_by.group(1).split(",")}
assert grouped == set(MEMORY_UNIT_COLUMNS)
@@ -47,20 +47,22 @@ async def test_activation_preserves_additive_score_across_fact_types(monkeypatch
monkeypatch.setattr(retriever, "_expand_combined", fake_expand_combined)
pool = SimpleNamespace(ops=object())
world_results, _ = await retriever.retrieve(
retrieved = await retriever.retrieve(
pool,
query_embedding_str="unused",
bank_id="bank",
fact_type="world",
budget=2,
)
experience_results, _ = await retriever.retrieve(
world_results = retrieved.results
retrieved = await retriever.retrieve(
pool,
query_embedding_str="unused",
bank_id="bank",
fact_type="experience",
budget=2,
)
experience_results = retrieved.results
combined = world_results + experience_results
combined.sort(key=lambda result: result.activation or 0.0, reverse=True)
@@ -89,7 +91,7 @@ async def test_preselected_semantic_seeds_skip_seed_query(monkeypatch):
monkeypatch.setattr(link_expansion_retrieval, "_find_semantic_seeds", fail_find_semantic_seeds)
monkeypatch.setattr(retriever, "_expand_combined", fake_expand_combined)
results, timings = await retriever.retrieve(
retrieved = await retriever.retrieve(
SimpleNamespace(ops=object()),
query_embedding_str="unused",
bank_id="bank",
@@ -100,6 +102,8 @@ async def test_preselected_semantic_seeds_skip_seed_query(monkeypatch):
RetrievalResult(id="seed-b", text="seed", fact_type="world"),
],
)
results = retrieved.results
timings = retrieved.timings
assert [result.id for result in results] == ["result"]
assert timings is not None
@@ -120,7 +124,7 @@ async def test_empty_preselected_semantic_seeds_do_not_fall_back(monkeypatch):
monkeypatch.setattr(link_expansion_retrieval, "acquire_with_retry", fake_acquire_with_retry)
monkeypatch.setattr(link_expansion_retrieval, "_find_semantic_seeds", fail_find_semantic_seeds)
results, timings = await retriever.retrieve(
retrieved = await retriever.retrieve(
SimpleNamespace(ops=object()),
query_embedding_str="unused",
bank_id="bank",
@@ -128,6 +132,8 @@ async def test_empty_preselected_semantic_seeds_do_not_fall_back(monkeypatch):
budget=2,
preselected_semantic_seeds=[],
)
results = retrieved.results
timings = retrieved.timings
assert results == []
assert timings is not None
@@ -0,0 +1,75 @@
"""MCP tool error payloads must be valid JSON, whatever the message contains.
The bank-id-parameter variants of the MCP tools declare ``-> str`` and return
JSON text, so their error branch has to return JSON too. Every one of them used
to build it by interpolation::
return f'{{"error": "{e}"}}'
which emits invalid JSON as soon as the exception message contains a double
quote, a backslash or a newline. That is not a corner case: PostgreSQL quotes
identifiers with double quotes, so a plain ``relation "memory_units" does not
exist`` already produced something the caller could not parse. The bug lived
only in the bank-id half of each duplicated tool -- the single-bank half
returned a dict and was always correct -- which is exactly why no test caught
it.
"""
import json
import pytest
from hindsight_api.mcp_tools import _error_json
# Messages a real backend actually produces. The PG ones are the reason this
# matters; the others cover the remaining JSON metacharacters.
HOSTILE_MESSAGES = [
'relation "memory_units" does not exist',
'column "tags" is of type text[] but expression is of type text',
'duplicate key value violates unique constraint "banks_pkey"',
"back\\slash",
"line one\nline two",
'tab\there and "quotes"',
"unicode: — é 中文",
"",
]
@pytest.mark.parametrize("message", HOSTILE_MESSAGES)
def test_error_payload_is_parseable_json(message):
parsed = json.loads(_error_json(Exception(message)))
assert parsed == {"error": message}
@pytest.mark.parametrize("message", HOSTILE_MESSAGES)
def test_error_payload_keeps_its_empty_collection(message):
"""Tools that return a collection include an empty one so callers keep their shape."""
parsed = json.loads(_error_json(Exception(message), results=[]))
assert parsed == {"error": message, "results": []}
def test_error_payload_accepts_a_plain_string():
"""The 'no bank configured' branch passes a str, not an exception."""
assert json.loads(_error_json("No bank_id configured")) == {"error": "No bank_id configured"}
def test_no_tool_builds_error_json_by_interpolation():
"""Structural guard: the f-string form must not come back.
A single reintroduced ``f'{{"error": "{e}"}}'`` is invisible until a message
with a quote reaches that one branch, so this asserts over the whole file
rather than over any one tool.
"""
import pathlib
import re
source = pathlib.Path(__file__).resolve().parent.parent / "hindsight_api" / "mcp_tools.py"
text = source.read_text()
# Skip the docstring that quotes the old form as the thing not to do.
body = text.replace('``f\'{{"error": "{e}"}}\'``', "")
offenders = [
(i + 1, line.strip())
for i, line in enumerate(body.split("\n"))
if re.search(r"""return\s+f['"]\{\{["']error""", line)
]
assert not offenders, f"mcp_tools.py builds error JSON by interpolation at {offenders}"
@@ -0,0 +1,110 @@
"""Each MCP tool is registered twice; the two copies must not diverge.
Every tool is declared once with an explicit ``bank_id`` parameter (returning
JSON text) and once resolving the bank from the session (returning a dict).
FastMCP builds each tool's schema from the literal signature and docstring, so
those two declarations genuinely have to exist twice but the *logic* does not,
and when it did, the copies drifted: the bank-id half built its error JSON by
string interpolation and emitted invalid JSON for any message containing a
double quote, while the session half returned a dict and was always correct.
These tests pin what must stay identical between the copies, and assert the
shared wrapper is actually used, so a future tool cannot quietly grow a second
implementation of the same behaviour.
"""
import ast
import pathlib
import pytest
MCP_TOOLS = pathlib.Path(__file__).resolve().parent.parent / "hindsight_api" / "mcp_tools.py"
#: Tools that legitimately keep bespoke bodies, and why. Named rather than
#: pattern-matched so adding a fifth is a deliberate act, not an accident.
BESPOKE = {
# The two copies use different pydantic serializers -- model_dump_json() for
# the JSON variant, model_dump() for the dict one -- which do not render
# datetimes and enums identically. Routing both through the shared wrapper
# would change the payload clients receive.
"recall",
"reflect",
# Retain resolves its bank per content item and reports partial success, so
# it has no single "resolve one bank, make one call" shape to share.
"retain",
"sync_retain",
}
def _tree():
return ast.parse(MCP_TOOLS.read_text())
def _tool_pairs():
pairs = {}
for node in ast.walk(_tree()):
if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
continue
if not node.name.startswith("_register_"):
continue
by_name: dict[str, list] = {}
for inner in ast.walk(node):
if isinstance(inner, (ast.FunctionDef, ast.AsyncFunctionDef)) and inner is not node:
by_name.setdefault(inner.name, []).append(inner)
for name, fns in by_name.items():
if len(fns) == 2:
pairs[name] = sorted(fns, key=lambda f: f.lineno)
return pairs
def _params(fn):
return [
(a.arg, ast.unparse(a.annotation) if a.annotation else None)
for a in list(fn.args.args) + list(fn.args.kwonlyargs)
]
def test_the_tool_family_is_non_trivial():
"""Guard the guard: an empty family would make every assertion below vacuous."""
assert len(_tool_pairs()) >= 30
@pytest.mark.parametrize("name", sorted(_tool_pairs()))
def test_both_copies_declare_the_same_parameters(name):
"""Beyond bank_id, the two copies must offer the caller the same surface.
A parameter added to one copy only is invisible: each registration is
exercised by whichever deployment mode uses it, so the mode nobody tested
silently lacks the capability.
"""
multi, single = _tool_pairs()[name]
multi_params = [p for p in _params(multi) if p[0] != "bank_id"]
assert multi_params == _params(single), (
f"{name}: bank-id copy declares {multi_params}, session copy declares {_params(single)}"
)
@pytest.mark.parametrize("name", sorted(_tool_pairs()))
def test_only_the_bank_id_copy_takes_bank_id(name):
multi, single = _tool_pairs()[name]
assert "bank_id" in dict(_params(multi)), f"{name}: bank-id copy is missing bank_id"
assert "bank_id" not in dict(_params(single)), f"{name}: session copy should not take bank_id"
@pytest.mark.parametrize("name", sorted(_tool_pairs()))
def test_the_copies_share_one_implementation(name):
"""Both copies must delegate to _run_tool rather than re-implement the flow."""
if name in BESPOKE:
pytest.skip(f"{name} is a documented exception (see BESPOKE)")
for fn in _tool_pairs()[name]:
body = "\n".join(ast.unparse(s) for s in fn.body)
assert "_run_tool" in body, (
f"{name} (line {fn.lineno}) does not use the shared _run_tool wrapper; "
f"if that is deliberate, add it to BESPOKE with the reason"
)
def test_bespoke_list_has_no_stale_entries():
"""An exemption that no longer names a real tool hides a rule that stopped applying."""
pairs = set(_tool_pairs())
assert BESPOKE <= pairs, f"BESPOKE names tools that no longer exist: {sorted(BESPOKE - pairs)}"
@@ -0,0 +1,174 @@
"""Test doubles must return the named result type the real function returns.
The refactor that replaced bare tuple returns with named types was verified by
sweeping for *callers* that still unpacked a tuple. That missed the opposite
direction, and CI caught it: five upstream tests stub one of these functions and
returned the old tuple, so production code doing ``result.results`` got an
``AttributeError`` or, worse, hung, because the stubbed pipeline never reached
the failure the test was waiting for.
A stub is invisible to a caller-side sweep and to the type checker, so this
asserts over the whole test suite instead: any double that stands in for one of
these functions must produce the same shape the real one does.
Two ways a double is recognised, because both occur:
* ``monkeypatch.setattr(module, "<name>", replacement)`` the target names the
function, and the replacement is looked up by name in the same file.
* a method named after the function (``FakeGraphRetriever.retrieve``) installed
by handing the whole object over, so the setattr target says nothing useful.
"""
import ast
import pathlib
import pytest
TESTS_DIR = pathlib.Path(__file__).resolve().parent
#: Function name -> the named type it returns. Keep in step with the dataclasses;
#: `test_retain_result_shape.py` asserts the functions themselves still return them.
NAMED_RESULTS = {
"retain_batch": "RetainBatchResult",
"_streaming_retain_batch": "RetainBatchResult",
"_try_delta_retain": "RetainBatchResult",
"_delta_metadata_only": "RetainBatchResult",
"_retain_batch_async_internal": "RetainBatchResult",
"_retain_batch_with_append_retry": "RetainBatchResult",
"extract_facts_from_contents": "ExtractionResult",
"extract_facts_from_contents_batch_api": "ExtractionResult",
"_extract_facts_chunks": "ExtractionResult",
"_extract_and_embed": "_EmbeddedExtraction",
"_prepare_facts_for_entity_processing": "PreparedFactEntities",
"build_tags_where_clause": "TagClause",
"build_tag_groups_where_clause": "TagClause",
"_build_group_clause": "TagClause",
"_parse_tags_match": "TagMatchSemantics",
"retrieve": "GraphRetrieval",
"_fit_structured_delta_prompt_parts": "FittedDeltaPrompt",
"_validate_operations_list": "ValidatedOperations",
}
def _test_modules():
return sorted(p for p in TESTS_DIR.glob("test_*.py"))
def _returns_bare_tuple(fn: ast.AST) -> list[int]:
"""Line numbers where ``fn`` itself returns a tuple literal (nested defs excluded)."""
nested = [n for n in ast.walk(fn) if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef)) and n is not fn]
return [
r.lineno
for r in ast.walk(fn)
if isinstance(r, ast.Return)
and isinstance(r.value, ast.Tuple)
and not any(i.lineno <= r.lineno <= i.end_lineno for i in nested)
]
def _doubles_in(tree: ast.AST) -> list[tuple[str, ast.AST]]:
"""(function-it-stands-in-for, the double's def) pairs found in one module."""
by_name = {n.name: n for n in ast.walk(tree) if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))}
found: list[tuple[str, ast.AST]] = []
# 1. monkeypatch.setattr(module, "<name>", replacement)
for node in ast.walk(tree):
if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)):
continue
if node.func.attr != "setattr" or len(node.args) < 3:
continue
target, replacement = node.args[1], node.args[2]
if not (isinstance(target, ast.Constant) and isinstance(target.value, str)):
continue
if target.value not in NAMED_RESULTS:
continue
if isinstance(replacement, ast.Name) and replacement.id in by_name:
found.append((target.value, by_name[replacement.id]))
# 2. a method named after the function, on any class in the module
for cls in [n for n in ast.walk(tree) if isinstance(n, ast.ClassDef)]:
for item in cls.body:
if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)) and item.name in NAMED_RESULTS:
found.append((item.name, item))
return found
def test_the_scan_finds_doubles_at_all():
"""Guard the guard: if the recognisers stop matching, every assertion below is vacuous."""
total = 0
for path in _test_modules():
try:
total += len(_doubles_in(ast.parse(path.read_text())))
except SyntaxError:
continue
assert total >= 5, f"only found {total} doubles — the recognisers have probably stopped matching"
@pytest.mark.parametrize("path", _test_modules(), ids=lambda p: p.stem)
def test_doubles_return_the_named_result(path):
try:
tree = ast.parse(path.read_text())
except SyntaxError:
pytest.skip("module does not parse")
offenders = []
for stands_for, fn in _doubles_in(tree):
for lineno in _returns_bare_tuple(fn):
offenders.append(
f"{path.name}:{lineno} {fn.name}() stands in for "
f"{stands_for}() and must return {NAMED_RESULTS[stands_for]}, not a tuple"
)
assert not offenders, "\n".join(offenders)
#: Both trees are scanned: a caller that unpacks is as broken as a stub that supplies.
_SCANNED_ROOTS = ("hindsight_api", "tests")
def _callee_name(call: ast.Call) -> str | None:
func = call.func
if isinstance(func, ast.Name):
return func.id
if isinstance(func, ast.Attribute):
return func.attr
return None
def _source_files():
root = TESTS_DIR.parent
for sub in _SCANNED_ROOTS:
yield from sorted((root / sub).rglob("*.py"))
def test_nothing_unpacks_a_named_result_as_a_tuple():
"""The caller-side half: a named result must not be destructured positionally.
This is what a diff-based review does catch but only for the call sites that
exist when the refactor lands. A rebase brings new ones, which is exactly how
``test_tag_resolution.py`` reached CI with ``clause, params, _ = ...``. Asserting
it over both trees means a new caller cannot arrive un-swept.
"""
offenders = []
for path in _source_files():
try:
tree = ast.parse(path.read_text())
except (SyntaxError, UnicodeDecodeError):
continue
rel = path.relative_to(TESTS_DIR.parent)
for node in ast.walk(tree):
calls: list[ast.Call] = []
if isinstance(node, ast.Assign) and node.targets and isinstance(node.targets[0], (ast.Tuple, ast.List)):
calls = [c for c in ast.walk(node.value) if isinstance(c, ast.Call)]
elif isinstance(node, ast.For) and isinstance(node.target, (ast.Tuple, ast.List)):
calls = [c for c in ast.walk(node.iter) if isinstance(c, ast.Call)]
elif isinstance(node, ast.Subscript):
inner = node.value.value if isinstance(node.value, ast.Await) else node.value
calls = [inner] if isinstance(inner, ast.Call) else []
for call in calls:
name = _callee_name(call)
if name in NAMED_RESULTS:
offenders.append(
f"{rel}:{node.lineno} treats {name}() as a tuple; "
f"it returns {NAMED_RESULTS[name]} — read its fields by name"
)
assert not offenders, "\n".join(sorted(set(offenders)))
@@ -0,0 +1,79 @@
"""A provider must accept a capability's parameters only if it implements the capability.
Prompt caching is opt-in: ``supports_prompt_caching()`` defaults False and
``get_or_create_cached_prefix()`` defaults to returning None, so callers only
ever pass ``cached_prefix=`` to a provider that overrode them. Two providers had
copied the parameter out of the interface signature without any of the
behaviour -- they declared ``cached_prefix``/``cached_prefix_message_count``,
never referenced them, and could never be handed one. Dead surface that reads
like a supported feature.
This asserts over the whole provider family rather than any one provider,
because the failure mode is a member that *omits* something (or, here, keeps
something it shouldn't) and the member that forgot is by construction the one
without a test.
"""
import ast
import pathlib
import pytest
PROVIDERS_DIR = pathlib.Path(__file__).resolve().parent.parent / "hindsight_api" / "engine" / "providers"
CACHE_PARAMS = {"cached_prefix", "cached_prefix_message_count"}
#: Overriding any of these is what "implements prompt caching" means.
CACHE_METHODS = {
"get_or_create_cached_prefix",
"create_incremental_cache",
"supports_prompt_caching",
"supports_incremental_prompt_cache",
}
def _provider_modules():
return sorted(p for p in PROVIDERS_DIR.glob("*_llm.py"))
def _classes(path):
tree = ast.parse(path.read_text())
return [n for n in ast.walk(tree) if isinstance(n, ast.ClassDef)]
def test_the_provider_family_is_non_trivial():
"""Guard the guard: an empty family would make every assertion below vacuous."""
assert len(_provider_modules()) >= 10
@pytest.mark.parametrize("path", _provider_modules(), ids=lambda p: p.stem)
def test_cache_parameters_only_where_caching_is_implemented(path):
for cls in _classes(path):
methods = {b.name: b for b in cls.body if isinstance(b, (ast.FunctionDef, ast.AsyncFunctionDef))}
implements = bool(CACHE_METHODS & set(methods))
for name in ("call", "call_with_tools"):
fn = methods.get(name)
if fn is None:
continue
params = {a.arg for a in list(fn.args.args) + list(fn.args.kwonlyargs)}
declared = CACHE_PARAMS & params
if declared and not implements:
pytest.fail(
f"{path.name}::{cls.name}.{name}() declares {sorted(declared)} but the class "
f"overrides none of {sorted(CACHE_METHODS)} — callers gate on "
f"get_or_create_cached_prefix() returning non-None, so it can never be passed one"
)
@pytest.mark.parametrize("path", _provider_modules(), ids=lambda p: p.stem)
def test_declared_cache_parameters_are_actually_used(path):
"""A parameter present in the signature must be read somewhere in the body."""
for cls in _classes(path):
for b in cls.body:
if not isinstance(b, (ast.FunctionDef, ast.AsyncFunctionDef)):
continue
params = {a.arg for a in list(b.args.args) + list(b.args.kwonlyargs)}
for param in CACHE_PARAMS & params:
loads = [
n for n in ast.walk(b) if isinstance(n, ast.Name) and n.id == param and isinstance(n.ctx, ast.Load)
]
assert loads, f"{path.name}::{cls.name}.{b.name}() declares {param} but never reads it"
@@ -18,6 +18,7 @@ from types import SimpleNamespace
import pytest
from hindsight_api.engine import memory_engine as memory_engine_module
from hindsight_api.engine.search.types import GraphRetrieval
from hindsight_api.engine.search import retrieval as retrieval_module
_QUERY = "[0.1,0.2,0.3]"
@@ -52,7 +53,7 @@ def stub_retrieval(monkeypatch):
class FakeGraphRetriever:
async def retrieve(self, **kwargs):
calls["graph"] += 1
return [], None
return GraphRetrieval([], None)
fake_config = SimpleNamespace(graph_seed_min_similarity=0.3, temporal_semantic_min_similarity=0.24)
+10 -3
View File
@@ -2642,13 +2642,16 @@ def test_chunks_extraction_mode():
RetainContent(content="Bob fixed the critical bug in the payment service."),
]
facts, chunks, usage = asyncio.run(
extraction = asyncio.run(
extract_facts_from_contents(
contents=contents,
llm_config=None, # Must not be called
config=_get_raw_config(),
)
)
facts = extraction.facts
chunks = extraction.chunks
usage = extraction.usage
# One fact per chunk (both contents fit in one chunk each)
assert len(facts) == len(chunks) == 2
@@ -2709,11 +2712,13 @@ async def test_verbatim_extraction_mode():
content=text, event_date=datetime(2024, 3, 10, tzinfo=timezone.utc), context="onboarding notes"
)
]
facts, chunks, _ = await extract_facts_from_contents(
extraction = await extract_facts_from_contents(
contents=contents,
llm_config=llm_config,
config=_get_raw_config(),
)
facts = extraction.facts
chunks = extraction.chunks
logger.info(f"Verbatim mode extracted {len(facts)} facts from {len(chunks)} chunks")
for i, f in enumerate(facts):
@@ -2960,13 +2965,15 @@ def test_strategy_overrides_extraction_mode_for_chunks():
RetainContent(content="Bob reviewed the pull request."),
]
facts, chunks, usage = asyncio.run(
extraction = asyncio.run(
extract_facts_from_contents(
contents=contents,
llm_config=None, # chunks must not call the LLM
config=strategy_config,
)
)
facts = extraction.facts
usage = extraction.usage
assert len(facts) == 2
assert facts[0].fact_text == contents[0].content
@@ -59,14 +59,14 @@ async def test_consumer_failure_cancels_in_flight_extractions(monkeypatch):
nonlocal calls, cancelled
calls += 1
if calls == 1:
return [], [], [], TokenUsage()
return orchestrator._EmbeddedExtraction([], [], [], TokenUsage())
hanging_started.set()
try:
await asyncio.Event().wait()
except asyncio.CancelledError:
cancelled += 1
raise
return [], [], [], TokenUsage()
return orchestrator._EmbeddedExtraction([], [], [], TokenUsage())
monkeypatch.setattr(orchestrator, "_extract_and_embed", fake_extract_and_embed)
@@ -0,0 +1,163 @@
"""Every retain pipeline path must return the named ``RetainBatchResult``.
A structural guard, not a behavioural one, because the defect it prevents is
invisible to behavioural tests. ``_streaming_retain_batch`` was annotated
``tuple[list[list[str]], TokenUsage]`` while returning three values, and
``retain_batch`` passed that straight back as its own 3-tuple. Every test still
passed the extra element flowed through positionally and ``ty`` cannot catch
it because ``invalid-return-type`` is disabled in ``pyproject.toml``. The
mismatch would only have surfaced as a ``ValueError`` the day someone unpacked
the annotated two names.
Mirrors ``test_migration_shape.py``: enumerate the family from the source and
assert every member satisfies the contract, so the next path added to the
pipeline cannot quietly reintroduce a bare tuple.
"""
import ast
import pathlib
import pytest
from hindsight_api.engine.retain.types import RetainBatchResult, merge_processed_content_tokens
API_ROOT = pathlib.Path(__file__).resolve().parent.parent / "hindsight_api"
#: The pipeline entry points, and the engine wrappers that forward their result.
RETAIN_RESULT_FUNCTIONS = {
"engine/retain/orchestrator.py": [
"retain_batch",
"_streaming_retain_batch",
"_try_delta_retain",
"_delta_metadata_only",
],
"engine/memory_engine.py": [
"_retain_batch_async_internal",
"_retain_batch_with_append_retry",
],
}
#: Other families where a bare tuple was the contract, and what replaced it.
#: Each entry is (module, function names, allowed return annotations).
NAMED_RESULT_FAMILIES = [
(
"engine/retain/fact_extraction.py",
["extract_facts_from_contents", "extract_facts_from_contents_batch_api", "_extract_facts_chunks"],
{"ExtractionResult"},
),
(
"engine/search/tags.py",
["build_tags_where_clause", "build_tag_groups_where_clause", "_build_group_clause"],
{"TagClause"},
),
(
"engine/search/graph_retrieval.py",
["retrieve"],
{"GraphRetrieval"},
),
(
"engine/search/link_expansion_retrieval.py",
["retrieve"],
{"GraphRetrieval"},
),
(
"engine/reflect/prompts.py",
["_fit_structured_delta_prompt_parts"],
{"FittedDeltaPrompt"},
),
(
"engine/reflect/delta_ops.py",
["_validate_operations_list"],
{"ValidatedOperations"},
),
(
"engine/retain/entity_processing.py",
["_prepare_facts_for_entity_processing"],
{"PreparedFactEntities"},
),
]
def _function_nodes(rel_path: str, names: list[str]) -> dict[str, ast.AST]:
tree = ast.parse((API_ROOT / rel_path).read_text())
found = {
node.name: node
for node in ast.walk(tree)
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name in names
}
missing = set(names) - set(found)
assert not missing, f"{rel_path}: expected functions vanished or were renamed: {sorted(missing)}"
return found
def _own_returns(func: ast.AST) -> list[ast.Return]:
"""Return statements belonging to ``func`` itself, excluding nested defs."""
nested = [n for n in ast.walk(func) if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef)) and n is not func]
return [
r
for r in ast.walk(func)
if isinstance(r, ast.Return) and not any(inner.lineno <= r.lineno <= inner.end_lineno for inner in nested)
]
@pytest.mark.parametrize(
("rel_path", "name"),
[(p, n) for p, names in RETAIN_RESULT_FUNCTIONS.items() for n in names],
)
def test_retain_path_is_annotated_with_the_named_result(rel_path, name):
func = _function_nodes(rel_path, RETAIN_RESULT_FUNCTIONS[rel_path])[name]
assert func.returns is not None, f"{name}() has no return annotation"
annotation = ast.unparse(func.returns).replace('"', "").replace("'", "")
assert annotation in ("RetainBatchResult", "RetainBatchResult | None"), (
f"{name}() returns {annotation!r}; retain paths must return the named "
f"RetainBatchResult so their arity cannot drift"
)
@pytest.mark.parametrize(
("rel_path", "name"),
[(p, n) for p, names in RETAIN_RESULT_FUNCTIONS.items() for n in names],
)
def test_retain_path_never_returns_a_bare_tuple(rel_path, name):
func = _function_nodes(rel_path, RETAIN_RESULT_FUNCTIONS[rel_path])[name]
offenders = [(r.lineno, ast.unparse(r.value)[:60]) for r in _own_returns(func) if isinstance(r.value, ast.Tuple)]
assert not offenders, f"{name}() returns bare tuples at {offenders}"
def test_retain_batch_result_fields_are_ordered_and_named():
"""Positional construction is used throughout, so field order is part of the contract."""
result = RetainBatchResult([["a"], []], usage=None, processed_content_tokens=7)
assert result.memory_ids == [["a"], []]
assert result.processed_content_tokens == 7
def test_merge_processed_content_tokens_treats_none_as_contagious():
assert merge_processed_content_tokens(5, 7) == 12
assert merge_processed_content_tokens(0, 0) == 0
# Unknown on either side makes the total unknown — not zero, and not the
# other side's value, both of which would under-bill the content.
assert merge_processed_content_tokens(None, 10) is None
assert merge_processed_content_tokens(10, None) is None
@pytest.mark.parametrize(
("rel_path", "name", "allowed"),
[(p, n, a) for p, names, a in NAMED_RESULT_FAMILIES for n in names],
)
def test_named_result_families_keep_their_type(rel_path, name, allowed):
"""The other converted families must not drift back to bare tuples either.
Each of these is a set of interchangeable implementations three extraction
routes that dispatch to one another, three tag-clause builders that compose
each other's parameter offsets, an abstract graph retriever and its
implementation. A tuple lets any one of them silently disagree with its
siblings about arity or order.
"""
func = _function_nodes(rel_path, [name])[name]
assert func.returns is not None, f"{name}() has no return annotation"
annotation = ast.unparse(func.returns).replace('"', "").replace("'", "")
assert annotation in allowed or annotation in {f"{a} | None" for a in allowed}, (
f"{rel_path}::{name}() returns {annotation!r}, expected one of {sorted(allowed)}"
)
offenders = [(r.lineno, ast.unparse(r.value)[:60]) for r in _own_returns(func) if isinstance(r.value, ast.Tuple)]
assert not offenders, f"{rel_path}::{name}() returns bare tuples at {offenders}"
@@ -27,7 +27,7 @@ def test_build_structured_delta_prompt_truncates_huge_document():
def test_fit_structured_delta_keeps_small_prompt_unchanged():
doc_out, cand_out, facts_out, truncated = _fit_structured_delta_prompt_parts(
fitted = _fit_structured_delta_prompt_parts(
source_query="q",
current_document_json='{"sections": []}',
candidate_markdown="hello",
@@ -36,7 +36,34 @@ def test_fit_structured_delta_keeps_small_prompt_unchanged():
task_footer="## Task\nDo it.",
max_input_tokens=24_000,
)
assert not truncated
assert doc_out == '{"sections": []}'
assert cand_out == "hello"
assert facts_out == "one line"
assert not fitted.truncated
assert fitted.document_json == '{"sections": []}'
assert fitted.candidate == "hello"
assert fitted.facts == "one line"
def test_retraction_prompt_does_not_transpose_surviving_and_retracted():
"""The two fact lists must land under their own headings.
``build_structured_retraction_prompt`` reuses ``_fit_structured_delta_prompt_parts``
with the surviving facts in the ``candidate`` slot and the retracted ones in the
``facts`` slot. Both are ``str``, so transposing them type-checks and produces a
grammatical prompt one that tells the model to strip content resting on facts
that are still valid and keep content resting on facts that were withdrawn.
This asserts the mapping directly, since no type can.
"""
from hindsight_api.engine.reflect.prompts import build_structured_retraction_prompt
prompt = build_structured_retraction_prompt(
current_document_json='{"sections": []}',
retracted_facts=[{"id": "r1", "text": "WITHDRAWN_MARKER", "type": "world", "context": ""}],
surviving_facts=[{"id": "v1", "text": "STILL_VALID_MARKER", "type": "world", "context": ""}],
source_query="topic",
)
surviving_heading = prompt.index("## STILL-SUPPORTED FACTS")
retracted_heading = prompt.index("## RETRACTED FACTS")
assert surviving_heading < prompt.index("STILL_VALID_MARKER") < retracted_heading, (
"the surviving fact must appear under STILL-SUPPORTED, before the RETRACTED heading"
)
assert prompt.index("WITHDRAWN_MARKER") > retracted_heading, "the retracted fact must appear under RETRACTED FACTS"
@@ -133,9 +133,9 @@ def test_unmatched_token_stays_unsatisfiable_not_empty():
[TagGroupLeaf(tags=["nosuchtag"], match="any_strict", resolve="fuzzy")],
VOCABULARY,
)
clause, params, _ = build_tag_groups_where_clause(resolved, 1)
assert clause != ""
assert params == [["nosuchtag"]]
built = build_tag_groups_where_clause(resolved, 1)
assert built.sql != ""
assert built.params == [["nosuchtag"]]
assert filter_results_by_tag_groups([_Result(["typescript"])], resolved) == []
@@ -172,7 +172,10 @@ class TestTagsWhereClauseBuilder:
def test_tags_where_clause_exact_empty_scope_keeps_param_offset(self, tags):
"""The parameterized builder must not consume a bind index for the empty scope,
so following clauses stay aligned with their params."""
clause, params, next_offset = build_tags_where_clause(tags, param_offset=4, match="exact")
built = build_tags_where_clause(tags, param_offset=4, match="exact")
clause = built.sql
params = built.params
next_offset = built.next_param_offset
assert clause == "AND (tags IS NULL OR tags = '{}')"
assert params == []
assert next_offset == 4
@@ -379,14 +382,20 @@ class TestBuildTagGroupsWhereClause:
def test_none_returns_empty(self):
"""None tag_groups returns empty clause."""
clause, params, next_offset = build_tag_groups_where_clause(None, 3)
built = build_tag_groups_where_clause(None, 3)
clause = built.sql
params = built.params
next_offset = built.next_param_offset
assert clause == ""
assert params == []
assert next_offset == 3
def test_empty_list_returns_empty(self):
"""Empty tag_groups list returns empty clause."""
clause, params, next_offset = build_tag_groups_where_clause([], 3)
built = build_tag_groups_where_clause([], 3)
clause = built.sql
params = built.params
next_offset = built.next_param_offset
assert clause == ""
assert params == []
assert next_offset == 3
@@ -394,7 +403,10 @@ class TestBuildTagGroupsWhereClause:
def test_single_leaf_any_strict(self):
"""Single any_strict leaf generates correct SQL."""
groups = [TagGroupLeaf(tags=["step:5", "step:8"], match="any_strict")]
clause, params, next_offset = build_tag_groups_where_clause(groups, 3)
built = build_tag_groups_where_clause(groups, 3)
clause = built.sql
params = built.params
next_offset = built.next_param_offset
assert clause.startswith("AND ")
assert "$3" in clause
assert "IS NOT NULL" in clause
@@ -406,7 +418,10 @@ class TestBuildTagGroupsWhereClause:
def test_single_leaf_all_strict(self):
"""Single all_strict leaf generates @> operator."""
groups = [TagGroupLeaf(tags=["user:alice"], match="all_strict")]
clause, params, next_offset = build_tag_groups_where_clause(groups, 1)
built = build_tag_groups_where_clause(groups, 1)
clause = built.sql
params = built.params
next_offset = built.next_param_offset
assert "@>" in clause
assert "IS NOT NULL" in clause
assert params == [["user:alice"]]
@@ -415,7 +430,10 @@ class TestBuildTagGroupsWhereClause:
def test_single_leaf_any_includes_untagged(self):
"""Single any (non-strict) leaf generates NULL-inclusive clause."""
groups = [TagGroupLeaf(tags=["user:alice"], match="any")]
clause, params, next_offset = build_tag_groups_where_clause(groups, 1)
built = build_tag_groups_where_clause(groups, 1)
clause = built.sql
params = built.params
next_offset = built.next_param_offset
assert "IS NULL" in clause
assert "= '{}'" in clause
assert "&&" in clause
@@ -434,7 +452,10 @@ class TestBuildTagGroupsWhereClause:
}
)
]
clause, params, next_offset = build_tag_groups_where_clause(groups, 3)
built = build_tag_groups_where_clause(groups, 3)
clause = built.sql
params = built.params
next_offset = built.next_param_offset
assert "AND" in clause
assert "$3" in clause
assert "$4" in clause
@@ -455,7 +476,10 @@ class TestBuildTagGroupsWhereClause:
}
)
]
clause, params, next_offset = build_tag_groups_where_clause(groups, 1)
built = build_tag_groups_where_clause(groups, 1)
clause = built.sql
params = built.params
next_offset = built.next_param_offset
assert "OR" in clause
assert "$1" in clause
assert "$2" in clause
@@ -465,7 +489,10 @@ class TestBuildTagGroupsWhereClause:
def test_not_wraps_with_not(self):
"""NOT group wraps child clause with NOT."""
groups = [TagGroupNot.model_validate({"not": {"tags": ["archived"], "match": "any_strict"}})]
clause, params, next_offset = build_tag_groups_where_clause(groups, 2)
built = build_tag_groups_where_clause(groups, 2)
clause = built.sql
params = built.params
next_offset = built.next_param_offset
assert "NOT" in clause
assert "$2" in clause
assert len(params) == 1
@@ -488,7 +515,10 @@ class TestBuildTagGroupsWhereClause:
}
)
]
clause, params, next_offset = build_tag_groups_where_clause(groups, 1)
built = build_tag_groups_where_clause(groups, 1)
clause = built.sql
params = built.params
next_offset = built.next_param_offset
assert "AND" in clause
assert "OR" in clause
assert len(params) == 3
@@ -507,7 +537,10 @@ class TestBuildTagGroupsWhereClause:
}
)
]
clause, params, next_offset = build_tag_groups_where_clause(groups, 5)
built = build_tag_groups_where_clause(groups, 5)
clause = built.sql
params = built.params
next_offset = built.next_param_offset
assert "$5" in clause
assert "$6" in clause
assert "$7" in clause
@@ -517,7 +550,10 @@ class TestBuildTagGroupsWhereClause:
def test_table_alias_applied_to_leaves(self):
"""Table alias is prefixed to column name in all leaf clauses."""
groups = [TagGroupLeaf(tags=["user:alice"], match="any_strict")]
clause, params, next_offset = build_tag_groups_where_clause(groups, 1, table_alias="mu.")
built = build_tag_groups_where_clause(groups, 1, table_alias="mu.")
clause = built.sql
params = built.params
next_offset = built.next_param_offset
assert "mu.tags" in clause
def test_table_alias_propagates_to_nested(self):
@@ -532,7 +568,10 @@ class TestBuildTagGroupsWhereClause:
}
)
]
clause, params, next_offset = build_tag_groups_where_clause(groups, 1, table_alias="mu.")
built = build_tag_groups_where_clause(groups, 1, table_alias="mu.")
clause = built.sql
params = built.params
next_offset = built.next_param_offset
# Each leaf of type any_strict references mu.tags three times (IS NOT NULL, != '{}', &&)
# We verify that 'tags' without alias is NOT present, proving the alias is always used
assert "mu.tags" in clause
@@ -548,7 +587,10 @@ class TestBuildTagGroupsWhereClause:
TagGroupLeaf(tags=["step:5"], match="any_strict"),
TagGroupLeaf(tags=["user:ep_42"], match="all_strict"),
]
clause, params, next_offset = build_tag_groups_where_clause(groups, 1)
built = build_tag_groups_where_clause(groups, 1)
clause = built.sql
params = built.params
next_offset = built.next_param_offset
# Should start with AND and have two param refs joined by AND
assert clause.startswith("AND ")
assert " AND " in clause[4:] # after the leading "AND "
@@ -560,7 +602,10 @@ class TestBuildTagGroupsWhereClause:
def test_exact_leaf_empty_scope_matches_untagged_only(self):
"""An exact leaf with [] becomes an untagged-only clause with no bind param."""
groups = [TagGroupLeaf(tags=[], match="exact")]
clause, params, next_offset = build_tag_groups_where_clause(groups, 5)
built = build_tag_groups_where_clause(groups, 5)
clause = built.sql
params = built.params
next_offset = built.next_param_offset
assert "IS NULL" in clause
assert "= '{}'" in clause
assert "$5" not in clause # param-free
@@ -17,6 +17,7 @@ from types import SimpleNamespace
import pytest
import hindsight_api.engine.search.retrieval as retrieval_module
from hindsight_api.engine.search.types import GraphRetrieval
from hindsight_api.engine.search.retrieval import _select_with_temporal_coverage, retrieve_temporal_combined_sql
from hindsight_api.engine.task_backend import fq_table
@@ -194,7 +195,7 @@ async def test_min_semantic_does_not_tighten_temporal_seed_threshold(monkeypatch
class FakeGraphRetriever:
async def retrieve(self, **kwargs):
graph_call_kwargs.append(set(kwargs))
return [], None
return GraphRetrieval([], None)
fake_config = SimpleNamespace(graph_seed_min_similarity=0.3, temporal_semantic_min_similarity=0.24)
+4 -2
View File
@@ -1593,13 +1593,14 @@ class TestRetainCompletedWebhook:
"""
from hindsight_api.engine.response_models import TokenUsage
from hindsight_api.engine.retain import fact_extraction
from hindsight_api.engine.retain.types import ExtractionResult
bank_id = f"wh-zerofact-{uuid.uuid4().hex[:8]}"
webhook_id = uuid.uuid4()
original_manager = memory._webhook_manager
async def _extract_no_facts(*args, **kwargs):
return [], [], TokenUsage()
return ExtractionResult([], [], TokenUsage())
try:
memory._webhook_manager = WebhookManager(backend=memory._backend, global_webhooks=[])
@@ -1709,13 +1710,14 @@ class TestRetainCompletedWebhook:
"""
from hindsight_api.engine.response_models import TokenUsage
from hindsight_api.engine.retain import fact_extraction
from hindsight_api.engine.retain.types import ExtractionResult
bank_id = f"wh-zerocount-{uuid.uuid4().hex[:8]}"
webhook_id = uuid.uuid4()
original_manager = memory._webhook_manager
async def _extract_no_facts(*args, **kwargs):
return [], [], TokenUsage()
return ExtractionResult([], [], TokenUsage())
try:
memory._webhook_manager = WebhookManager(backend=memory._backend, global_webhooks=[])