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

chapter hinzugefügt

parent 98cd81ec
Show whitespace changes
Inline Side-by-side
math-tutor/backend/app/deterministic_services/vector_store.py
View file @ c9d1de0b
...@@ -59,16 +59,19 @@ def parse_markdown_with_frontmatter(text: str) -> Tuple[Dict[str, Any], str]: ...@@ -59,16 +59,19 @@ def parse_markdown_with_frontmatter(text: str) -> Tuple[Dict[str, Any], str]:
def stable_uid(doc_type: str, path: str, meta: Dict[str, Any]) -> str: def stable_uid(doc_type: str, path: str, meta: Dict[str, Any]) -> str:
chapter = meta.get("chapter_index")
sec = meta.get("section_index") sec = meta.get("section_index")
sub = meta.get("subsection_index") sub = meta.get("subsection_index")
child = meta.get("child_index") child = meta.get("child_index")
if doc_type == "chapter" and sub is not None:
raw = f"chapter|c{int(chapter):03d}"
if doc_type == "section" and sec is not None: if doc_type == "section" and sec is not None:
raw = f"section|s{int(sec):03d}" raw = f"section|c{int(chapter):03d}|s{int(sec):03d}"
elif doc_type == "subsection" and sec is not None and sub is not None: elif doc_type == "subsection" and sec is not None and sub is not None:
raw = f"subsection|s{int(sec):03d}|ss{int(sub):03d}" raw = f"subsection|c{int(chapter):03d}|s{int(sec):03d}|ss{int(sub):03d}"
elif doc_type == "child" and sec is not None and sub is not None and child is not None: elif doc_type == "child" and sec is not None and sub is not None and child is not None:
raw = f"child|s{int(sec):03d}|ss{int(sub):03d}|c{int(child):03d}" raw = f"child|c{int(chapter):03d}|s{int(sec):03d}|ss{int(sub):03d}|c{int(child):03d}"
else: else:
raw = f"{doc_type}|{path}" raw = f"{doc_type}|{path}"
...@@ -79,6 +82,7 @@ def stable_uid(doc_type: str, path: str, meta: Dict[str, Any]) -> str: ...@@ -79,6 +82,7 @@ def stable_uid(doc_type: str, path: str, meta: Dict[str, Any]) -> str:
def load_docs(base_dir: Path) -> List[DocRecord]: def load_docs(base_dir: Path) -> List[DocRecord]:
docs: List[DocRecord] = [] docs: List[DocRecord] = []
mapping = [ mapping = [
("chapter", [base_dir / "chapter", base_dir / "chapter"]),
("section", [base_dir / "sections", base_dir / "section"]), ("section", [base_dir / "sections", base_dir / "section"]),
("subsection", [base_dir / "subsections", base_dir / "subsection"]), ("subsection", [base_dir / "subsections", base_dir / "subsection"]),
("child", [base_dir / "childs"]), ("child", [base_dir / "childs"]),
...@@ -110,7 +114,7 @@ def load_docs(base_dir: Path) -> List[DocRecord]: ...@@ -110,7 +114,7 @@ def load_docs(base_dir: Path) -> List[DocRecord]:
if not docs: if not docs:
raise FileNotFoundError( raise FileNotFoundError(
f"No markdown docs found under {base_dir}. Expected sections/subsections/childs." f"No markdown docs found under {base_dir}. Expected chapter/sections/subsections/childs."
) )
return docs return docs
...@@ -126,10 +130,12 @@ CREATE EXTENSION IF NOT EXISTS vector; ...@@ -126,10 +130,12 @@ CREATE EXTENSION IF NOT EXISTS vector;
CREATE TABLE IF NOT EXISTS docs ( CREATE TABLE IF NOT EXISTS docs (
uid TEXT PRIMARY KEY, uid TEXT PRIMARY KEY,
doc_type TEXT NOT NULL, doc_type TEXT NOT NULL,
chapter_index INT NULL,
section_index INT NULL, section_index INT NULL,
subsection_index INT NULL, subsection_index INT NULL,
child_index INT NULL, child_index INT NULL,
chapter_title TEXT NULL,
section_title TEXT NULL, section_title TEXT NULL,
subsection_title TEXT NULL, subsection_title TEXT NULL,
title TEXT NULL, title TEXT NULL,
...@@ -146,8 +152,9 @@ CREATE INDEX IF NOT EXISTS docs_embedding_cos_idx ...@@ -146,8 +152,9 @@ CREATE INDEX IF NOT EXISTS docs_embedding_cos_idx
CREATE INDEX IF NOT EXISTS docs_source_type_idx ON docs(source_type); CREATE INDEX IF NOT EXISTS docs_source_type_idx ON docs(source_type);
CREATE INDEX IF NOT EXISTS docs_doc_type_idx ON docs(doc_type); CREATE INDEX IF NOT EXISTS docs_doc_type_idx ON docs(doc_type);
CREATE INDEX IF NOT EXISTS docs_section_idx ON docs(section_index); CREATE INDEX IF NOT EXISTS docs_chapter_idx ON docs(chapter_index);
CREATE INDEX IF NOT EXISTS docs_subsection_idx ON docs(section_index, subsection_index); CREATE INDEX IF NOT EXISTS docs_section_idx ON docs(chapter_index, section_index);
CREATE INDEX IF NOT EXISTS docs_subsection_idx ON docs(chapter_index, section_index, subsection_index);
""" """
...@@ -166,20 +173,22 @@ def init_db(pg_url: str) -> None: ...@@ -166,20 +173,22 @@ def init_db(pg_url: str) -> None:
UPSERT_SQL = """ UPSERT_SQL = """
INSERT INTO docs ( INSERT INTO docs (
uid, doc_type, uid, doc_type,
section_index, subsection_index, child_index, chapter_index, section_index, subsection_index, child_index,
section_title, subsection_title, title, source_type, chapter_title, section_title, subsection_title, title, source_type,
path, markdown, embedding path, markdown, embedding
) VALUES ( ) VALUES (
%(uid)s, %(doc_type)s, %(uid)s, %(doc_type)s,
%(section_index)s, %(subsection_index)s, %(child_index)s, %(chapter_index)s, %(section_index)s, %(subsection_index)s, %(child_index)s,
%(section_title)s, %(subsection_title)s, %(title)s, %(source_type)s, %(chapter_title)s, %(section_title)s, %(subsection_title)s, %(title)s, %(source_type)s,
%(path)s, %(markdown)s, %(embedding)s %(path)s, %(markdown)s, %(embedding)s
) )
ON CONFLICT (uid) DO UPDATE SET ON CONFLICT (uid) DO UPDATE SET
doc_type = EXCLUDED.doc_type, doc_type = EXCLUDED.doc_type,
chapter_index = EXCLUDED.chapter_index,
section_index = EXCLUDED.section_index, section_index = EXCLUDED.section_index,
subsection_index = EXCLUDED.subsection_index, subsection_index = EXCLUDED.subsection_index,
child_index = EXCLUDED.child_index, child_index = EXCLUDED.child_index,
chapter_title = EXCLUDED.chapter_title,
section_title = EXCLUDED.section_title, section_title = EXCLUDED.section_title,
subsection_title = EXCLUDED.subsection_title, subsection_title = EXCLUDED.subsection_title,
title = EXCLUDED.title, title = EXCLUDED.title,
...@@ -212,9 +221,11 @@ def upsert_docs(pg_url: str, docs: List[DocRecord], embeddings: List[List[float] ...@@ -212,9 +221,11 @@ def upsert_docs(pg_url: str, docs: List[DocRecord], embeddings: List[List[float]
{ {
"uid": doc.uid, "uid": doc.uid,
"doc_type": doc.doc_type, "doc_type": doc.doc_type,
"chapter_index": _meta_int(m, "chapter_index"),
"section_index": _meta_int(m, "section_index"), "section_index": _meta_int(m, "section_index"),
"subsection_index": _meta_int(m, "subsection_index"), "subsection_index": _meta_int(m, "subsection_index"),
"child_index": _meta_int(m, "child_index"), "child_index": _meta_int(m, "child_index"),
"chapter_title": m.get("chapter_title"),
"section_title": m.get("section_title"), "section_title": m.get("section_title"),
"subsection_title": m.get("subsection_title"), "subsection_title": m.get("subsection_title"),
"title": m.get("title"), "title": m.get("title"),
...@@ -274,6 +285,7 @@ def _row_to_retrieved(row: Dict[str, Any], source_type: Optional[str] = None) -> ...@@ -274,6 +285,7 @@ def _row_to_retrieved(row: Dict[str, Any], source_type: Optional[str] = None) ->
meta = { meta = {
"uid": row["uid"], "uid": row["uid"],
"doc_type": row["doc_type"], "doc_type": row["doc_type"],
"chapter_index": row["chapter_index"],
"section_index": row["section_index"], "section_index": row["section_index"],
"subsection_index": row["subsection_index"], "subsection_index": row["subsection_index"],
"child_index": row["child_index"], "child_index": row["child_index"],
...@@ -297,6 +309,7 @@ def retrieve( ...@@ -297,6 +309,7 @@ def retrieve(
embedder: EmbeddingLike, embedder: EmbeddingLike,
query: str, query: str,
k: int = 4, k: int = 4,
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,
...@@ -308,6 +321,10 @@ def retrieve( ...@@ -308,6 +321,10 @@ def retrieve(
where = ["doc_type = 'child'"] where = ["doc_type = 'child'"]
params: Dict[str, Any] = {"qvec": qvec, "k": k} params: Dict[str, Any] = {"qvec": qvec, "k": k}
if chapter_index is not None:
where.append("chapter_index = %(chapter_index)s")
params["chapter_index"] = chapter_index
if section_index is not None: if 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"] = section_index
...@@ -325,8 +342,8 @@ def retrieve( ...@@ -325,8 +342,8 @@ def retrieve(
sql = f""" sql = f"""
SELECT SELECT
uid, doc_type, uid, doc_type,
section_index, subsection_index, child_index, chapter_index, section_index, subsection_index, child_index,
section_title, subsection_title, title, source_type, chapter_title, section_title, subsection_title, title, source_type,
path, markdown, path, markdown,
1 - (embedding <=> %(qvec)s) AS score 1 - (embedding <=> %(qvec)s) AS score
FROM docs FROM docs
...@@ -337,6 +354,7 @@ def retrieve( ...@@ -337,6 +354,7 @@ def retrieve(
children: List[Retrieved] = [] children: List[Retrieved] = []
neighbors: List[Retrieved] = [] neighbors: List[Retrieved] = []
chapter: List[Retrieved] = []
subsections: List[Retrieved] = [] subsections: List[Retrieved] = []
sections_docs: List[Retrieved] = [] sections_docs: List[Retrieved] = []
children_direct: List[Retrieved] = [] children_direct: List[Retrieved] = []
...@@ -351,34 +369,36 @@ def retrieve( ...@@ -351,34 +369,36 @@ def retrieve(
children_direct = list(children) children_direct = list(children)
if expand_links and children: if expand_links and children:
sec_sub_counts: Dict[tuple[int, int], int] = {} cpt_sec_sub_counts: Dict[tuple[int, int, int], int] = {}
for child in children: for child in children:
cpt_idx = child.metadata.get("chapter_index")
sec_idx = child.metadata.get("section_index") sec_idx = child.metadata.get("section_index")
sub_idx = child.metadata.get("subsection_index") sub_idx = child.metadata.get("subsection_index")
if sec_idx is not None and sub_idx is not None: if cpt_idx is not None and sec_idx is not None and sub_idx is not None:
key = (int(sec_idx), int(sub_idx)) key = (int(cpt_idx), int(sec_idx), int(sub_idx))
sec_sub_counts[key] = sec_sub_counts.get(key, 0) + 1 cpt_sec_sub_counts[key] = cpt_sec_sub_counts.get(key, 0) + 1
if sec_sub_counts: if cpt_sec_sub_counts:
most_common_sec_sub = max( most_common_cpt_sec_sub = max(
sec_sub_counts.items(), key=lambda x: x[1])[0] cpt_sec_sub_counts.items(), key=lambda x: x[1])[0]
most_common_sec, most_common_sub = most_common_sec_sub most_common_cpt, most_common_sec, most_common_sub = most_common_cpt_sec_sub
cur.execute( cur.execute(
""" """
SELECT SELECT
d.uid, d.doc_type, d.uid, d.doc_type,
d.section_index, d.subsection_index, d.child_index, d.chapter_index, d.section_index, d.subsection_index, d.child_index,
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,
1.0 AS score 1.0 AS score
FROM docs d FROM docs d
WHERE d.doc_type = ANY(%(sub_doc_types)s) WHERE d.doc_type = ANY(%(sub_doc_types)s)
AND d.chapter_index = %(cpt)s
AND d.section_index = %(sec)s AND d.section_index = %(sec)s
AND d.subsection_index = %(sub)s AND d.subsection_index = %(sub)s
""", """,
{"sec": most_common_sec, "sub": most_common_sub, {"cpt": most_common_cpt, "sec": most_common_sec, "sub": most_common_sub,
"sub_doc_types": ["subsection", "chapter"]}, "sub_doc_types": ["chapter", "section", "subsection"]},
) )
subsections = [_row_to_retrieved( subsections = [_row_to_retrieved(
row, source_type="subsection") for row in cur.fetchall()] row, source_type="subsection") for row in cur.fetchall()]
...@@ -387,17 +407,18 @@ def retrieve( ...@@ -387,17 +407,18 @@ def retrieve(
cur.execute( cur.execute(
""" """
SELECT SELECT
uid, doc_type, d.uid, d.doc_type,
section_index, subsection_index, child_index, d.chapter_index, d.section_index, d.subsection_index, d.child_index,
section_title, subsection_title, title, source_type, d.chapter_title, d.section_title, d.subsection_title, d.title, d.source_type,
path, markdown, d.path, d.markdown,
1.0 AS score 1.0 AS score
FROM docs FROM docs d
WHERE doc_type = ANY(%(sec_doc_types)s) WHERE d.doc_type = ANY(%(sec_doc_types)s)
AND section_index = %(sec)s AND d.chapter_index = %(cpt)s
AND d.section_index = %(sec)s
""", """,
{"sec": most_common_sec, "sec_doc_types": [ {"cpt": most_common_cpt, "sec": most_common_sec, "sec_doc_types": [
"section", "oberchapter"]}, "section", "chapter"]},
) )
sections_docs = [_row_to_retrieved( sections_docs = [_row_to_retrieved(
row, source_type="section") for row in cur.fetchall()] row, source_type="section") for row in cur.fetchall()]
...@@ -406,18 +427,37 @@ def retrieve( ...@@ -406,18 +427,37 @@ def retrieve(
""" """
SELECT SELECT
d.uid, d.doc_type, d.uid, d.doc_type,
d.section_index, d.subsection_index, d.child_index, d.chapter_index, d.section_index, d.subsection_index, d.child_index,
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,
1.0 AS score
FROM docs d
WHERE d.doc_type = ANY(%(cpt_doc_types)s)
AND d.chapter_index = %(cpt)s
""",
{"cpt": most_common_cpt, "cpt_doc_types": [
"chapter"]},
)
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, d.path, d.markdown,
1 - (d.embedding <=> %(qvec)s) AS score 1 - (d.embedding <=> %(qvec)s) AS score
FROM docs d FROM docs d
WHERE d.doc_type = 'child' WHERE d.doc_type = 'child'
AND chapter_index = %(cpt)s
AND d.section_index = %(sec)s AND d.section_index = %(sec)s
AND d.subsection_index = %(sub)s AND d.subsection_index = %(sub)s
ORDER BY d.embedding <=> %(qvec)s ORDER BY d.embedding <=> %(qvec)s
LIMIT 5; LIMIT 5;
""", """,
{"qvec": qvec, "sec": most_common_sec, {"qvec": qvec, "cpt": most_common_cpt, "sec": most_common_sec,
"sub": most_common_sub}, "sub": most_common_sub},
) )
...@@ -432,36 +472,38 @@ def retrieve( ...@@ -432,36 +472,38 @@ def retrieve(
if neighbor_expand and neighbor_expand > 0: if neighbor_expand and neighbor_expand > 0:
wanted: set[tuple[int, int, int]] = set() wanted: set[tuple[int, int, int]] = set()
for child in children: for child in children:
cpti = child.metadata.get("chapter_index")
si = child.metadata.get("section_index") si = child.metadata.get("section_index")
ssi = child.metadata.get("subsection_index") ssi = child.metadata.get("subsection_index")
ci = child.metadata.get("child_index") ci = child.metadata.get("child_index")
if si is None or ssi is None or ci is None: if cpti is None or si is None or ssi is None or ci is None:
continue continue
for dx in range(-neighbor_expand, neighbor_expand + 1): for dx in range(-neighbor_expand, neighbor_expand + 1):
if dx == 0: if dx == 0:
continue continue
wanted.add((int(si), int(ssi), int(ci) + dx)) wanted.add((int(cpti), int(si), int(ssi), int(ci) + dx))
if wanted: if wanted:
triples = sorted(wanted) quadruple = sorted(wanted)
sec_arr = [a for (a, b, cidx) in triples] cpt_arr = [a for (a, b, c, cidx) in quadruple]
sub_arr = [b for (a, b, cidx) in triples] sec_arr = [b for (a, b, c, cidx) in quadruple]
child_arr = [cidx for (a, b, cidx) in triples] sub_arr = [c for (a, b, c, cidx) in quadruple]
child_arr = [cidx for (a, b, c, cidx) in quadruple]
cur.execute( cur.execute(
""" """
SELECT SELECT
d.uid, d.doc_type, d.uid, d.doc_type,
d.section_index, d.subsection_index, d.child_index, d.chapter_index, d.section_index, d.subsection_index, d.child_index,
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(%(sec_arr)s::int[], %(sub_arr)s::int[], %(child_arr)s::int[]) AS u(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.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'
""", """,
{"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},
) )
neighbors = [_row_to_retrieved( neighbors = [_row_to_retrieved(
...@@ -473,6 +515,7 @@ def retrieve( ...@@ -473,6 +515,7 @@ def retrieve(
retrivla_dict = { retrivla_dict = {
"children_direct": children_direct, "children_direct": children_direct,
"children_expanded": children_expanded, "children_expanded": children_expanded,
"chapters": chapter_docs,
"subsections": subsections, "subsections": subsections,
"sections": sections_docs, "sections": sections_docs,
"neighbors": neighbors, "neighbors": neighbors,
...@@ -650,7 +693,7 @@ def _retrivla_to_sources(retrievd: Dict[str, List[Retrieved]]) -> List[Source]: ...@@ -650,7 +693,7 @@ def _retrivla_to_sources(retrievd: Dict[str, List[Retrieved]]) -> List[Source]:
for retrival in retrieved_grouep: for retrival in retrieved_grouep:
sources.append(Source( sources.append(Source(
source_id=SourceID( source_id=SourceID(
chapter_title="none", chapter_title=retrival.metadata.get("chapter_title"),
section_title=retrival.metadata.get("section_title"), section_title=retrival.metadata.get("section_title"),
subsection_title=retrival.metadata.get("subsection_title"), subsection_title=retrival.metadata.get("subsection_title"),
title=retrival.metadata.get("title"), title=retrival.metadata.get("title"),
......
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