type to source_type und pydantic hinzugefügt

math-tutor/backend/app/deterministic_services/__init__.py

 # Package marker for services
-from app.deterministic_services.vector_store import Source
-__all__ = ["Source"]
+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

@@ -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

@@ -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,
     )