Скелет проекта

src/rag_agent/retrieval/__init__.py

@@ -0,0 +1 @@
+__all__ = []

src/rag_agent/retrieval/search.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import psycopg
+
+from rag_agent.index.postgres import fetch_similar
+
+
+@dataclass(frozen=True)
+class SearchResult:
+    path: str
+    content: str
+    distance: float
+
+
+def search_similar(
+    conn: psycopg.Connection,
+    query_embedding: list[float],
+    top_k: int = 5,
+    story_id: int | None = None,
+) -> list[SearchResult]:
+    rows = fetch_similar(conn, query_embedding, top_k, story_id=story_id)
+    return [SearchResult(path=row[0], content=row[1], distance=row[2]) for row in rows]