Skip to content
GitLab
Commit d93bc3e6 authored by Kantz's avatar Kantz
Browse files

aufräumen des vector stores

parent 9a38c477
Hide whitespace changes
Inline Side-by-side
math-tutor/backend/app/deterministic_services/vector_store.py
View file @ d93bc3e6
...@@ -3,7 +3,7 @@ from __future__ import annotations ...@@ -3,7 +3,7 @@ from __future__ import annotations
import hashlib import hashlib
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import Any, Dict, List, Optional, Protocol, Tuple from typing import Any, Dict, List, Literal, Optional, Protocol, Tuple
import psycopg import psycopg
from psycopg.rows import dict_row from psycopg.rows import dict_row
...@@ -305,42 +305,88 @@ def _row_to_retrieved(row: Dict[str, Any], source_type: Optional[str] = None) -> ...@@ -305,42 +305,88 @@ def _row_to_retrieved(row: Dict[str, Any], source_type: Optional[str] = None) ->
) )
def retrieve( @dataclass
pg_url: str, class RetrievalPipelineConfig:
embedder: EmbeddingLike, k: int
query: str, vector_k: int
expand_links: bool = True
neighbor_expand: int = 0
enable_global_search: bool = True
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
source_type_filter: Optional[List[str]] = None
@dataclass
class DominantScope:
chapter_index: int
section_index: int
subsection_index: Optional[int]
level: Literal["subsection", "section"]
def build_default_pipeline_config(
k: int = 4, k: int = 4,
chapter_index: Optional[int] = None, chapter_index: Optional[int] = None,
section_index: Optional[int] = None, section_index: Optional[int] = None,
subsection_index: Optional[int] = None, subsection_index: Optional[int] = None,
source_type_filter: Optional[List[str]] = None, source_type_filter: Optional[List[str]] = None,
expand_links: bool = True, expand_links: bool = True,
neighbor_expand: int = 0, neighbor_expand: int = 0,
) -> List[Source]: ) -> RetrievalPipelineConfig:
qvec = Vector(embed_query(embedder, query))
vector_k = max(k * 4, k + 16) vector_k = max(k * 4, k + 16)
return RetrievalPipelineConfig(
k=k,
vector_k=vector_k,
expand_links=expand_links,
neighbor_expand=neighbor_expand,
enable_global_search=True,
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,
source_type_filter=source_type_filter,
)
def run_global_child_vector_search(
pg_url: str,
qvec: Vector,
config: RetrievalPipelineConfig,
) -> List[Retrieved]:
if not config.enable_global_search:
return []
where = ["doc_type = 'child'"] where = ["doc_type = 'child'"]
params: Dict[str, Any] = {"qvec": qvec, "vector_k": vector_k} params: Dict[str, Any] = {"qvec": qvec, "vector_k": max(1, int(config.vector_k))}
if chapter_index is not None: if config.chapter_index is not None:
where.append("chapter_index = %(chapter_index)s") where.append("chapter_index = %(chapter_index)s")
params["chapter_index"] = chapter_index params["chapter_index"] = config.chapter_index
if section_index is not None: if config.section_index is not None:
where.append("section_index = %(section_index)s") where.append("section_index = %(section_index)s")
params["section_index"] = section_index params["section_index"] = config.section_index
if subsection_index is not None: if config.subsection_index is not None:
where.append("subsection_index = %(subsection_index)s") where.append("subsection_index = %(subsection_index)s")
params["subsection_index"] = subsection_index params["subsection_index"] = config.subsection_index
if source_type_filter: if config.source_type_filter:
where.append("source_type = ANY(%(source_type_filter)s)") where.append("source_type = ANY(%(source_type_filter)s)")
params["source_type_filter"] = source_type_filter params["source_type_filter"] = config.source_type_filter
where_sql = " AND ".join(where) where_sql = " AND ".join(where)
sql = f""" sql = f"""
SELECT SELECT
uid, doc_type, uid, doc_type,
...@@ -353,192 +399,362 @@ def retrieve( ...@@ -353,192 +399,362 @@ def retrieve(
ORDER BY embedding <=> %(qvec)s ORDER BY embedding <=> %(qvec)s
LIMIT %(vector_k)s; LIMIT %(vector_k)s;
""" """
with psycopg.connect(pg_url, row_factory=dict_row) as conn:
register_vector(conn)
with conn.cursor() as cur:
cur.execute(sql, params)
rows = cur.fetchall()
return [_row_to_retrieved(row) for row in rows]
children: List[Retrieved] = []
neighbors: List[Retrieved] = [] def select_top_k_unique(values: List[Retrieved], wanted_k: int) -> List[Retrieved]:
chapter_docs: List[Retrieved] = [] selected: List[Retrieved] = []
subsections: List[Retrieved] = [] seen: set[str] = set()
sections_docs: List[Retrieved] = [] for item in values:
children_direct: List[Retrieved] = [] if item.uid in seen:
children_expanded: List[Retrieved] = [] continue
selected.append(item)
def _top_k_unique(values: List[Retrieved], wanted_k: int) -> List[Retrieved]: seen.add(item.uid)
selected: List[Retrieved] = [] if len(selected) >= wanted_k:
seen: set[str] = set() break
for item in values: return selected
if item.uid in seen:
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]] = {}
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:
continue
if level == "subsection":
if subsection_index is None:
continue continue
selected.append(item) key = (int(chapter_index), int(section_index), int(subsection_index))
seen.add(item.uid) else:
if len(selected) >= wanted_k: key = (int(chapter_index), int(section_index))
break count, score_sum = counts.get(key, (0, 0.0))
return selected counts[key] = (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
avg_score = score_sum / count if count else 0.0
return (count, avg_score, tuple([-part for part in key]))
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",
)
def fetch_scope_context_docs(
pg_url: str,
scope: DominantScope,
) -> Dict[str, List[Retrieved]]:
chapter_docs: List[Retrieved] = []
section_docs: List[Retrieved] = []
subsection_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:
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.path, d.markdown,
1.0 AS score
FROM docs d
WHERE d.doc_type = 'subsection'
AND d.chapter_index = %(cpt)s
AND d.section_index = %(sec)s
AND d.subsection_index = %(sub)s
""",
{"cpt": scope.chapter_index, "sec": scope.section_index,
"sub": scope.subsection_index},
)
subsection_docs = [_row_to_retrieved(
row, source_type="subsection") 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.path, d.markdown,
1.0 AS score
FROM docs d
WHERE d.doc_type = 'section'
AND d.chapter_index = %(cpt)s
AND d.section_index = %(sec)s
""",
{"cpt": scope.chapter_index, "sec": scope.section_index},
)
section_docs = [_row_to_retrieved(
row, source_type="section") 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.path, d.markdown,
1.0 AS score
FROM docs d
WHERE d.doc_type = 'chapter'
AND d.chapter_index = %(cpt)s
""",
{"cpt": scope.chapter_index},
)
chapter_docs = [_row_to_retrieved(
row, source_type="chapter") for row in cur.fetchall()]
return {
"chapters": chapter_docs,
"sections": section_docs,
"subsections": subsection_docs,
}
def run_scoped_child_vector_search(
pg_url: str,
qvec: Vector,
scope: DominantScope,
fill_k: int,
) -> List[Retrieved]:
where = [
"d.doc_type = 'child'",
"d.chapter_index = %(cpt)s",
"d.section_index = %(sec)s",
]
params: Dict[str, Any] = {
"qvec": qvec,
"cpt": scope.chapter_index,
"sec": scope.section_index,
"fill_k": max(1, int(fill_k)),
}
if scope.level == "subsection" and scope.subsection_index is not None:
where.append("d.subsection_index = %(sub)s")
params["sub"] = scope.subsection_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.path, d.markdown,
1 - (d.embedding <=> %(qvec)s) AS score
FROM docs d
WHERE {where_sql}
ORDER BY d.embedding <=> %(qvec)s
LIMIT %(fill_k)s;
"""
with psycopg.connect(pg_url, row_factory=dict_row) as conn: with psycopg.connect(pg_url, row_factory=dict_row) as conn:
register_vector(conn) register_vector(conn)
with conn.cursor() as cur: with conn.cursor() as cur:
cur.execute(sql, params) cur.execute(sql, params)
rows = cur.fetchall() rows = cur.fetchall()
children = [_row_to_retrieved(row) for row in rows] return [_row_to_retrieved(row) for row in rows]
children_direct = _top_k_unique(children, k)
if expand_links and children_direct:
cpt_sec_sub_counts: Dict[tuple[int, int, int], int] = {}
for child in children_direct:
cpt_idx = child.metadata.get("chapter_index")
sec_idx = child.metadata.get("section_index")
sub_idx = child.metadata.get("subsection_index")
if cpt_idx is not None and sec_idx is not None and sub_idx is not None:
key = (int(cpt_idx), int(sec_idx), int(sub_idx))
cpt_sec_sub_counts[key] = cpt_sec_sub_counts.get(key, 0) + 1
if cpt_sec_sub_counts:
most_common_cpt_sec_sub = max(
cpt_sec_sub_counts.items(), key=lambda x: x[1])[0]
most_common_cpt, most_common_sec, most_common_sub = most_common_cpt_sec_sub
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.path, d.markdown,
1.0 AS score
FROM docs d
WHERE d.doc_type = 'subsection'
AND d.chapter_index = %(cpt)s
AND d.section_index = %(sec)s
AND d.subsection_index = %(sub)s
""",
{"cpt": most_common_cpt, "sec": most_common_sec, "sub": most_common_sub},
)
subsections = [_row_to_retrieved(
row, source_type="subsection") 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.path, d.markdown,
1.0 AS score
FROM docs d
WHERE d.doc_type = ANY(%(sec_doc_types)s)
AND d.chapter_index = %(cpt)s
AND d.section_index = %(sec)s
""",
{"cpt": most_common_cpt, "sec": most_common_sec, "sec_doc_types": [
"section"]},
)
sections_docs = [_row_to_retrieved(
row, source_type="section") 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.path, d.markdown,
1.0 AS score
FROM docs d
WHERE d.doc_type = 'chapter'
AND d.chapter_index = %(cpt)s
""",
{"cpt": most_common_cpt},
)
chapter_docs = [_row_to_retrieved(
row, source_type="chapter") 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.path, d.markdown,
1 - (d.embedding <=> %(qvec)s) AS score
FROM docs d
WHERE d.doc_type = 'child'
AND chapter_index = %(cpt)s
AND d.section_index = %(sec)s
AND d.subsection_index = %(sub)s
ORDER BY d.embedding <=> %(qvec)s
LIMIT %(fill_k)s;
""",
{"qvec": qvec, "cpt": most_common_cpt, "sec": most_common_sec,
"sub": most_common_sub, "fill_k": max(k, 5)},
)
additional_children = [_row_to_retrieved(
row) for row in cur.fetchall()] def expand_neighbor_children(
existing_uids = {child.uid for child in children_direct} pg_url: str,
new_children = [ children: List[Retrieved],
child for child in additional_children if child.uid not in existing_uids] neighbor_expand: int,
needed = max(0, k - len(children_direct)) ) -> List[Retrieved]:
if needed: if neighbor_expand <= 0:
fill_children = new_children[:needed] return []
children_expanded.extend(fill_children)
children = [*children_direct, *fill_children] wanted: set[Tuple[int, int, int, int]] = set()
else: for child in children:
children = list(children_direct) cpti = child.metadata.get("chapter_index")
si = child.metadata.get("section_index")
if neighbor_expand and neighbor_expand > 0: ssi = child.metadata.get("subsection_index")
wanted: set[tuple[int, int, int, int]] = set() ci = child.metadata.get("child_index")
for child in children: if cpti is None or si is None or ssi is None or ci is None:
cpti = child.metadata.get("chapter_index") continue
si = child.metadata.get("section_index") for dx in range(-neighbor_expand, neighbor_expand + 1):
ssi = child.metadata.get("subsection_index") if dx == 0:
ci = child.metadata.get("child_index") continue
if cpti is None or si is None or ssi is None or ci is None: wanted.add((int(cpti), int(si), int(ssi), int(ci) + dx))
continue
for dx in range(-neighbor_expand, neighbor_expand + 1): if not wanted:
if dx == 0: return []
continue
wanted.add((int(cpti), int(si), int(ssi), int(ci) + dx)) quadruple = sorted(wanted)
cpt_arr = [a for (a, b, c, cidx) in quadruple]
if wanted: sec_arr = [b for (a, b, c, cidx) in quadruple]
quadruple = sorted(wanted) sub_arr = [c for (a, b, c, cidx) in quadruple]
cpt_arr = [a for (a, b, c, cidx) in quadruple] child_arr = [cidx 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] with psycopg.connect(pg_url, row_factory=dict_row) as conn:
child_arr = [cidx for (a, b, c, cidx) in quadruple] with conn.cursor() as cur:
cur.execute(
cur.execute( """
""" SELECT
SELECT d.uid, d.doc_type,
d.uid, d.doc_type, d.chapter_index, d.section_index, d.subsection_index, d.child_index,
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_title, d.section_title, d.subsection_title, d.title, d.source_type, d.path, d.markdown,
d.path, d.markdown, 0 AS score
0 AS score FROM docs d
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)
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
ON d.chapter_index = u.cpt AND d.section_index = u.sec AND d.subsection_index = u.sub AND d.child_index = u.child WHERE d.doc_type = '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, "child_arr": child_arr},
"child_arr": child_arr}, )
) rows = cur.fetchall()
neighbors = [_row_to_retrieved( return [_row_to_retrieved(row) for row in rows]
row) for row in cur.fetchall()]
children = [*children_direct, *children_expanded] def merge_retrieval_groups(
child_uids = {child.uid for child in children} groups: Dict[str, List[Retrieved]],
neighbors = [ k: int,
neighbor for neighbor in neighbors if neighbor.uid not in child_uids] ) -> Dict[str, List[Retrieved]]:
retrivla_dict = { if not groups:
"children_direct": children_direct, return {}
"children_expanded": children_expanded,
"chapters": chapter_docs, result: Dict[str, List[Retrieved]] = {name: [] for name in groups}
"subsections": subsections, first_group_for_uid: Dict[str, str] = {}
"sections": sections_docs, best_for_uid: Dict[str, Retrieved] = {}
"neighbors": neighbors,
for group_name, items in groups.items():
for item in items:
existing = best_for_uid.get(item.uid)
if existing is None:
first_group_for_uid[item.uid] = group_name
best_for_uid[item.uid] = item
continue
if item.score > existing.score:
best_for_uid[item.uid] = item
for uid, item in best_for_uid.items():
group_name = first_group_for_uid[uid]
result[group_name].append(item)
child_groups = [name for name in ("children_direct", "children_expanded") if name in result]
if child_groups:
all_children = [item for group_name in child_groups for item in result[group_name]]
all_children = sorted(all_children, key=lambda item: item.score, reverse=True)[:max(0, k)]
allowed_uids = {item.uid for item in all_children}
for group_name in child_groups:
result[group_name] = [
item for item in result[group_name] if item.uid in allowed_uids]
return result
def run_child_retrieval_pipeline(
pg_url: str,
embedder: EmbeddingLike,
query: str,
config: RetrievalPipelineConfig,
) -> Dict[str, List[Retrieved]]:
qvec = Vector(embed_query(embedder, query))
groups: Dict[str, List[Retrieved]] = {
"children_direct": [],
"children_expanded": [],
"chapters": [],
"subsections": [],
"sections": [],
"neighbors": [],
} }
sources = _retrivla_to_sources(retrivla_dict)
global_children = run_global_child_vector_search(pg_url, qvec, config)
children_direct = select_top_k_unique(global_children, config.k)
groups["children_direct"] = children_direct
final_children = list(children_direct)
dominant_scope: Optional[DominantScope] = None
if config.expand_links and config.enable_dominant_scope:
dominant_scope = select_dominant_scope(children_direct, config.dominance_level)
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", [])
if dominant_scope and config.expand_links and config.enable_scoped_child_search:
additional_children = run_scoped_child_vector_search(
pg_url=pg_url,
qvec=qvec,
scope=dominant_scope,
fill_k=config.scope_fill_k,
)
existing_uids = {child.uid for child in children_direct}
needed = max(0, config.k - len(children_direct))
if needed > 0:
fill_children = [child for child in additional_children if child.uid not in existing_uids][:needed]
groups["children_expanded"] = fill_children
final_children = [*children_direct, *fill_children]
if config.neighbor_expand > 0:
neighbors = expand_neighbor_children(
pg_url=pg_url,
children=final_children,
neighbor_expand=config.neighbor_expand,
)
groups["neighbors"] = neighbors
return merge_retrieval_groups(groups, config.k)
def retrieve(
pg_url: str,
embedder: EmbeddingLike,
query: str,
k: int = 4,
chapter_index: Optional[int] = None,
section_index: Optional[int] = None,
subsection_index: Optional[int] = None,
source_type_filter: Optional[List[str]] = None,
expand_links: bool = True,
neighbor_expand: int = 0,
) -> List[Source]:
config = build_default_pipeline_config(
k=k,
chapter_index=chapter_index,
section_index=section_index,
subsection_index=subsection_index,
source_type_filter=source_type_filter,
expand_links=expand_links,
neighbor_expand=neighbor_expand,
)
retrieval_groups = run_child_retrieval_pipeline(
pg_url=pg_url,
embedder=embedder,
query=query,
config=config,
)
sources = _retrivla_to_sources(retrieval_groups)
return sorted(sources, key=lambda source: source.score, reverse=True) return sorted(sources, key=lambda source: source.score, reverse=True)
# -------------------------------------------------------------------------------------------------------------------- # --------------------------------------------------------------------------------------------------------------------
...@@ -602,7 +818,7 @@ def load_children_for_subsections( ...@@ -602,7 +818,7 @@ def load_children_for_subsections(
def merge_sources(primary: List[Source], additional: List[Source]) -> List[Source]: def merge_sources(primary: List[Source], additional: List[Source]) -> List[Source]:
merged: List[Source] = [] merged: List[Source] = []
seen: set[Tuple[str, str, str, str, str, str]] = set() index_by_key: Dict[Tuple[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]:
return ( return (
...@@ -616,10 +832,15 @@ def merge_sources(primary: List[Source], additional: List[Source]) -> List[Sourc ...@@ -616,10 +832,15 @@ def merge_sources(primary: List[Source], additional: List[Source]) -> List[Sourc
for source in [*primary, *additional]: for source in [*primary, *additional]:
key = source_key(source) key = source_key(source)
if key in seen: existing_index = index_by_key.get(key)
if existing_index is None:
index_by_key[key] = len(merged)
merged.append(source)
continue continue
seen.add(key)
merged.append(source) # If a source appears both as vector hit and task-linked child, prefer task_childs.
if source.retrieved_as == "task_childs":
merged[existing_index] = source
return merged return merged
......
math-tutor/backend/test/vector_store_pipeline_unit_test.py 0 → 100644
View file @ d93bc3e6
from __future__ import annotations
import unittest
from app.deterministic_services.vector_store import (
Retrieved,
Source,
SourceID,
build_default_pipeline_config,
expand_neighbor_children,
merge_retrieval_groups,
merge_sources,
select_dominant_scope,
)
def _mk_retrieved(
uid: str,
score: float,
chapter_index: int,
section_index: int,
subsection_index: int | None,
child_index: int = 1,
) -> Retrieved:
return Retrieved(
uid=uid,
doc_type="child",
score=score,
metadata={
"chapter_index": chapter_index,
"section_index": section_index,
"subsection_index": subsection_index,
"child_index": child_index,
"chapter_title": "C",
"section_title": "S",
"subsection_title": "SS",
"title": uid,
"source_type": "child",
"path": "",
"doc_type": "child",
"uid": uid,
},
markdown=f"md-{uid}",
)
class VectorStorePipelineUnitTest(unittest.TestCase):
def test_merge_sources_prefers_task_childs_on_duplicate(self) -> None:
child_direct = Source(
source_id=SourceID(
chapter_title="Kapitel",
section_title="Section",
subsection_title="Subsection",
title="Child A",
doc_type="child",
),
retrieved_as="children_direct",
source_type="child",
score=0.7,
markdown="same-md",
)
task_child = Source(
source_id=SourceID(
chapter_title="Kapitel",
section_title="Section",
subsection_title="Subsection",
title="Child A",
doc_type="child",
),
retrieved_as="task_childs",
source_type="child",
score=1.0,
markdown="same-md",
)
merged = merge_sources([child_direct], [task_child])
self.assertEqual(len(merged), 1)
self.assertEqual(merged[0].retrieved_as, "task_childs")
def test_build_default_pipeline_config(self) -> None:
cfg = build_default_pipeline_config(k=4, expand_links=True, neighbor_expand=2)
self.assertEqual(cfg.vector_k, 20)
self.assertEqual(cfg.scope_fill_k, 5)
self.assertTrue(cfg.enable_global_search)
self.assertTrue(cfg.enable_dominant_scope)
self.assertTrue(cfg.enable_scoped_child_search)
self.assertTrue(cfg.enable_context_docs)
def test_select_dominant_scope_prefers_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")
self.assertIsNotNone(scope)
assert scope is not None
self.assertEqual((scope.chapter_index, scope.section_index, scope.subsection_index), (1, 1, 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),
]
scope = select_dominant_scope(children, level="subsection")
self.assertIsNotNone(scope)
assert scope is not None
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),
]
scope = select_dominant_scope(children, level="subsection")
self.assertIsNotNone(scope)
assert scope is not None
self.assertEqual((scope.chapter_index, scope.section_index, scope.subsection_index), (1, 2, 3))
def test_merge_retrieval_groups_dedup_max_score_and_trim_children(self) -> None:
groups = {
"children_direct": [
_mk_retrieved("u1", 0.5, 1, 1, 1),
_mk_retrieved("u2", 0.7, 1, 1, 1),
],
"children_expanded": [
_mk_retrieved("u1", 0.9, 1, 1, 1),
_mk_retrieved("u3", 0.6, 1, 1, 1),
],
"chapters": [],
"subsections": [],
"sections": [],
"neighbors": [_mk_retrieved("u2", 0.1, 1, 1, 1)],
}
merged = merge_retrieval_groups(groups, k=2)
self.assertEqual(sorted(item.uid for item in merged["children_direct"]), ["u1", "u2"])
score_u1 = [item.score for item in merged["children_direct"] if item.uid == "u1"][0]
self.assertEqual(score_u1, 0.9)
self.assertEqual(merged["children_expanded"], [])
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)]
result = expand_neighbor_children("postgresql://unused", children, neighbor_expand=0)
self.assertEqual(result, [])
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

Dies ist die Gitlab-Instanz des Transferportals der Hochschule für Technik Stuttgart. Hier geht es zurück zum Portal