File size: 5,138 Bytes
6bff5d9 b9dfc76 6bff5d9 b9dfc76 6bff5d9 b9dfc76 81e5fe7 0721bb4 81e5fe7 b9dfc76 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """CatalogReader — loads + filters catalog by source_hint.
For typical users (≤50 tables), returns the FULL catalog with no slicing.
Catalog-level search is added later if catalog grows past the limit.
"""
from datetime import UTC, datetime
from typing import Literal
from .models import Catalog, Source
from .store import CatalogStore
SourceHint = Literal["chat", "unstructured", "structured"]
def _filter_sources(catalog: Catalog, source_hint: SourceHint) -> Catalog:
"""Return a copy of `catalog` keeping only the sources matching `source_hint`."""
filtered: list[Source]
if source_hint == "chat":
filtered = []
elif source_hint == "structured":
filtered = [s for s in catalog.sources if s.source_type in {"schema", "tabular"}]
else: # "unstructured"
filtered = [s for s in catalog.sources if s.source_type == "unstructured"]
return catalog.model_copy(update={"sources": filtered})
class CatalogReader:
"""Loads the user's catalog and filters by source_hint.
On miss, returns an empty Catalog (never raises) — query path is
responsible for handling "no data registered yet" gracefully.
Returned Catalog is always a copy; the underlying stored catalog
is never mutated.
"""
def __init__(self, store: CatalogStore) -> None:
self._store = store
async def read(self, user_id: str, source_hint: SourceHint) -> Catalog:
catalog = await self._store.get(user_id)
if catalog is None:
return Catalog(user_id=user_id, generated_at=datetime.now(UTC))
return _filter_sources(catalog, source_hint)
class MemoizingCatalogReader(CatalogReader):
"""Request-scoped CatalogReader that caches each ``read`` by source_hint.
One per request. The same per-user catalog is otherwise fetched from the
catalog DB 4-5x during a single slow-path run (planner load, then
check_data's structured read + check_knowledge's unstructured read, then
retrieve_data's structured read). Wrapping the base reader collapses those
to one round-trip
per distinct source_hint and pins a single consistent snapshot for the whole
request (plan-time and execution-time catalogs can no longer diverge).
"""
def __init__(self, inner: CatalogReader) -> None:
# `read` is fully overridden below and delegates to `inner`, so the parent's
# `_store` is never used — carry it through only so this stays a real
# CatalogReader (any inner with a `read` works, including test fakes).
super().__init__(getattr(inner, "_store", None))
self._inner = inner
self._cache: dict[SourceHint, Catalog] = {}
async def read(self, user_id: str, source_hint: SourceHint) -> Catalog:
cached = self._cache.get(source_hint)
if cached is None:
cached = await self._inner.read(user_id, source_hint)
self._cache[source_hint] = cached
return cached
class AnalysisScopedCatalogReader(CatalogReader):
"""Reads the analysis-scope catalog, falling back to the user-scope reader.
Used by the `check` skill so "what data do I have" inside a room reflects
that analysis's bound sources — structured AND documents — with their real
names. A database shows as "xl test" (analysis-scope) instead of the
auto-generated `postgres_<hash>` placeholder, and documents show at all
(the user-scope catalog holds no `unstructured` sources, so reading them from
user-scope always came back empty). When the analysis has no catalog row
(legacy / not yet bound) or the read fails, it degrades to the wrapped
user-scope reader, so unbound rooms behave exactly as before.
"""
def __init__(self, inner: CatalogReader, analysis_id: str | None) -> None:
# `inner` is a real CatalogReader (constructed at the check call site), so
# its `_store` is the live CatalogStore we need for the analysis read.
super().__init__(inner._store)
self._inner = inner
self._analysis_id = analysis_id
async def read(self, user_id: str, source_hint: SourceHint) -> Catalog:
# Read analysis-scope for BOTH structured and unstructured. Verified via
# the dedorch `data_catalog` table: the analysis-scope rows carry the real
# DB names AND the room's documents (`source_type='unstructured'`), whereas
# the user-scope rows hold only structured sources with `postgres_<hash>`
# placeholder names and NO documents at all — so reading documents from
# user-scope always returned empty ("not listed"). Fall back to the
# user-scope reader only when the analysis has no catalog row (legacy /
# unbound room).
if self._analysis_id:
try:
catalog = await self._store.get_by_analysis(self._analysis_id)
except Exception: # noqa: BLE001 — never block check on the analysis read
catalog = None
if catalog is not None:
return _filter_sources(catalog, source_hint)
return await self._inner.read(user_id, source_hint)
|