Commit 75a26eb5 authored by Kantz's avatar Kantz
Browse files

anpassung aug neue preprocessing struktur

parent a27a4ff5
......@@ -21,6 +21,7 @@ def _apply_selected_subsection(
def _retrieve_context_for_subsection(state: base.ChatState, query_text: str) -> int:
# Deprecated: socratic retrieval still uses subsection refs until it is migrated to the new parent model.
refs = task_catalog.get_selected_subsection_refs(state.sheet)
if not refs:
return 0
......
......@@ -55,6 +55,7 @@ def _ensure_context_task_fields(state: base.ChatState, query_text: str) -> tuple
def _retrieve_context_for_task(state: base.ChatState, query_text: str) -> int:
# Deprecated: this still uses subsection refs until task retrieval is migrated to the new parent model.
refs = task_catalog.get_selected_task_subsection_refs(state.sheet)
if not refs:
return 0
......
......@@ -12,6 +12,7 @@ class _SourceIDLike(Protocol):
chapter_title: str | None
section_title: str | None
subsection_title: str | None
subsubsection_title: str | None
title: str | None
doc_type: str | None
......@@ -27,6 +28,7 @@ def _source_id_key(source_id: _SourceIDLike) -> str:
source_id.chapter_title or "",
source_id.section_title or "",
source_id.subsection_title or "",
source_id.subsubsection_title or "",
source_id.title or "",
source_id.doc_type or "",
]
......@@ -58,6 +60,7 @@ def _build_source_index(sources: Iterable[_SourceLike]) -> Tuple[Dict[str, str],
source.source_id.chapter_title,
source.source_id.section_title,
source.source_id.subsection_title,
source.source_id.subsubsection_title,
source.source_id.title,
]
for start in range(len(parts)):
......
......@@ -19,6 +19,7 @@ def retrieve(
chapter_index: int | None = None,
section_index: int | None = None,
subsection_index: int | None = None,
subsubsection_index: int | None = None,
source_type_filter: list[str] | None = None,
expand_links: bool = True,
neighbor_expand: int = 0,
......@@ -32,6 +33,7 @@ def retrieve(
chapter_index=chapter_index,
section_index=section_index,
subsection_index=subsection_index,
subsubsection_index=subsubsection_index,
source_type_filter=source_type_filter,
expand_links=expand_links,
neighbor_expand=neighbor_expand,
......@@ -45,6 +47,7 @@ def retrieve(
chapter_index=chapter_index,
section_index=section_index,
subsection_index=subsection_index,
subsubsection_index=subsubsection_index,
source_type_filter=source_type_filter,
expand_links=expand_links,
neighbor_expand=neighbor_expand,
......@@ -60,6 +63,7 @@ def retrieve_with_subsections(
chapter_index: int | None = None,
section_index: int | None = None,
subsection_index: int | None = None,
subsubsection_index: int | None = None,
source_type_filter: list[str] | None = None,
expand_links: bool = True,
neighbor_expand: int = 0,
......@@ -74,6 +78,7 @@ def retrieve_with_subsections(
chapter_index=chapter_index,
section_index=section_index,
subsection_index=subsection_index,
subsubsection_index=subsubsection_index,
source_type_filter=source_type_filter,
expand_links=expand_links,
neighbor_expand=neighbor_expand,
......@@ -88,6 +93,7 @@ def retrieve_with_subsections(
chapter_index=chapter_index,
section_index=section_index,
subsection_index=subsection_index,
subsubsection_index=subsubsection_index,
source_type_filter=source_type_filter,
expand_links=expand_links,
neighbor_expand=neighbor_expand,
......@@ -98,6 +104,7 @@ def retrieve_for_subsections(
pg_url: str,
subsection_refs: list[vector_store.SubsectionRef] | None = None,
) -> List[Source]:
# Deprecated: subsection-ref retrieval is kept only for legacy task/socratic flows.
return vector_store.load_children_for_subsections(
pg_url=pg_url,
subsection_refs=subsection_refs,
......
......@@ -62,6 +62,7 @@ def stable_uid(doc_type: str, path: str, meta: Dict[str, Any]) -> str:
chapter = meta.get("chapter_index")
sec = meta.get("section_index")
sub = meta.get("subsection_index")
subsub = meta.get("subsubsection_index")
child = meta.get("child_index")
if doc_type == "chapter" and chapter is not None:
......@@ -70,8 +71,10 @@ def stable_uid(doc_type: str, path: str, meta: Dict[str, Any]) -> str:
raw = f"section|c{int(chapter):03d}|s{int(sec):03d}"
elif doc_type == "subsection" and chapter is not None and sec is not None and sub is not None:
raw = f"subsection|c{int(chapter):03d}|s{int(sec):03d}|ss{int(sub):03d}"
elif doc_type == "child" and chapter is not None and sec is not None and sub is not None and child is not None:
raw = f"child|c{int(chapter):03d}|s{int(sec):03d}|ss{int(sub):03d}|c{int(child):03d}"
elif doc_type == "subsubsection" and chapter is not None and sec is not None and sub is not None and subsub is not None:
raw = f"subsubsection|c{int(chapter):03d}|s{int(sec):03d}|ss{int(sub):03d}|sss{int(subsub):03d}"
elif doc_type == "child" and chapter is not None and sec is not None and sub is not None and subsub is not None and child is not None:
raw = f"child|c{int(chapter):03d}|s{int(sec):03d}|ss{int(sub):03d}|sss{int(subsub):03d}|c{int(child):03d}"
else:
raw = f"{doc_type}|{path}"
......@@ -85,6 +88,7 @@ def load_docs(base_dir: Path) -> List[DocRecord]:
("chapter", [base_dir / "chapters", base_dir / "chapter"]),
("section", [base_dir / "sections", base_dir / "section"]),
("subsection", [base_dir / "subsections", base_dir / "subsection"]),
("subsubsection", [base_dir / "subsubsections", base_dir / "subsubsection"]),
("child", [base_dir / "childs"]),
]
......@@ -133,11 +137,13 @@ CREATE TABLE IF NOT EXISTS docs (
chapter_index INT NULL,
section_index INT NULL,
subsection_index INT NULL,
subsubsection_index INT NULL,
child_index INT NULL,
chapter_title TEXT NULL,
section_title TEXT NULL,
subsection_title TEXT NULL,
subsubsection_title TEXT NULL,
title TEXT NULL,
source_type TEXT NULL,
......@@ -155,6 +161,7 @@ CREATE INDEX IF NOT EXISTS docs_doc_type_idx ON docs(doc_type);
CREATE INDEX IF NOT EXISTS docs_chapter_idx ON docs(chapter_index);
CREATE INDEX IF NOT EXISTS docs_section_idx ON docs(chapter_index, section_index);
CREATE INDEX IF NOT EXISTS docs_subsection_idx ON docs(chapter_index, section_index, subsection_index);
CREATE INDEX IF NOT EXISTS docs_subsubsection_idx ON docs(chapter_index, section_index, subsection_index, subsubsection_index);
"""
......@@ -173,13 +180,13 @@ def init_db(pg_url: str) -> None:
UPSERT_SQL = """
INSERT INTO docs (
uid, doc_type,
chapter_index, section_index, subsection_index, child_index,
chapter_title, section_title, subsection_title, title, source_type,
chapter_index, section_index, subsection_index, subsubsection_index, child_index,
chapter_title, section_title, subsection_title, subsubsection_title, title, source_type,
path, markdown, embedding
) VALUES (
%(uid)s, %(doc_type)s,
%(chapter_index)s, %(section_index)s, %(subsection_index)s, %(child_index)s,
%(chapter_title)s, %(section_title)s, %(subsection_title)s, %(title)s, %(source_type)s,
%(chapter_index)s, %(section_index)s, %(subsection_index)s, %(subsubsection_index)s, %(child_index)s,
%(chapter_title)s, %(section_title)s, %(subsection_title)s, %(subsubsection_title)s, %(title)s, %(source_type)s,
%(path)s, %(markdown)s, %(embedding)s
)
ON CONFLICT (uid) DO UPDATE SET
......@@ -187,10 +194,12 @@ ON CONFLICT (uid) DO UPDATE SET
chapter_index = EXCLUDED.chapter_index,
section_index = EXCLUDED.section_index,
subsection_index = EXCLUDED.subsection_index,
subsubsection_index = EXCLUDED.subsubsection_index,
child_index = EXCLUDED.child_index,
chapter_title = EXCLUDED.chapter_title,
section_title = EXCLUDED.section_title,
subsection_title = EXCLUDED.subsection_title,
subsubsection_title = EXCLUDED.subsubsection_title,
title = EXCLUDED.title,
source_type = EXCLUDED.source_type,
path = EXCLUDED.path,
......@@ -228,10 +237,12 @@ def upsert_docs(pg_url: str, docs: List[DocRecord], embeddings: List[List[float]
"chapter_index": _meta_int(m, "chapter_index"),
"section_index": _meta_int(m, "section_index"),
"subsection_index": _meta_int(m, "subsection_index"),
"subsubsection_index": _meta_int(m, "subsubsection_index"),
"child_index": _meta_int(m, "child_index"),
"chapter_title": m.get("chapter_title"),
"section_title": m.get("section_title"),
"subsection_title": m.get("subsection_title"),
"subsubsection_title": m.get("subsubsection_title"),
"title": m.get("title"),
"source_type": m.get("source_type"),
"path": doc.path,
......@@ -292,10 +303,12 @@ def _row_to_retrieved(row: Dict[str, Any], source_type: Optional[str] = None) ->
"chapter_index": row["chapter_index"],
"section_index": row["section_index"],
"subsection_index": row["subsection_index"],
"subsubsection_index": row["subsubsection_index"],
"child_index": row["child_index"],
"chapter_title": row["chapter_title"],
"section_title": row["section_title"],
"subsection_title": row["subsection_title"],
"subsubsection_title": row["subsubsection_title"],
"title": row["title"],
"source_type": row["source_type"] or source_type,
"path": row["path"],
......@@ -319,11 +332,11 @@ class RetrievalPipelineConfig:
enable_dominant_scope: bool = True
enable_scoped_child_search: bool = True
enable_context_docs: bool = True
dominance_level: Literal["subsection", "section"] = "subsection"
scope_fill_k: int = 5
chapter_index: Optional[int] = None
section_index: Optional[int] = None
subsection_index: Optional[int] = None
subsubsection_index: Optional[int] = None
source_type_filter: Optional[List[str]] = None
......@@ -332,7 +345,8 @@ class DominantScope:
chapter_index: int
section_index: int
subsection_index: Optional[int]
level: Literal["subsection", "section"]
subsubsection_index: Optional[int]
level: Literal["section", "subsection", "subsubsection"]
def build_default_pipeline_config(
......@@ -340,6 +354,7 @@ def build_default_pipeline_config(
chapter_index: Optional[int] = None,
section_index: Optional[int] = None,
subsection_index: Optional[int] = None,
subsubsection_index: Optional[int] = None,
source_type_filter: Optional[List[str]] = None,
expand_links: bool = True,
neighbor_expand: int = 0,
......@@ -354,11 +369,11 @@ def build_default_pipeline_config(
enable_dominant_scope=True,
enable_scoped_child_search=True,
enable_context_docs=True,
dominance_level="subsection",
scope_fill_k=max(k, 5),
chapter_index=chapter_index,
section_index=section_index,
subsection_index=subsection_index,
subsubsection_index=subsubsection_index,
source_type_filter=source_type_filter,
)
......@@ -386,6 +401,10 @@ def run_global_child_vector_search(
where.append("subsection_index = %(subsection_index)s")
params["subsection_index"] = config.subsection_index
if config.subsubsection_index is not None:
where.append("subsubsection_index = %(subsubsection_index)s")
params["subsubsection_index"] = config.subsubsection_index
if config.source_type_filter:
where.append("source_type = ANY(%(source_type_filter)s)")
params["source_type_filter"] = config.source_type_filter
......@@ -394,8 +413,8 @@ def run_global_child_vector_search(
sql = f"""
SELECT
uid, doc_type,
chapter_index, section_index, subsection_index, child_index,
chapter_title, section_title, subsection_title, title, source_type,
chapter_index, section_index, subsection_index, subsubsection_index, child_index,
chapter_title, section_title, subsection_title, subsubsection_title, title, source_type,
path, markdown,
1 - (embedding <=> %(qvec)s) AS score
FROM docs
......@@ -424,51 +443,94 @@ def select_top_k_unique(values: List[Retrieved], wanted_k: int) -> List[Retrieve
return selected
def _normalized_level_index(value: Any) -> int:
if value is None:
return 0
try:
return int(value)
except Exception:
return 0
def _build_scope_from_indices(
chapter_index: Any,
section_index: Any,
subsection_index: Any,
subsubsection_index: Any,
) -> Optional[DominantScope]:
if chapter_index is None or section_index is None:
return None
sub = _normalized_level_index(subsection_index)
subsub = _normalized_level_index(subsubsection_index)
if sub <= 0:
return DominantScope(
chapter_index=int(chapter_index),
section_index=int(section_index),
subsection_index=None,
subsubsection_index=None,
level="section",
)
if subsub <= 0:
return DominantScope(
chapter_index=int(chapter_index),
section_index=int(section_index),
subsection_index=sub,
subsubsection_index=None,
level="subsection",
)
return DominantScope(
chapter_index=int(chapter_index),
section_index=int(section_index),
subsection_index=sub,
subsubsection_index=subsub,
level="subsubsection",
)
def _scope_key(scope: DominantScope) -> Tuple[int, ...]:
if scope.level == "section":
return (scope.chapter_index, scope.section_index)
if scope.level == "subsection":
return (scope.chapter_index, scope.section_index, int(scope.subsection_index or 0))
return (
scope.chapter_index,
scope.section_index,
int(scope.subsection_index or 0),
int(scope.subsubsection_index or 0),
)
def select_dominant_scope(
children_direct: List[Retrieved],
level: Literal["subsection", "section"] = "subsection",
) -> Optional[DominantScope]:
if not children_direct:
return None
counts: Dict[Tuple[int, ...], Tuple[int, float]] = {}
counts: Dict[Tuple[str, Tuple[int, ...]], Tuple[DominantScope, int, float]] = {}
for child in children_direct:
chapter_index = child.metadata.get("chapter_index")
section_index = child.metadata.get("section_index")
subsection_index = child.metadata.get("subsection_index")
if chapter_index is None or section_index is None:
scope = _build_scope_from_indices(
child.metadata.get("chapter_index"),
child.metadata.get("section_index"),
child.metadata.get("subsection_index"),
child.metadata.get("subsubsection_index"),
)
if scope is None:
continue
if level == "subsection":
if subsection_index is None:
continue
key = (int(chapter_index), int(section_index), int(subsection_index))
else:
key = (int(chapter_index), int(section_index))
count, score_sum = counts.get(key, (0, 0.0))
counts[key] = (count + 1, score_sum + float(child.score))
key = (scope.level, _scope_key(scope))
_, count, score_sum = counts.get(key, (scope, 0, 0.0))
counts[key] = (scope, count + 1, score_sum + float(child.score))
if not counts:
return None
def _rank(item: Tuple[Tuple[int, ...], Tuple[int, float]]) -> Tuple[int, float, Tuple[int, ...]]:
key, (count, score_sum) = item
def _rank(item: Tuple[Tuple[str, Tuple[int, ...]], Tuple[DominantScope, int, float]]) -> Tuple[int, float, Tuple[int, ...]]:
_, (scope, count, score_sum) = item
avg_score = score_sum / count if count else 0.0
return (count, avg_score, tuple([-part for part in key]))
return (count, avg_score, tuple([-part for part in _scope_key(scope)]))
winner_key, _ = max(counts.items(), key=_rank)
if level == "subsection":
return DominantScope(
chapter_index=winner_key[0],
section_index=winner_key[1],
subsection_index=winner_key[2],
level="subsection",
)
return DominantScope(
chapter_index=winner_key[0],
section_index=winner_key[1],
subsection_index=None,
level="section",
)
_, (winner_scope, _, _) = max(counts.items(), key=_rank)
return winner_scope
def fetch_scope_context_docs(
......@@ -478,16 +540,17 @@ def fetch_scope_context_docs(
chapter_docs: List[Retrieved] = []
section_docs: List[Retrieved] = []
subsection_docs: List[Retrieved] = []
subsubsection_docs: List[Retrieved] = []
with psycopg.connect(pg_url, row_factory=dict_row) as conn:
with conn.cursor() as cur:
if scope.level == "subsection" and scope.subsection_index is not None:
if scope.level in ("subsection", "subsubsection") and scope.subsection_index is not None:
cur.execute(
"""
SELECT
d.uid, d.doc_type,
d.chapter_index, d.section_index, d.subsection_index, d.child_index,
d.chapter_title, d.section_title, d.subsection_title, d.title, d.source_type,
d.chapter_index, d.section_index, d.subsection_index, d.subsubsection_index, d.child_index,
d.chapter_title, d.section_title, d.subsection_title, d.subsubsection_title, d.title, d.source_type,
d.path, d.markdown,
1.0 AS score
FROM docs d
......@@ -502,12 +565,33 @@ def fetch_scope_context_docs(
subsection_docs = [_row_to_retrieved(
row, source_type="subsection") for row in cur.fetchall()]
if scope.level == "subsubsection" and scope.subsection_index is not None and scope.subsubsection_index is not None:
cur.execute(
"""
SELECT
d.uid, d.doc_type,
d.chapter_index, d.section_index, d.subsection_index, d.subsubsection_index, d.child_index,
d.chapter_title, d.section_title, d.subsection_title, d.subsubsection_title, d.title, d.source_type,
d.path, d.markdown,
1.0 AS score
FROM docs d
WHERE d.doc_type = 'subsubsection'
AND d.chapter_index = %(cpt)s
AND d.section_index = %(sec)s
AND d.subsection_index = %(sub)s
AND d.subsubsection_index = %(subsub)s
""",
{"cpt": scope.chapter_index, "sec": scope.section_index, "sub": scope.subsection_index, "subsub": scope.subsubsection_index},
)
subsubsection_docs = [_row_to_retrieved(
row, source_type="subsubsection") for row in cur.fetchall()]
cur.execute(
"""
SELECT
d.uid, d.doc_type,
d.chapter_index, d.section_index, d.subsection_index, d.child_index,
d.chapter_title, d.section_title, d.subsection_title, d.title, d.source_type,
d.chapter_index, d.section_index, d.subsection_index, d.subsubsection_index, d.child_index,
d.chapter_title, d.section_title, d.subsection_title, d.subsubsection_title, d.title, d.source_type,
d.path, d.markdown,
1.0 AS score
FROM docs d
......@@ -524,8 +608,8 @@ def fetch_scope_context_docs(
"""
SELECT
d.uid, d.doc_type,
d.chapter_index, d.section_index, d.subsection_index, d.child_index,
d.chapter_title, d.section_title, d.subsection_title, d.title, d.source_type,
d.chapter_index, d.section_index, d.subsection_index, d.subsubsection_index, d.child_index,
d.chapter_title, d.section_title, d.subsection_title, d.subsubsection_title, d.title, d.source_type,
d.path, d.markdown,
1.0 AS score
FROM docs d
......@@ -541,6 +625,7 @@ def fetch_scope_context_docs(
"chapters": chapter_docs,
"sections": section_docs,
"subsections": subsection_docs,
"subsubsections": subsubsection_docs,
}
......@@ -561,16 +646,19 @@ def run_scoped_child_vector_search(
"sec": scope.section_index,
"fill_k": max(1, int(fill_k)),
}
if scope.level == "subsection" and scope.subsection_index is not None:
if scope.level in ("subsection", "subsubsection") and scope.subsection_index is not None:
where.append("d.subsection_index = %(sub)s")
params["sub"] = scope.subsection_index
if scope.level == "subsubsection" and scope.subsubsection_index is not None:
where.append("d.subsubsection_index = %(subsub)s")
params["subsub"] = scope.subsubsection_index
where_sql = " AND ".join(where)
sql = f"""
SELECT
d.uid, d.doc_type,
d.chapter_index, d.section_index, d.subsection_index, d.child_index,
d.chapter_title, d.section_title, d.subsection_title, d.title, d.source_type,
d.chapter_index, d.section_index, d.subsection_index, d.subsubsection_index, d.child_index,
d.chapter_title, d.section_title, d.subsection_title, d.subsubsection_title, d.title, d.source_type,
d.path, d.markdown,
1 - (d.embedding <=> %(qvec)s) AS score
FROM docs d
......@@ -594,27 +682,29 @@ def expand_neighbor_children(
if neighbor_expand <= 0:
return []
wanted: set[Tuple[int, int, int, int]] = set()
wanted: set[Tuple[int, int, int, int, int]] = set()
for child in children:
cpti = child.metadata.get("chapter_index")
si = child.metadata.get("section_index")
ssi = child.metadata.get("subsection_index")
sssi = child.metadata.get("subsubsection_index")
ci = child.metadata.get("child_index")
if cpti is None or si is None or ssi is None or ci is None:
if cpti is None or si is None or ssi is None or sssi is None or ci is None:
continue
for dx in range(-neighbor_expand, neighbor_expand + 1):
if dx == 0:
continue
wanted.add((int(cpti), int(si), int(ssi), int(ci) + dx))
wanted.add((int(cpti), int(si), int(ssi), int(sssi), int(ci) + dx))
if not wanted:
return []
quadruple = sorted(wanted)
cpt_arr = [a for (a, b, c, cidx) in quadruple]
sec_arr = [b for (a, b, c, cidx) in quadruple]
sub_arr = [c for (a, b, c, cidx) in quadruple]
child_arr = [cidx for (a, b, c, cidx) in quadruple]
quintuple = sorted(wanted)
cpt_arr = [a for (a, b, c, d, cidx) in quintuple]
sec_arr = [b for (a, b, c, d, cidx) in quintuple]
sub_arr = [c for (a, b, c, d, cidx) in quintuple]
subsub_arr = [d for (a, b, c, d, cidx) in quintuple]
child_arr = [cidx for (a, b, c, d, cidx) in quintuple]
with psycopg.connect(pg_url, row_factory=dict_row) as conn:
with conn.cursor() as cur:
......@@ -622,16 +712,16 @@ def expand_neighbor_children(
"""
SELECT
d.uid, d.doc_type,
d.chapter_index, d.section_index, d.subsection_index, d.child_index,
d.chapter_title, d.section_title, d.subsection_title, d.title, d.source_type,
d.chapter_index, d.section_index, d.subsection_index, d.subsubsection_index, d.child_index,
d.chapter_title, d.section_title, d.subsection_title, d.subsubsection_title, d.title, d.source_type,
d.path, d.markdown,
0 AS score
FROM docs d
JOIN unnest(%(cpt_arr)s::int[],%(sec_arr)s::int[], %(sub_arr)s::int[], %(child_arr)s::int[]) AS u(cpt, sec, sub, child)
ON d.chapter_index = u.cpt AND d.section_index = u.sec AND d.subsection_index = u.sub AND d.child_index = u.child
JOIN unnest(%(cpt_arr)s::int[],%(sec_arr)s::int[], %(sub_arr)s::int[], %(subsub_arr)s::int[], %(child_arr)s::int[]) AS u(cpt, sec, sub, subsub, child)
ON d.chapter_index = u.cpt AND d.section_index = u.sec AND d.subsection_index = u.sub AND d.subsubsection_index = u.subsub AND d.child_index = u.child
WHERE d.doc_type = 'child'
""",
{"cpt_arr": cpt_arr, "sec_arr": sec_arr, "sub_arr": sub_arr,
{"cpt_arr": cpt_arr, "sec_arr": sec_arr, "sub_arr": sub_arr, "subsub_arr": subsub_arr,
"child_arr": child_arr},
)
rows = cur.fetchall()
......@@ -687,6 +777,7 @@ def run_child_retrieval_pipeline(
"children_expanded": [],
"chapters": [],
"subsections": [],
"subsubsections": [],
"sections": [],
"neighbors": [],
}
......@@ -698,13 +789,14 @@ def run_child_retrieval_pipeline(
dominant_scope: Optional[DominantScope] = None
if config.expand_links and config.enable_dominant_scope:
dominant_scope = select_dominant_scope(children_direct, config.dominance_level)
dominant_scope = select_dominant_scope(children_direct)
if dominant_scope and config.expand_links and config.enable_context_docs:
context_docs = fetch_scope_context_docs(pg_url, dominant_scope)
groups["chapters"] = context_docs.get("chapters", [])
groups["sections"] = context_docs.get("sections", [])
groups["subsections"] = context_docs.get("subsections", [])
groups["subsubsections"] = context_docs.get("subsubsections", [])
if dominant_scope and config.expand_links and config.enable_scoped_child_search:
additional_children = run_scoped_child_vector_search(
......@@ -739,6 +831,7 @@ def retrieve(
chapter_index: Optional[int] = None,
section_index: Optional[int] = None,
subsection_index: Optional[int] = None,
subsubsection_index: Optional[int] = None,
source_type_filter: Optional[List[str]] = None,
expand_links: bool = True,
neighbor_expand: int = 0,
......@@ -748,6 +841,7 @@ def retrieve(
chapter_index=chapter_index,
section_index=section_index,
subsection_index=subsection_index,
subsubsection_index=subsubsection_index,
source_type_filter=source_type_filter,
expand_links=expand_links,
neighbor_expand=neighbor_expand,
......@@ -766,12 +860,14 @@ def retrieve(
# --------------------------------------------------------------------------------------------------------------------
# Deprecated: subsection refs only represent the legacy subsection-centric task/socratic flow.
SubsectionRef = Tuple[int, int, int]
def _normalize_subsection_refs(
subsection_refs: Optional[List[SubsectionRef]],
) -> List[SubsectionRef]:
# Deprecated compatibility helper for subsection-only retrieval.
if not subsection_refs:
return []
normalized = {
......@@ -784,6 +880,7 @@ def load_children_for_subsections(
pg_url: str,
subsection_refs: Optional[List[SubsectionRef]],
) -> List[Source]:
# Deprecated compatibility path for subsection-only task/socratic retrieval.
refs = _normalize_subsection_refs(subsection_refs)
if not refs:
return []
......@@ -794,8 +891,8 @@ def load_children_for_subsections(
sql = """
SELECT
d.uid, d.doc_type,
d.chapter_index, d.section_index, d.subsection_index, d.child_index,
d.chapter_title, d.section_title, d.subsection_title, d.title, d.source_type,
d.chapter_index, d.section_index, d.subsection_index, d.subsubsection_index, d.child_index,
d.chapter_title, d.section_title, d.subsection_title, d.subsubsection_title, d.title, d.source_type,
d.path, d.markdown,
1.0 AS score
FROM docs d
......@@ -822,13 +919,14 @@ def load_children_for_subsections(
def merge_sources(primary: List[Source], additional: List[Source]) -> List[Source]:
merged: List[Source] = []
index_by_key: Dict[Tuple[str, str, str, str, str, str], int] = {}
index_by_key: Dict[Tuple[str, str, str, str, str, str, str], int] = {}
def source_key(source: Source) -> Tuple[str, str, str, str, str, str]:
def source_key(source: Source) -> Tuple[str, str, str, str, str, str, str]:
return (
source.source_id.doc_type or "",
source.source_id.section_title or "",
source.source_id.subsection_title or "",
source.source_id.subsubsection_title or "",
source.source_id.title or "",
source.source_type or "",
source.markdown or "",
......@@ -857,10 +955,12 @@ def retrieve_with_subsections(
chapter_index: Optional[int] = None,
section_index: Optional[int] = None,
subsection_index: Optional[int] = None,
subsubsection_index: Optional[int] = None,
source_type_filter: Optional[List[str]] = None,
expand_links: bool = True,
neighbor_expand: int = 0,
) -> List[Source]:
# Deprecated compatibility path for subsection-only retrieval composition.
# Oversampling improves recall with ivfflat when additional filters exclude
# close hits. We trim back to k after retrieval.
vector_k = max(k * 4, k + 16)
......@@ -872,6 +972,7 @@ def retrieve_with_subsections(
chapter_index=chapter_index,
section_index=section_index,
subsection_index=subsection_index,
subsubsection_index=subsubsection_index,
source_type_filter=source_type_filter,
expand_links=expand_links,
neighbor_expand=neighbor_expand,
......@@ -919,6 +1020,7 @@ class SourceID(BaseModel):
chapter_title: Optional[str] = None
section_title: Optional[str] = None
subsection_title: Optional[str] = None
subsubsection_title: Optional[str] = None
title: str
doc_type: str
......@@ -927,12 +1029,15 @@ class SourceID(BaseModel):
"chapter_title": self.chapter_title,
"section_title": self.section_title,
"subsection_title": self.subsection_title,
"subsubsection_title": self.subsubsection_title,
"title": self.title,
"doc_type": self.doc_type,
}
def to_string(self) -> str:
string_rep = self.title
if self.subsubsection_title:
string_rep = f"{self.subsubsection_title}|{string_rep}"
if self.subsection_title:
string_rep = f"{self.subsection_title}|{string_rep}"
if self.section_title:
......@@ -951,6 +1056,7 @@ def _retrivla_to_sources(retrievd: Dict[str, List[Retrieved]]) -> List[Source]:
chapter_title=retrival.metadata.get("chapter_title"),
section_title=retrival.metadata.get("section_title"),
subsection_title=retrival.metadata.get("subsection_title"),
subsubsection_title=retrival.metadata.get("subsubsection_title"),
title=retrival.metadata.get("title"),
doc_type=retrival.doc_type
),
......@@ -991,7 +1097,7 @@ def list_subsections(pg_url: str, section_idx: Optional[int] = None) -> List[Dic
section_index, subsection_index,
COALESCE(NULLIF(subsection_title, ''), '') AS subsection_title
FROM docs
WHERE subsection_index IS NOT NULL
WHERE subsection_index IS NOT NULL AND subsection_index > 0
ORDER BY section_index, subsection_index,
CASE WHEN subsection_title IS NULL OR subsection_title = '' THEN 1 ELSE 0 END,
subsection_title
......@@ -1003,7 +1109,7 @@ def list_subsections(pg_url: str, section_idx: Optional[int] = None) -> List[Dic
section_index, subsection_index,
COALESCE(NULLIF(subsection_title, ''), '') AS subsection_title
FROM docs
WHERE subsection_index IS NOT NULL AND section_index = %(sec)s
WHERE subsection_index IS NOT NULL AND subsection_index > 0 AND section_index = %(sec)s
ORDER BY section_index, subsection_index,
CASE WHEN subsection_title IS NULL OR subsection_title = '' THEN 1 ELSE 0 END,
subsection_title
......
......@@ -18,6 +18,8 @@ from app.deterministic_services.vector_store import (
merge_sources,
)
# Deprecated: this module keeps the legacy subsection-centric retrieval path for compatibility.
def retrieve(
pg_url: str,
......@@ -27,13 +29,15 @@ def retrieve(
chapter_index: Optional[int] = None,
section_index: Optional[int] = None,
subsection_index: Optional[int] = None,
subsubsection_index: Optional[int] = None,
source_type_filter: Optional[List[str]] = None,
expand_links: bool = False,
neighbor_expand: int = 0,
) -> List[Source]:
# Parameters kept for drop-in compatibility with child-level retrieve.
# Deprecated compatibility path. Parameters kept for drop-in compatibility with child-level retrieve.
_ = expand_links
_ = neighbor_expand
_ = subsubsection_index
qvec = Vector(embed_query(embedder, query))
......@@ -64,8 +68,8 @@ def retrieve(
sql = f"""
SELECT
uid, doc_type,
chapter_index, section_index, subsection_index, child_index,
chapter_title, section_title, subsection_title, title, source_type,
chapter_index, section_index, subsection_index, subsubsection_index, child_index,
chapter_title, section_title, subsection_title, subsubsection_title, title, source_type,
path, markdown,
1 - (embedding <=> %(qvec)s) AS score
FROM docs
......@@ -94,10 +98,12 @@ def retrieve_with_subsections(
chapter_index: Optional[int] = None,
section_index: Optional[int] = None,
subsection_index: Optional[int] = None,
subsubsection_index: Optional[int] = None,
source_type_filter: Optional[List[str]] = None,
expand_links: bool = False,
neighbor_expand: int = 0,
) -> List[Source]:
# Deprecated compatibility path for subsection-only retrieval composition.
vector_k = max(k * 4, k + 16)
vector_sources = retrieve(
pg_url=pg_url,
......@@ -107,6 +113,7 @@ def retrieve_with_subsections(
chapter_index=chapter_index,
section_index=section_index,
subsection_index=subsection_index,
subsubsection_index=subsubsection_index,
source_type_filter=source_type_filter,
expand_links=expand_links,
neighbor_expand=neighbor_expand,
......
......@@ -24,6 +24,7 @@ def _mk_retrieved(
chapter_index: int,
section_index: int,
subsection_index: int | None,
subsubsection_index: int | None = 0,
child_index: int = 1,
) -> Retrieved:
return Retrieved(
......@@ -34,10 +35,12 @@ def _mk_retrieved(
"chapter_index": chapter_index,
"section_index": section_index,
"subsection_index": subsection_index,
"subsubsection_index": subsubsection_index,
"child_index": child_index,
"chapter_title": "C",
"section_title": "S",
"subsection_title": "SS",
"subsubsection_title": "SSS" if subsubsection_index else "",
"title": uid,
"source_type": "child",
"path": "",
......@@ -91,37 +94,77 @@ class VectorStorePipelineUnitTest(unittest.TestCase):
self.assertTrue(cfg.enable_scoped_child_search)
self.assertTrue(cfg.enable_context_docs)
def test_select_dominant_scope_prefers_count(self) -> None:
def test_select_dominant_scope_prefers_subsection_count(self) -> None:
children = [
_mk_retrieved("a", 0.9, 1, 1, 1),
_mk_retrieved("b", 0.8, 1, 1, 1),
_mk_retrieved("c", 0.95, 1, 1, 2),
]
scope = select_dominant_scope(children, level="subsection")
scope = select_dominant_scope(children)
self.assertIsNotNone(scope)
assert scope is not None
self.assertEqual(scope.level, "subsection")
self.assertEqual((scope.chapter_index, scope.section_index, scope.subsection_index), (1, 1, 1))
def test_select_dominant_scope_resolves_section_parent(self) -> None:
children = [
_mk_retrieved("a", 0.9, 1, 1, 0, 0),
_mk_retrieved("b", 0.8, 1, 1, 0, 0),
_mk_retrieved("c", 0.95, 1, 1, 2, 0),
]
scope = select_dominant_scope(children)
self.assertIsNotNone(scope)
assert scope is not None
self.assertEqual(scope.level, "section")
self.assertEqual((scope.chapter_index, scope.section_index), (1, 1))
self.assertIsNone(scope.subsection_index)
self.assertIsNone(scope.subsubsection_index)
def test_select_dominant_scope_resolves_subsubsection_parent(self) -> None:
children = [
_mk_retrieved("a", 0.9, 1, 1, 2, 1),
_mk_retrieved("b", 0.8, 1, 1, 2, 1),
_mk_retrieved("c", 0.95, 1, 1, 2, 0),
]
scope = select_dominant_scope(children)
self.assertIsNotNone(scope)
assert scope is not None
self.assertEqual(scope.level, "subsubsection")
self.assertEqual((scope.chapter_index, scope.section_index, scope.subsection_index, scope.subsubsection_index), (1, 1, 2, 1))
def test_select_dominant_scope_tiebreak_avg_score(self) -> None:
children = [
_mk_retrieved("a", 0.7, 1, 1, 1),
_mk_retrieved("b", 0.8, 1, 1, 2),
_mk_retrieved("a", 0.7, 1, 1, 1, 0),
_mk_retrieved("b", 0.8, 1, 1, 2, 0),
]
scope = select_dominant_scope(children, level="subsection")
scope = select_dominant_scope(children)
self.assertIsNotNone(scope)
assert scope is not None
self.assertEqual(scope.level, "subsection")
self.assertEqual((scope.chapter_index, scope.section_index, scope.subsection_index), (1, 1, 2))
def test_select_dominant_scope_tiebreak_lexicographic(self) -> None:
children = [
_mk_retrieved("a", 0.8, 2, 1, 1),
_mk_retrieved("b", 0.8, 1, 2, 3),
_mk_retrieved("a", 0.8, 2, 1, 1, 0),
_mk_retrieved("b", 0.8, 1, 2, 3, 0),
]
scope = select_dominant_scope(children, level="subsection")
scope = select_dominant_scope(children)
self.assertIsNotNone(scope)
assert scope is not None
self.assertEqual((scope.chapter_index, scope.section_index, scope.subsection_index), (1, 2, 3))
def test_select_dominant_scope_keeps_exact_parent_levels_separate(self) -> None:
children = [
_mk_retrieved("section-hit", 0.91, 1, 1, 2, 0),
_mk_retrieved("subsub-hit", 0.92, 1, 1, 2, 1),
]
scope = select_dominant_scope(children)
self.assertIsNotNone(scope)
assert scope is not None
self.assertEqual(scope.level, "subsubsection")
self.assertEqual(scope.subsection_index, 2)
self.assertEqual(scope.subsubsection_index, 1)
def test_merge_retrieval_groups_dedup_max_score_and_trim_children(self) -> None:
groups = {
"children_direct": [
......@@ -146,10 +189,55 @@ class VectorStorePipelineUnitTest(unittest.TestCase):
self.assertEqual(merged["neighbors"], [])
def test_expand_neighbor_children_returns_empty_for_zero_expand(self) -> None:
children = [_mk_retrieved("u1", 0.8, 1, 1, 1, child_index=3)]
children = [_mk_retrieved("u1", 0.8, 1, 1, 1, 4, child_index=3)]
result = expand_neighbor_children("postgresql://unused", children, neighbor_expand=0)
self.assertEqual(result, [])
def test_source_id_to_string_includes_subsubsection_title(self) -> None:
source_id = SourceID(
chapter_title="Kapitel",
section_title="Section",
subsection_title="Subsection",
subsubsection_title="Subsubsection",
title="Child",
doc_type="child",
)
self.assertEqual(source_id.to_string(), "[Kapitel|Section|Subsection|Subsubsection|Child|child]")
def test_merge_sources_distinguishes_subsubsection_title(self) -> None:
first = Source(
source_id=SourceID(
chapter_title="Kapitel",
section_title="Section",
subsection_title="Subsection",
subsubsection_title="A",
title="Child A",
doc_type="child",
),
retrieved_as="children_direct",
source_type="child",
score=0.7,
markdown="same-md",
)
second = Source(
source_id=SourceID(
chapter_title="Kapitel",
section_title="Section",
subsection_title="Subsection",
subsubsection_title="B",
title="Child A",
doc_type="child",
),
retrieved_as="children_direct",
source_type="child",
score=0.6,
markdown="same-md",
)
merged = merge_sources([first], [second])
self.assertEqual(len(merged), 2)
if __name__ == "__main__":
unittest.main()
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment