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

moving retrieval to retrieval store

parent 8c122e25
Hide whitespace changes
Inline Side-by-side
math-tutor/backend/app/deterministic_services/__init__.py
View file @ 0d61c1f7
# Package marker for services # Package marker for services
from app.deterministic_services.vector_store import Source, SourceID from app.deterministic_services.retrieval_store import Source, SourceID
__all__ = ["Source", "SourceID"] __all__ = ["Source", "SourceID"]
\ No newline at end of file
math-tutor/backend/app/deterministic_services/retrieval_store.py
View file @ 0d61c1f7
# Mask for Vector Store Retrieval. # Retrieval Store to fetch relevant documents based on query and context, using vector search and other heuristics.
# dont know why this is here?
from __future__ import annotations from __future__ import annotations
from typing import List from dataclasses import dataclass
from typing import Any, Dict, List, Literal, Optional, Tuple
import psycopg
from pgvector import Vector
from pgvector.psycopg import register_vector
from psycopg.rows import dict_row
from pydantic import BaseModel
from app.deterministic_services.vector_store import EmbeddingLike, embed_query
def _load_subsection_store():
from app.deterministic_services import vector_store_subsection
return vector_store_subsection
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
def to_dict(self) -> Dict[str, Any]:
return {
"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:
string_rep = f"{self.section_title}|{string_rep}"
if self.chapter_title:
string_rep = f"{self.chapter_title}|{string_rep}"
return f"[{string_rep}|{self.doc_type}]"
class Source(BaseModel):
source_id: SourceID
retrieved_as: str
source_type: str
score: float
markdown: str
def to_dict(self) -> Dict[str, Any]:
return {
"source_id": self.source_id.to_dict(),
"retrieved_as": self.retrieved_as,
"source_type": self.source_type,
"score": self.score,
"markdown": self.markdown,
}
def to_string(self) -> str:
return (
f"Source(source_id={self.source_id.to_string()},\n"
f" retrieved_as={self.retrieved_as},\n"
f" source_type={self.source_type},\n"
f" score={self.score},\n"
f" markdown={self.markdown}"
)
@dataclass
class Retrieved:
uid: str
doc_type: str
score: float
metadata: Dict[str, Any]
markdown: str
def to_dict(self) -> Dict[str, Any]:
return {
"uid": self.uid,
"doc_type": self.doc_type,
"score": self.score,
"metadata": self.metadata,
"markdown": self.markdown,
}
def _row_to_retrieved(row: Dict[str, Any], source_type: Optional[str] = None) -> Retrieved:
meta = {
"uid": row["uid"],
"doc_type": row["doc_type"],
"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"],
}
return Retrieved(
uid=row["uid"],
doc_type=row["doc_type"],
score=float(row["score"]),
metadata=meta,
markdown=row["markdown"],
)
@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
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
@dataclass
class DominantScope:
chapter_index: int
section_index: int
subsection_index: Optional[int]
subsubsection_index: Optional[int]
level: Literal["section", "subsection", "subsubsection"]
SubsectionRef = Tuple[int, int, int]
ParentRef = Tuple[int, int, int, int]
def build_default_pipeline_config(
k: int = 4,
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,
) -> 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,
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,
)
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": max(1, int(config.vector_k))}
if config.chapter_index is not None:
where.append("chapter_index = %(chapter_index)s")
params["chapter_index"] = config.chapter_index
if config.section_index is not None:
where.append("section_index = %(section_index)s")
params["section_index"] = config.section_index
if config.subsection_index is not None:
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
where_sql = " AND ".join(where)
sql = f"""
SELECT
uid, doc_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
WHERE {where_sql}
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]
def select_top_k_unique(values: List[Retrieved], wanted_k: int) -> List[Retrieved]:
selected: List[Retrieved] = []
seen: set[str] = set()
for item in values:
if item.uid in seen:
continue
selected.append(item)
seen.add(item.uid)
if len(selected) >= wanted_k:
break
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],
) -> Optional[DominantScope]:
if not children_direct:
return None
counts: Dict[Tuple[str, Tuple[int, ...]], Tuple[DominantScope, int, float]] = {}
for child in children_direct:
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
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[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 _scope_key(scope)]))
_, (winner_scope, _, _) = max(counts.items(), key=_rank)
return winner_scope
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] = []
subsubsection_docs: List[Retrieved] = []
with psycopg.connect(pg_url, row_factory=dict_row) as conn:
with conn.cursor() as cur:
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.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 = '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()]
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.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 = '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()]
from app.deterministic_services import vector_store cur.execute(
from app.deterministic_services.vector_store import EmbeddingLike, Source """
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 = '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,
"subsubsections": subsubsection_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 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.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
WHERE {where_sql}
ORDER BY d.embedding <=> %(qvec)s
LIMIT %(fill_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]
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, 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 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(sssi), int(ci) + dx))
if not wanted:
return []
quintuple = sorted(wanted)
cpt_arr = [a for (a, _, _, _, _) in quintuple]
sec_arr = [b for (_, b, _, _, _) in quintuple]
sub_arr = [c for (_, _, c, _, _) in quintuple]
subsub_arr = [d for (_, _, _, d, _) in quintuple]
child_arr = [cidx for (_, _, _, _, cidx) in quintuple]
with psycopg.connect(pg_url, row_factory=dict_row) as conn:
with conn.cursor() as cur:
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,
0 AS score
FROM docs d
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,
"subsub_arr": subsub_arr,
"child_arr": child_arr,
},
)
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": [],
"subsubsections": [],
"sections": [],
"neighbors": [],
}
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)
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(
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 _retrivla_to_sources(retrievd: Dict[str, List[Retrieved]]) -> List[Source]:
sources = []
for name, retrieved_group in retrievd.items():
for retrieval in retrieved_group:
sources.append(
Source(
source_id=SourceID(
chapter_title=retrieval.metadata.get("chapter_title"),
section_title=retrieval.metadata.get("section_title"),
subsection_title=retrieval.metadata.get("subsection_title"),
subsubsection_title=retrieval.metadata.get("subsubsection_title"),
title=retrieval.metadata.get("title"),
doc_type=retrieval.doc_type,
),
retrieved_as=name,
source_type=retrieval.metadata.get("source_type") or "unknown",
score=retrieval.score,
markdown=retrieval.markdown,
)
)
return sources
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, str], int] = {}
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 "",
)
for source in [*primary, *additional]:
key = source_key(source)
existing_index = index_by_key.get(key)
if existing_index is None:
index_by_key[key] = len(merged)
merged.append(source)
continue
if source.retrieved_as == "task_childs":
merged[existing_index] = source
return merged
def retrieve( def retrieve(
pg_url: str, pg_url: str,
embedder: EmbeddingLike, embedder: EmbeddingLike,
query: str, query: str,
k: int = 4, k: int = 4,
chapter_index: int | None = None, chapter_index: Optional[int] = None,
section_index: int | None = None, section_index: Optional[int] = None,
subsection_index: int | None = None, subsection_index: Optional[int] = None,
subsubsection_index: int | None = None, subsubsection_index: Optional[int] = None,
source_type_filter: list[str] | None = 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]: ) -> List[Source]:
config_obj = build_default_pipeline_config(
return vector_store.retrieve(
pg_url=pg_url,
embedder=embedder,
query=query,
k=k, k=k,
chapter_index=chapter_index, chapter_index=chapter_index,
section_index=section_index, section_index=section_index,
...@@ -39,27 +700,128 @@ def retrieve( ...@@ -39,27 +700,128 @@ def retrieve(
expand_links=expand_links, expand_links=expand_links,
neighbor_expand=neighbor_expand, neighbor_expand=neighbor_expand,
) )
retrieval_groups = run_child_retrieval_pipeline(
pg_url=pg_url,
embedder=embedder,
query=query,
config=config_obj,
)
sources = _retrivla_to_sources(retrieval_groups)
return sorted(sources, key=lambda source: source.score, reverse=True)
def _normalize_parent_refs(
parent_refs: Optional[List[ParentRef]],
) -> List[ParentRef]:
if not parent_refs:
return []
normalized: set[ParentRef] = set()
for chap, sec, sub, subsub in parent_refs:
chap_i, sec_i, sub_i, subsub_i = int(chap), int(sec), int(sub), int(subsub)
if sub_i <= 0:
normalized.add((chap_i, sec_i, 0, 0))
elif subsub_i <= 0:
normalized.add((chap_i, sec_i, sub_i, 0))
else:
normalized.add((chap_i, sec_i, sub_i, subsub_i))
return sorted(normalized)
def _parent_ref_scope(ref: ParentRef) -> DominantScope:
chap, sec, sub, subsub = ref
if sub <= 0:
return DominantScope(chapter_index=chap, section_index=sec, subsection_index=None, subsubsection_index=None, level="section")
if subsub <= 0:
return DominantScope(chapter_index=chap, section_index=sec, subsection_index=sub, subsubsection_index=None, level="subsection")
return DominantScope(chapter_index=chap, section_index=sec, subsection_index=sub, subsubsection_index=subsub, level="subsubsection")
def _normalize_subsection_refs(
subsection_refs: Optional[List[SubsectionRef]],
) -> List[SubsectionRef]:
if not subsection_refs:
return []
normalized = {(int(chap), int(sec), int(sub)) for chap, sec, sub in subsection_refs}
return sorted(normalized)
def load_children_for_subsections(
pg_url: str,
subsection_refs: Optional[List[SubsectionRef]],
) -> List[Source]:
refs = _normalize_subsection_refs(subsection_refs)
if not refs:
return []
chap_arr = [chap for chap, _, _ in refs]
sec_arr = [sec for _, sec, _ in refs]
sub_arr = [sub for _, _, sub in refs]
sql = """
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
JOIN unnest(%(chap_arr)s::int[], %(sec_arr)s::int[], %(sub_arr)s::int[]) AS u(chap, sec, sub)
ON d.chapter_index = u.chap AND d.section_index = u.sec AND d.subsection_index = u.sub
WHERE d.doc_type = 'child'
ORDER BY d.chapter_index, d.section_index, d.subsection_index, d.child_index
"""
with psycopg.connect(pg_url, row_factory=dict_row) as conn:
with conn.cursor() as cur:
cur.execute(
sql,
{
"chap_arr": chap_arr,
"sec_arr": sec_arr,
"sub_arr": sub_arr,
},
)
rows = cur.fetchall()
children = [_row_to_retrieved(row) for row in rows]
return _retrivla_to_sources({"task_childs": children})
def retrieve_with_subsections( def retrieve_with_subsections(
pg_url: str, pg_url: str,
embedder: EmbeddingLike, embedder: EmbeddingLike,
query: str, query: str,
subsection_refs: Optional[List[SubsectionRef]] = None,
k: int = 4, k: int = 4,
chapter_index: int | None = None, chapter_index: Optional[int] = None,
section_index: int | None = None, section_index: Optional[int] = None,
subsection_index: int | None = None, subsection_index: Optional[int] = None,
subsubsection_index: int | None = None, subsubsection_index: Optional[int] = None,
source_type_filter: list[str] | None = 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]: ) -> List[Source]:
if not subsection_refs:
# ponytail: no separate default-path behavior without subsection refs; reuse retrieve until proven otherwise.
return retrieve(
pg_url=pg_url,
embedder=embedder,
query=query,
k=k,
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,
)
return vector_store.retrieve_with_subsections( vector_k = max(k * 4, k + 16)
vector_sources = retrieve(
pg_url=pg_url, pg_url=pg_url,
embedder=embedder, embedder=embedder,
query=query, query=query,
k=k, k=vector_k,
chapter_index=chapter_index, chapter_index=chapter_index,
section_index=section_index, section_index=section_index,
subsection_index=subsection_index, subsection_index=subsection_index,
...@@ -68,13 +830,166 @@ def retrieve_with_subsections( ...@@ -68,13 +830,166 @@ def retrieve_with_subsections(
expand_links=expand_links, expand_links=expand_links,
neighbor_expand=neighbor_expand, neighbor_expand=neighbor_expand,
) )
vector_sources = vector_sources[:k]
subsection_children = load_children_for_subsections(
pg_url=pg_url,
subsection_refs=subsection_refs,
)
return merge_sources(vector_sources, subsection_children)
def retrieve_for_subsections(
pg_url: str,
subsection_refs: Optional[List[SubsectionRef]] = None,
) -> List[Source]:
return load_children_for_subsections(
pg_url=pg_url,
subsection_refs=subsection_refs,
)
def load_sources_for_parent_refs(
pg_url: str,
parent_refs: Optional[List[ParentRef]],
) -> List[Source]:
refs = _normalize_parent_refs(parent_refs)
if not refs:
return []
groups: Dict[str, List[Retrieved]] = {
"sections": [],
"subsections": [],
"subsubsections": [],
"task_childs": [],
}
with psycopg.connect(pg_url, row_factory=dict_row) as conn:
with conn.cursor() as cur:
for ref in refs:
scope = _parent_ref_scope(ref)
if scope.level == "section":
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 = 'section'
AND d.chapter_index = %(cpt)s
AND d.section_index = %(sec)s
""",
{"cpt": scope.chapter_index, "sec": scope.section_index},
)
groups["sections"].extend([_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.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 = 'child'
AND d.chapter_index = %(cpt)s
AND d.section_index = %(sec)s
AND d.subsection_index = 0
AND d.subsubsection_index = 0
ORDER BY d.child_index
""",
{"cpt": scope.chapter_index, "sec": scope.section_index},
)
groups["task_childs"].extend([_row_to_retrieved(row) for row in cur.fetchall()])
continue
if scope.level == "subsection":
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 = '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},
)
groups["subsections"].extend([_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.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 = 'child'
AND d.chapter_index = %(cpt)s
AND d.section_index = %(sec)s
AND d.subsection_index = %(sub)s
AND d.subsubsection_index = 0
ORDER BY d.child_index
""",
{"cpt": scope.chapter_index, "sec": scope.section_index, "sub": scope.subsection_index},
)
groups["task_childs"].extend([_row_to_retrieved(row) for row in cur.fetchall()])
continue
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},
)
groups["subsubsections"].extend([_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.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 = 'child'
AND d.chapter_index = %(cpt)s
AND d.section_index = %(sec)s
AND d.subsection_index = %(sub)s
AND d.subsubsection_index = %(subsub)s
ORDER BY d.child_index
""",
{"cpt": scope.chapter_index, "sec": scope.section_index, "sub": scope.subsection_index, "subsub": scope.subsubsection_index},
)
groups["task_childs"].extend([_row_to_retrieved(row) for row in cur.fetchall()])
return merge_sources([], _retrivla_to_sources(groups))
def retrieve_for_parent_refs( def retrieve_for_parent_refs(
pg_url: str, pg_url: str,
parent_refs: list[vector_store.ParentRef] | None = None, parent_refs: Optional[List[ParentRef]] = None,
) -> List[Source]: ) -> List[Source]:
return vector_store.load_sources_for_parent_refs( return load_sources_for_parent_refs(
pg_url=pg_url, pg_url=pg_url,
parent_refs=parent_refs, parent_refs=parent_refs,
) )
math-tutor/backend/app/deterministic_services/vector_store.py
View file @ 0d61c1f7
# Vector Store implementation using PostgreSQL with pgvector extension. # Vector Store for storing document embeddings and metadata, and performing similarity search.
from __future__ import annotations 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, Literal, Optional, Protocol, Tuple from typing import Any, Dict, List, Optional, Protocol, Tuple
import psycopg import psycopg
from psycopg.rows import dict_row
from pgvector import Vector from pgvector import Vector
from pgvector.psycopg import register_vector from pgvector.psycopg import register_vector
from psycopg.rows import dict_row
import yaml import yaml
from pydantic import BaseModel
import app.config import app.config
embedding_dim = app.config.get_embedding_settings().target_dim embedding_dim = app.config.get_embedding_settings().target_dim
# -------------------------------------------------------------------------------------------------------------------- # --------------------------------------------------------------------------------------------------------------------
# Reading of Documents # Reading the Documents
# -------------------------------------------------------------------------------------------------------------------- # --------------------------------------------------------------------------------------------------------------------
...@@ -126,7 +125,7 @@ def load_docs(base_dir: Path) -> List[DocRecord]: ...@@ -126,7 +125,7 @@ def load_docs(base_dir: Path) -> List[DocRecord]:
return docs return docs
# -------------------------------------------------------------------------------------------------------------------- # --------------------------------------------------------------------------------------------------------------------
# Init of the database # Init the Database
# -------------------------------------------------------------------------------------------------------------------- # --------------------------------------------------------------------------------------------------------------------
...@@ -175,7 +174,7 @@ def init_db(pg_url: str) -> None: ...@@ -175,7 +174,7 @@ def init_db(pg_url: str) -> None:
conn.commit() conn.commit()
# -------------------------------------------------------------------------------------------------------------------- # --------------------------------------------------------------------------------------------------------------------
# Insert and update of documents and embeddings # Insertion of Documents and Embeddings
# -------------------------------------------------------------------------------------------------------------------- # --------------------------------------------------------------------------------------------------------------------
...@@ -276,873 +275,7 @@ def embed_query(embedder: EmbeddingLike, text: str) -> List[float]: ...@@ -276,873 +275,7 @@ def embed_query(embedder: EmbeddingLike, text: str) -> List[float]:
return embedder.embed_query(text) return embedder.embed_query(text)
# -------------------------------------------------------------------------------------------------------------------- # --------------------------------------------------------------------------------------------------------------------
# Retrieval of Documents # Listing for Filtering
# --------------------------------------------------------------------------------------------------------------------
@dataclass
class Retrieved:
uid: str
doc_type: str
score: float
metadata: Dict[str, Any]
markdown: str
def to_dict(self) -> Dict[str, Any]:
return {
"uid": self.uid,
"doc_type": self.doc_type,
"score": self.score,
"metadata": self.metadata,
"markdown": self.markdown,
}
def _row_to_retrieved(row: Dict[str, Any], source_type: Optional[str] = None) -> Retrieved:
meta = {
"uid": row["uid"],
"doc_type": row["doc_type"],
"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"],
}
return Retrieved(
uid=row["uid"],
doc_type=row["doc_type"],
score=float(row["score"]),
metadata=meta,
markdown=row["markdown"],
)
@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
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
@dataclass
class DominantScope:
chapter_index: int
section_index: int
subsection_index: Optional[int]
subsubsection_index: Optional[int]
level: Literal["section", "subsection", "subsubsection"]
def build_default_pipeline_config(
k: int = 4,
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,
) -> 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,
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,
)
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": max(1, int(config.vector_k))}
if config.chapter_index is not None:
where.append("chapter_index = %(chapter_index)s")
params["chapter_index"] = config.chapter_index
if config.section_index is not None:
where.append("section_index = %(section_index)s")
params["section_index"] = config.section_index
if config.subsection_index is not None:
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
where_sql = " AND ".join(where)
sql = f"""
SELECT
uid, doc_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
WHERE {where_sql}
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]
def select_top_k_unique(values: List[Retrieved], wanted_k: int) -> List[Retrieved]:
selected: List[Retrieved] = []
seen: set[str] = set()
for item in values:
if item.uid in seen:
continue
selected.append(item)
seen.add(item.uid)
if len(selected) >= wanted_k:
break
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],
) -> Optional[DominantScope]:
if not children_direct:
return None
counts: Dict[Tuple[str, Tuple[int, ...]], Tuple[DominantScope, int, float]] = {}
for child in children_direct:
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
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[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 _scope_key(scope)]))
_, (winner_scope, _, _) = max(counts.items(), key=_rank)
return winner_scope
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] = []
subsubsection_docs: List[Retrieved] = []
with psycopg.connect(pg_url, row_factory=dict_row) as conn:
with conn.cursor() as cur:
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.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 = '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()]
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.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 = '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.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 = '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,
"subsubsections": subsubsection_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 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.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
WHERE {where_sql}
ORDER BY d.embedding <=> %(qvec)s
LIMIT %(fill_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]
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, 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 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(sssi), int(ci) + dx))
if not wanted:
return []
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:
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,
0 AS score
FROM docs d
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, "subsub_arr": subsub_arr,
"child_arr": child_arr},
)
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": [],
"subsubsections": [],
"sections": [],
"neighbors": [],
}
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)
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(
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,
subsubsection_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,
subsubsection_index=subsubsection_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)
# --------------------------------------------------------------------------------------------------------------------
# Retrieval with parent refs
# --------------------------------------------------------------------------------------------------------------------
ParentRef = Tuple[int, int, int, int]
def _normalize_parent_refs(
parent_refs: Optional[List[ParentRef]],
) -> List[ParentRef]:
if not parent_refs:
return []
normalized: set[ParentRef] = set()
for chap, sec, sub, subsub in parent_refs:
chap_i, sec_i, sub_i, subsub_i = int(chap), int(sec), int(sub), int(subsub)
if sub_i <= 0:
normalized.add((chap_i, sec_i, 0, 0))
elif subsub_i <= 0:
normalized.add((chap_i, sec_i, sub_i, 0))
else:
normalized.add((chap_i, sec_i, sub_i, subsub_i))
return sorted(normalized)
def _parent_ref_scope(ref: ParentRef) -> DominantScope:
chap, sec, sub, subsub = ref
if sub <= 0:
return DominantScope(chapter_index=chap, section_index=sec, subsection_index=None, subsubsection_index=None, level="section")
if subsub <= 0:
return DominantScope(chapter_index=chap, section_index=sec, subsection_index=sub, subsubsection_index=None, level="subsection")
return DominantScope(chapter_index=chap, section_index=sec, subsection_index=sub, subsubsection_index=subsub, level="subsubsection")
def load_sources_for_parent_refs(
pg_url: str,
parent_refs: Optional[List[ParentRef]],
) -> List[Source]:
refs = _normalize_parent_refs(parent_refs)
if not refs:
return []
groups: Dict[str, List[Retrieved]] = {
"sections": [],
"subsections": [],
"subsubsections": [],
"task_childs": [],
}
with psycopg.connect(pg_url, row_factory=dict_row) as conn:
with conn.cursor() as cur:
for ref in refs:
scope = _parent_ref_scope(ref)
if scope.level == "section":
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 = 'section'
AND d.chapter_index = %(cpt)s
AND d.section_index = %(sec)s
""",
{"cpt": scope.chapter_index, "sec": scope.section_index},
)
groups["sections"].extend([_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.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 = 'child'
AND d.chapter_index = %(cpt)s
AND d.section_index = %(sec)s
AND d.subsection_index = 0
AND d.subsubsection_index = 0
ORDER BY d.child_index
""",
{"cpt": scope.chapter_index, "sec": scope.section_index},
)
groups["task_childs"].extend([_row_to_retrieved(row) for row in cur.fetchall()])
continue
if scope.level == "subsection":
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 = '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},
)
groups["subsections"].extend([_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.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 = 'child'
AND d.chapter_index = %(cpt)s
AND d.section_index = %(sec)s
AND d.subsection_index = %(sub)s
AND d.subsubsection_index = 0
ORDER BY d.child_index
""",
{"cpt": scope.chapter_index, "sec": scope.section_index, "sub": scope.subsection_index},
)
groups["task_childs"].extend([_row_to_retrieved(row) for row in cur.fetchall()])
continue
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},
)
groups["subsubsections"].extend([_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.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 = 'child'
AND d.chapter_index = %(cpt)s
AND d.section_index = %(sec)s
AND d.subsection_index = %(sub)s
AND d.subsubsection_index = %(subsub)s
ORDER BY d.child_index
""",
{"cpt": scope.chapter_index, "sec": scope.section_index, "sub": scope.subsection_index, "subsub": scope.subsubsection_index},
)
groups["task_childs"].extend([_row_to_retrieved(row) for row in cur.fetchall()])
return merge_sources([], _retrivla_to_sources(groups))
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, str], int] = {}
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 "",
)
for source in [*primary, *additional]:
key = source_key(source)
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
# --------------------------------------------------------------------------------------------------------------------
# transform retrieval results to sources
# --------------------------------------------------------------------------------------------------------------------
class Source(BaseModel):
source_id: SourceID
retrieved_as: str
source_type: str
score: float
markdown: str
def to_dict(self) -> Dict[str, Any]:
return {
"source_id": self.source_id.to_dict(),
"retrieved_as": self.retrieved_as,
"source_type": self.source_type,
"score": self.score,
"markdown": self.markdown,
}
def to_string(self) -> str:
return (
f"Source(source_id={self.source_id.to_string()},\n"
f" retrieved_as={self.retrieved_as},\n"
f" source_type={self.source_type},\n"
f" score={self.score},\n"
f" markdown={self.markdown}"
)
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
def to_dict(self) -> Dict[str, Any]:
return {
"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:
string_rep = f"{self.section_title}|{string_rep}"
if self.chapter_title:
string_rep = f"{self.chapter_title}|{string_rep}"
return f"[{string_rep}|{self.doc_type}]"
def _retrivla_to_sources(retrievd: Dict[str, List[Retrieved]]) -> List[Source]:
sources = []
for name, retrieved_grouep in retrievd.items():
for retrival in retrieved_grouep:
sources.append(Source(
source_id=SourceID(
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
),
retrieved_as=name,
source_type=retrival.metadata.get("source_type") or "unknown",
score=retrival.score,
markdown=retrival.markdown
))
return sources
# --------------------------------------------------------------------------------------------------------------------
# Lists for filtering
# -------------------------------------------------------------------------------------------------------------------- # --------------------------------------------------------------------------------------------------------------------
......
math-tutor/backend/test/vector_store_pipeline_unit_test.py
View file @ 0d61c1f7
...@@ -2,11 +2,13 @@ from __future__ import annotations ...@@ -2,11 +2,13 @@ from __future__ import annotations
import os import os
import unittest import unittest
from unittest.mock import patch
os.environ.setdefault("EMBEDDING_PROVIDER", "sentence-transformer") os.environ.setdefault("EMBEDDING_PROVIDER", "sentence-transformer")
os.environ.setdefault("EMBEDDING_TYPE", "sentence-transformer") os.environ.setdefault("EMBEDDING_TYPE", "sentence-transformer")
from app.deterministic_services.vector_store import ( from app.deterministic_services import Source as PackageSource, SourceID as PackageSourceID
from app.deterministic_services.retrieval_store import (
Retrieved, Retrieved,
Source, Source,
SourceID, SourceID,
...@@ -14,6 +16,8 @@ from app.deterministic_services.vector_store import ( ...@@ -14,6 +16,8 @@ from app.deterministic_services.vector_store import (
expand_neighbor_children, expand_neighbor_children,
merge_retrieval_groups, merge_retrieval_groups,
merge_sources, merge_sources,
retrieve,
retrieve_with_subsections,
select_dominant_scope, select_dominant_scope,
) )
...@@ -52,6 +56,10 @@ def _mk_retrieved( ...@@ -52,6 +56,10 @@ def _mk_retrieved(
class VectorStorePipelineUnitTest(unittest.TestCase): class VectorStorePipelineUnitTest(unittest.TestCase):
def test_package_exports_point_to_retrieval_models(self) -> None:
self.assertIs(PackageSource, Source)
self.assertIs(PackageSourceID, SourceID)
def test_merge_sources_prefers_task_childs_on_duplicate(self) -> None: def test_merge_sources_prefers_task_childs_on_duplicate(self) -> None:
child_direct = Source( child_direct = Source(
source_id=SourceID( source_id=SourceID(
...@@ -238,6 +246,47 @@ class VectorStorePipelineUnitTest(unittest.TestCase): ...@@ -238,6 +246,47 @@ class VectorStorePipelineUnitTest(unittest.TestCase):
merged = merge_sources([first], [second]) merged = merge_sources([first], [second])
self.assertEqual(len(merged), 2) self.assertEqual(len(merged), 2)
def test_retrieve_sorts_sources_by_score(self) -> None:
groups = {
"children_direct": [
_mk_retrieved("low", 0.2, 1, 1, 1),
_mk_retrieved("high", 0.9, 1, 1, 1),
]
}
with patch("app.deterministic_services.retrieval_store._use_subsection_retrieval", return_value=False):
with patch("app.deterministic_services.retrieval_store.run_child_retrieval_pipeline", return_value=groups):
result = retrieve(
pg_url="postgresql://unused",
embedder=object(), # type: ignore[arg-type]
query="query",
k=2,
)
self.assertEqual([source.source_id.title for source in result], ["high", "low"])
def test_retrieve_with_subsections_matches_retrieve_without_refs(self) -> None:
expected = [
Source(
source_id=SourceID(title="Child", doc_type="child"),
retrieved_as="children_direct",
source_type="child",
score=0.7,
markdown="md",
)
]
with patch("app.deterministic_services.retrieval_store._use_subsection_retrieval", return_value=False):
with patch("app.deterministic_services.retrieval_store.retrieve", return_value=expected):
result = retrieve_with_subsections(
pg_url="postgresql://unused",
embedder=object(), # type: ignore[arg-type]
query="query",
subsection_refs=None,
)
self.assertEqual(result, expected)
if __name__ == "__main__": if __name__ == "__main__":
unittest.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