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

type to source_type und pydantic hinzugefügt

parent da7d39aa
Hide whitespace changes
Inline Side-by-side
math-tutor/backend/app/deterministic_services/__init__.py
View file @ 5d1db923
# Package marker for services
from app.deterministic_services.vector_store import Source
__all__ = ["Source"]
\ No newline at end of file
from app.deterministic_services.vector_store import Source, SourceID
__all__ = ["Source", "SourceID"]
\ No newline at end of file
math-tutor/backend/app/deterministic_services/vector_store.py
View file @ 5d1db923
......@@ -10,7 +10,7 @@ from psycopg.rows import dict_row
from pgvector import Vector
from pgvector.psycopg import register_vector
import yaml
from enum import Enum
from pydantic import BaseModel
@dataclass
......@@ -130,7 +130,7 @@ CREATE TABLE IF NOT EXISTS docs (
section_title TEXT NULL,
subsection_title TEXT NULL,
title TEXT NULL,
type TEXT NULL,
source_type TEXT NULL,
path TEXT NOT NULL,
markdown TEXT NOT NULL,
......@@ -141,7 +141,7 @@ CREATE TABLE IF NOT EXISTS docs (
CREATE INDEX IF NOT EXISTS docs_embedding_cos_idx
ON docs USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);
CREATE INDEX IF NOT EXISTS docs_type_idx ON docs(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_section_idx ON docs(section_index);
CREATE INDEX IF NOT EXISTS docs_subsection_idx ON docs(section_index, subsection_index);
......@@ -151,12 +151,12 @@ UPSERT_SQL = """
INSERT INTO docs (
uid, doc_type,
section_index, subsection_index, child_index,
section_title, subsection_title, title, type,
section_title, subsection_title, title, source_type,
path, markdown, embedding
) VALUES (
%(uid)s, %(doc_type)s,
%(section_index)s, %(subsection_index)s, %(child_index)s,
%(section_title)s, %(subsection_title)s, %(title)s, %(type)s,
%(section_title)s, %(subsection_title)s, %(title)s, %(source_type)s,
%(path)s, %(markdown)s, %(embedding)s
)
ON CONFLICT (uid) DO UPDATE SET
......@@ -167,7 +167,7 @@ ON CONFLICT (uid) DO UPDATE SET
section_title = EXCLUDED.section_title,
subsection_title = EXCLUDED.subsection_title,
title = EXCLUDED.title,
type = EXCLUDED.type,
source_type = EXCLUDED.source_type,
path = EXCLUDED.path,
markdown = EXCLUDED.markdown,
embedding = EXCLUDED.embedding
......@@ -210,7 +210,7 @@ def upsert_docs(pg_url: str, docs: List[DocRecord], embeddings: List[List[float]
"section_title": m.get("section_title"),
"subsection_title": m.get("subsection_title"),
"title": m.get("title"),
"type": m.get("type"),
"source_type": m.get("source_type"),
"path": doc.path,
"markdown": doc.markdown,
"embedding": Vector(emb),
......@@ -248,12 +248,11 @@ class Retrieved:
"metadata": self.metadata,
"markdown": self.markdown,
}
@dataclass
class Source:
class Source(BaseModel):
source_id: SourceID
retrieved_as: str
type: str
source_type: str
score: float
markdown: str
......@@ -261,18 +260,24 @@ class Source:
return {
"source_id": self.source_id,
"retrieved_as": self.retrieved_as,
"type": self.type,
"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 retrieved_as={self.retrieved_as},\n type={self.type},\n score={self.score},\n markdown_length={len(self.markdown)})"
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_length={len(self.markdown)})"
)
@dataclass
class SourceID:
chapter_title: Optional[str]
section_title: Optional[str]
subsection_title: Optional[str]
class SourceID(BaseModel):
chapter_title: Optional[str] = None
section_title: Optional[str] = None
subsection_title: Optional[str] = None
title: str
doc_type: str
......@@ -297,7 +302,7 @@ def _row_to_retrieved(row: Dict[str, Any]) -> Retrieved:
"section_title": row["section_title"],
"subsection_title": row["subsection_title"],
"title": row["title"],
"type": row["type"],
"source_type": row["source_type"],
"path": row["path"],
}
return Retrieved(
......@@ -322,7 +327,7 @@ def _retrivla_to_sources(retrievd : Dict[str, List[Retrieved]]) -> List[Source]:
doc_type=retrival.doc_type
),
retrieved_as=name,
type=retrival.metadata.get("type"),
source_type=retrival.metadata.get("source_type") or "unknown",
score=retrival.score,
markdown=retrival.markdown
))
......@@ -336,7 +341,7 @@ def retrieve(
k: int = 4,
section_index: Optional[int] = None,
subsection_index: Optional[int] = None,
type_filter: Optional[List[str]] = None,
source_type_filter: Optional[List[str]] = None,
expand_links: bool = True,
neighbor_expand: int = 0,
) -> List[Source]:
......@@ -353,9 +358,9 @@ def retrieve(
where.append("subsection_index = %(subsection_index)s")
params["subsection_index"] = subsection_index
if type_filter:
where.append("type = ANY(%(type_filter)s)")
params["type_filter"] = type_filter
if source_type_filter:
where.append("source_type = ANY(%(source_type_filter)s)")
params["source_type_filter"] = source_type_filter
where_sql = " AND ".join(where)
......@@ -363,7 +368,7 @@ def retrieve(
SELECT
uid, doc_type,
section_index, subsection_index, child_index,
section_title, subsection_title, title, type,
section_title, subsection_title, title, source_type,
path, markdown,
1 - (embedding <=> %(qvec)s) AS score
FROM docs
......@@ -405,7 +410,7 @@ def retrieve(
SELECT
d.uid, d.doc_type,
d.section_index, d.subsection_index, d.child_index,
d.section_title, d.subsection_title, d.title, d.type,
d.section_title, d.subsection_title, d.title, d.source_type,
d.path, d.markdown,
1.0 AS score
FROM docs d
......@@ -422,7 +427,7 @@ def retrieve(
SELECT
uid, doc_type,
section_index, subsection_index, child_index,
section_title, subsection_title, title, type,
section_title, subsection_title, title, source_type,
path, markdown,
1.0 AS score
FROM docs
......@@ -438,7 +443,7 @@ def retrieve(
SELECT
d.uid, d.doc_type,
d.section_index, d.subsection_index, d.child_index,
d.section_title, d.subsection_title, d.title, d.type,
d.section_title, d.subsection_title, d.title, d.source_type,
d.path, d.markdown,
1 - (d.embedding <=> %(qvec)s) AS score
FROM docs d
......@@ -481,7 +486,7 @@ def retrieve(
SELECT
d.uid, d.doc_type,
d.section_index, d.subsection_index, d.child_index,
d.section_title, d.subsection_title, d.title, d.type,
d.section_title, d.subsection_title, d.title, d.source_type,
d.path, d.markdown,
0 AS score
FROM docs d
......
math-tutor/backend/test/vector_store_test.py
View file @ 5d1db923
......@@ -32,7 +32,7 @@ def main() -> None:
parser.add_argument("--expand", action="store_true")
parser.add_argument("--section-index", type=int, default=None)
parser.add_argument("--subsection-index", type=int, default=None)
parser.add_argument("--type-filter", nargs="*", default=None)
parser.add_argument("--source-type-filter", nargs="*", default=None)
parser.add_argument("--neighbor-expand", type=int, default=0)
args = parser.parse_args()
......@@ -48,7 +48,7 @@ def main() -> None:
expand_links=args.expand,
section_index=args.section_index,
subsection_index=args.subsection_index,
type_filter=args.type_filter,
source_type_filter=args.source_type_filter,
neighbor_expand=args.neighbor_expand,
)
......
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