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

aufräumen des vector stores

parent 9a38c477
Show 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
import hashlib
from dataclasses import dataclass
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
from psycopg.rows import dict_row
......@@ -305,10 +305,33 @@ def _row_to_retrieved(row: Dict[str, Any], source_type: Optional[str] = None) ->
)
def retrieve(
pg_url: str,
embedder: EmbeddingLike,
query: str,
@dataclass
class RetrievalPipelineConfig:
k: int
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,
chapter_index: Optional[int] = None,
section_index: Optional[int] = None,
......@@ -316,31 +339,54 @@ def retrieve(
source_type_filter: Optional[List[str]] = None,
expand_links: bool = True,
neighbor_expand: int = 0,
) -> List[Source]:
qvec = Vector(embed_query(embedder, query))
) -> RetrievalPipelineConfig:
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'"]
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")
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")
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")
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)")
params["source_type_filter"] = source_type_filter
params["source_type_filter"] = config.source_type_filter
where_sql = " AND ".join(where)
sql = f"""
SELECT
uid, doc_type,
......@@ -353,16 +399,15 @@ def retrieve(
ORDER BY embedding <=> %(qvec)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] = []
chapter_docs: List[Retrieved] = []
subsections: List[Retrieved] = []
sections_docs: List[Retrieved] = []
children_direct: List[Retrieved] = []
children_expanded: List[Retrieved] = []
def _top_k_unique(values: List[Retrieved], wanted_k: int) -> List[Retrieved]:
def select_top_k_unique(values: List[Retrieved], wanted_k: int) -> List[Retrieved]:
selected: List[Retrieved] = []
seen: set[str] = set()
for item in values:
......@@ -374,29 +419,65 @@ def retrieve(
break
return selected
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()
children = [_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] = {}
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:
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
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
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))
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
......@@ -411,12 +492,12 @@ def retrieve(
AND d.section_index = %(sec)s
AND d.subsection_index = %(sub)s
""",
{"cpt": most_common_cpt, "sec": most_common_sec, "sub": most_common_sub},
{"cpt": scope.chapter_index, "sec": scope.section_index,
"sub": scope.subsection_index},
)
subsections = [_row_to_retrieved(
subsection_docs = [_row_to_retrieved(
row, source_type="subsection") for row in cur.fetchall()]
cur.execute(
"""
SELECT
......@@ -426,14 +507,13 @@ def retrieve(
d.path, d.markdown,
1.0 AS score
FROM docs d
WHERE d.doc_type = ANY(%(sec_doc_types)s)
WHERE d.doc_type = 'section'
AND d.chapter_index = %(cpt)s
AND d.section_index = %(sec)s
""",
{"cpt": most_common_cpt, "sec": most_common_sec, "sec_doc_types": [
"section"]},
{"cpt": scope.chapter_index, "sec": scope.section_index},
)
sections_docs = [_row_to_retrieved(
section_docs = [_row_to_retrieved(
row, source_type="section") for row in cur.fetchall()]
cur.execute(
......@@ -448,13 +528,41 @@ def retrieve(
WHERE d.doc_type = 'chapter'
AND d.chapter_index = %(cpt)s
""",
{"cpt": most_common_cpt},
{"cpt": scope.chapter_index},
)
chapter_docs = [_row_to_retrieved(
row, source_type="chapter") for row in cur.fetchall()]
cur.execute(
"""
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,
......@@ -462,32 +570,27 @@ def retrieve(
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
WHERE {where_sql}
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)},
)
"""
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]
additional_children = [_row_to_retrieved(
row) for row in cur.fetchall()]
existing_uids = {child.uid for child in children_direct}
new_children = [
child for child in additional_children if child.uid not in existing_uids]
needed = max(0, k - len(children_direct))
if needed:
fill_children = new_children[:needed]
children_expanded.extend(fill_children)
children = [*children_direct, *fill_children]
else:
children = list(children_direct)
if neighbor_expand and neighbor_expand > 0:
wanted: set[tuple[int, int, int, int]] = set()
def expand_neighbor_children(
pg_url: str,
children: List[Retrieved],
neighbor_expand: int,
) -> List[Retrieved]:
if neighbor_expand <= 0:
return []
wanted: set[Tuple[int, int, int, int]] = set()
for child in children:
cpti = child.metadata.get("chapter_index")
si = child.metadata.get("section_index")
......@@ -500,13 +603,17 @@ def retrieve(
continue
wanted.add((int(cpti), int(si), int(ssi), int(ci) + dx))
if wanted:
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]
with psycopg.connect(pg_url, row_factory=dict_row) as conn:
with conn.cursor() as cur:
cur.execute(
"""
SELECT
......@@ -523,22 +630,131 @@ def retrieve(
{"cpt_arr": cpt_arr, "sec_arr": sec_arr, "sub_arr": sub_arr,
"child_arr": child_arr},
)
neighbors = [_row_to_retrieved(
row) for row in cur.fetchall()]
children = [*children_direct, *children_expanded]
child_uids = {child.uid for child in children}
neighbors = [
neighbor for neighbor in neighbors if neighbor.uid not in child_uids]
retrivla_dict = {
"children_direct": children_direct,
"children_expanded": children_expanded,
"chapters": chapter_docs,
"subsections": subsections,
"sections": sections_docs,
"neighbors": neighbors,
rows = cur.fetchall()
return [_row_to_retrieved(row) for row in rows]
def merge_retrieval_groups(
groups: Dict[str, List[Retrieved]],
k: int,
) -> Dict[str, List[Retrieved]]:
if not groups:
return {}
result: Dict[str, List[Retrieved]] = {name: [] for name in groups}
first_group_for_uid: Dict[str, str] = {}
best_for_uid: Dict[str, Retrieved] = {}
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)
# --------------------------------------------------------------------------------------------------------------------
......@@ -602,7 +818,7 @@ def load_children_for_subsections(
def merge_sources(primary: List[Source], additional: List[Source]) -> 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]:
return (
......@@ -616,10 +832,15 @@ def merge_sources(primary: List[Source], additional: List[Source]) -> List[Sourc
for source in [*primary, *additional]:
key = source_key(source)
if key in seen:
continue
seen.add(key)
existing_index = index_by_key.get(key)
if existing_index is None:
index_by_key[key] = len(merged)
merged.append(source)
continue
# 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
......
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