Commit 48afff57 authored by Kantz's avatar Kantz
Browse files

aufräumen 2

parent e5dc46a3
...@@ -5,7 +5,7 @@ from datetime import datetime ...@@ -5,7 +5,7 @@ from datetime import datetime
def _format_log_timestamp(created_at: str | None) -> str: def _format_log_timestamp(created_at: str | None) -> str:
if not created_at: if not created_at:
return datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ") return datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
try: try:
parsed = datetime.strptime(created_at, "%Y-%m-%dT%H:%M:%SZ") parsed = datetime.strptime(created_at, "%Y-%m-%dT%H:%M:%SZ")
return parsed.strftime("%Y-%m-%dT%H:%M:%SZ") return parsed.strftime("%Y-%m-%dT%H:%M:%SZ")
......
...@@ -16,6 +16,10 @@ import app.config ...@@ -16,6 +16,10 @@ import app.config
embedding_dim = app.config.get_embedding_settings().target_dim embedding_dim = app.config.get_embedding_settings().target_dim
# -----------------------------
# Einlesen der Dokumente
# -----------------------------
@dataclass @dataclass
class DocRecord: class DocRecord:
doc_type: str doc_type: str
...@@ -110,16 +114,11 @@ def load_docs(base_dir: Path) -> List[DocRecord]: ...@@ -110,16 +114,11 @@ def load_docs(base_dir: Path) -> List[DocRecord]:
return docs return docs
# -----------------------------
# Init der Databse
# -----------------------------
def embed_documents(embedder: EmbeddingLike, texts: List[str]) -> List[List[float]]: DATABASE_CREATION_SQL = f"""
return embedder.embed_documents(texts)
def embed_query(embedder: EmbeddingLike, text: str) -> List[float]:
return embedder.embed_query(text)
DDL = f"""
CREATE EXTENSION IF NOT EXISTS vector; CREATE EXTENSION IF NOT EXISTS vector;
CREATE TABLE IF NOT EXISTS docs ( CREATE TABLE IF NOT EXISTS docs (
...@@ -149,6 +148,17 @@ CREATE INDEX IF NOT EXISTS docs_section_idx ON docs(section_index); ...@@ -149,6 +148,17 @@ 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); CREATE INDEX IF NOT EXISTS docs_subsection_idx ON docs(section_index, subsection_index);
""" """
def init_db(pg_url: str) -> None:
with psycopg.connect(pg_url) as conn:
with conn.cursor() as cur:
cur.execute(DATABASE_CREATION_SQL)
register_vector(conn)
conn.commit()
# -----------------------------
# Einfügen der Dokumente und Embeddings
# -----------------------------
UPSERT_SQL = """ UPSERT_SQL = """
INSERT INTO docs ( INSERT INTO docs (
uid, doc_type, uid, doc_type,
...@@ -176,15 +186,6 @@ ON CONFLICT (uid) DO UPDATE SET ...@@ -176,15 +186,6 @@ ON CONFLICT (uid) DO UPDATE SET
; ;
""" """
def init_db(pg_url: str) -> None:
with psycopg.connect(pg_url) as conn:
with conn.cursor() as cur:
cur.execute(DDL)
register_vector(conn)
conn.commit()
def _meta_int(meta: Dict[str, Any], key: str) -> Optional[int]: def _meta_int(meta: Dict[str, Any], key: str) -> Optional[int]:
value = meta.get(key) value = meta.get(key)
if value is None or value == "": if value is None or value == "":
...@@ -233,6 +234,16 @@ def clear_docs(pg_url: str) -> None: ...@@ -233,6 +234,16 @@ def clear_docs(pg_url: str) -> None:
cur.execute("TRUNCATE TABLE docs;") cur.execute("TRUNCATE TABLE docs;")
conn.commit() conn.commit()
def embed_documents(embedder: EmbeddingLike, texts: List[str]) -> List[List[float]]:
return embedder.embed_documents(texts)
def embed_query(embedder: EmbeddingLike, text: str) -> List[float]:
return embedder.embed_query(text)
# -----------------------------
# Retrival der Dokumente
# -----------------------------
@dataclass @dataclass
class Retrieved: class Retrieved:
...@@ -251,58 +262,6 @@ class Retrieved: ...@@ -251,58 +262,6 @@ class Retrieved:
"markdown": self.markdown, "markdown": self.markdown,
} }
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
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,
"title": self.title,
"doc_type": self.doc_type,
}
def to_string(self) -> str:
string_rep = self.title
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 _row_to_retrieved(row: Dict[str, Any], source_type: Optional[str] = None) -> Retrieved: def _row_to_retrieved(row: Dict[str, Any], source_type: Optional[str] = None) -> Retrieved:
meta = { meta = {
"uid": row["uid"], "uid": row["uid"],
...@@ -325,26 +284,6 @@ def _row_to_retrieved(row: Dict[str, Any], source_type: Optional[str] = None) -> ...@@ -325,26 +284,6 @@ def _row_to_retrieved(row: Dict[str, Any], source_type: Optional[str] = None) ->
) )
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 = "none",
section_title=retrival.metadata.get("section_title"),
subsection_title=retrival.metadata.get("subsection_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
def retrieve( def retrieve(
pg_url: str, pg_url: str,
embedder: EmbeddingLike, embedder: EmbeddingLike,
...@@ -521,6 +460,83 @@ def retrieve( ...@@ -521,6 +460,83 @@ def retrieve(
sources = _retrivla_to_sources(retrivla_dict) sources = _retrivla_to_sources(retrivla_dict)
return sources return sources
# -----------------------------
# Retrival in Sources umwandeln
# -----------------------------
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
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,
"title": self.title,
"doc_type": self.doc_type,
}
def to_string(self) -> str:
string_rep = self.title
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 = "none",
section_title=retrival.metadata.get("section_title"),
subsection_title=retrival.metadata.get("subsection_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
# -----------------------------
# Listen für Filterung
# -----------------------------
def list_sections(pg_url: str) -> List[Dict[str, Any]]: def list_sections(pg_url: str) -> List[Dict[str, Any]]:
sql = """ sql = """
......
import argparse import argparse
from app import config from app import config
from app.deterministic_services.embeddings import OpenAILikeEmbeddings from app.deterministic_services.embeddings import EmbeddingFactory, OpenAILikeEmbeddings
from app.deterministic_services import vector_store from app.deterministic_services import vector_store
def _get_embedder() -> OpenAILikeEmbeddings:
settings = config.get_embedding_settings()
return OpenAILikeEmbeddings(
base_url=settings.base_url,
api_key=settings.api_key,
model=settings.model,
target_dim=settings.target_dim,
)
def _normalize_sources(result: object) -> list[vector_store.Source]: def _normalize_sources(result: object) -> list[vector_store.Source]:
if isinstance(result, list): if isinstance(result, list):
return result return result
...@@ -38,7 +28,7 @@ def main() -> None: ...@@ -38,7 +28,7 @@ def main() -> None:
args = parser.parse_args() args = parser.parse_args()
pg_url = args.pg or config.get_postgres_url() pg_url = args.pg or config.get_postgres_url()
embedder = _get_embedder() embedder = EmbeddingFactory.create(config.get_embedding_settings())
result = vector_store.retrieve( result = vector_store.retrieve(
pg_url=pg_url, pg_url=pg_url,
......
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