108 lines
3.2 KiB
Python
108 lines
3.2 KiB
Python
import os
|
|
import psycopg2
|
|
from psycopg2.extras import execute_values
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
def get_conn():
|
|
return psycopg2.connect(
|
|
host=os.getenv("DB_HOST"),
|
|
port=os.getenv("DB_PORT"),
|
|
dbname=os.getenv("DB_NAME"),
|
|
user=os.getenv("DB_USER"),
|
|
password=os.getenv("DB_PASSWORD"),
|
|
)
|
|
|
|
def init_db():
|
|
conn = get_conn()
|
|
cur = conn.cursor()
|
|
cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS document_chunks (
|
|
id SERIAL PRIMARY KEY,
|
|
filename TEXT NOT NULL,
|
|
page_num INTEGER,
|
|
chunk_index INTEGER,
|
|
text TEXT NOT NULL,
|
|
embedding vector(768),
|
|
nationality TEXT DEFAULT 'unknown',
|
|
corpus TEXT DEFAULT 'unknown',
|
|
created_at TIMESTAMPTZ DEFAULT NOW()
|
|
);
|
|
""")
|
|
# Migrate existing tables that predate nationality/corpus columns
|
|
for col, default in [("nationality", "unknown"), ("corpus", "unknown")]:
|
|
cur.execute(f"""
|
|
ALTER TABLE document_chunks
|
|
ADD COLUMN IF NOT EXISTS {col} TEXT DEFAULT '{default}';
|
|
""")
|
|
cur.execute("""
|
|
CREATE INDEX IF NOT EXISTS document_chunks_embedding_idx
|
|
ON document_chunks
|
|
USING ivfflat (embedding vector_cosine_ops)
|
|
WITH (lists = 100);
|
|
""")
|
|
cur.execute("""
|
|
CREATE INDEX IF NOT EXISTS document_chunks_nationality_idx
|
|
ON document_chunks (nationality);
|
|
""")
|
|
cur.execute("""
|
|
CREATE INDEX IF NOT EXISTS document_chunks_corpus_idx
|
|
ON document_chunks (corpus);
|
|
""")
|
|
conn.commit()
|
|
cur.close()
|
|
conn.close()
|
|
print("DB initialized.")
|
|
|
|
def upsert_chunks(rows: list[dict]):
|
|
conn = get_conn()
|
|
cur = conn.cursor()
|
|
execute_values(
|
|
cur,
|
|
"""
|
|
INSERT INTO document_chunks (filename, page_num, chunk_index, text, embedding, nationality, corpus)
|
|
VALUES %s
|
|
""",
|
|
[
|
|
(
|
|
r["filename"], r["page_num"], r["chunk_index"], r["text"], r["embedding"],
|
|
r.get("nationality", "unknown"), r.get("corpus", "unknown"),
|
|
)
|
|
for r in rows
|
|
],
|
|
template="(%s, %s, %s, %s, %s::vector, %s, %s)"
|
|
)
|
|
conn.commit()
|
|
cur.close()
|
|
conn.close()
|
|
|
|
def semantic_search(query_embedding: list[float], top_k: int = 5,
|
|
nationality: str = None, corpus: str = None):
|
|
conn = get_conn()
|
|
cur = conn.cursor()
|
|
filters = []
|
|
params = []
|
|
if nationality:
|
|
filters.append("nationality = %s")
|
|
params.append(nationality)
|
|
if corpus:
|
|
filters.append("corpus = %s")
|
|
params.append(corpus)
|
|
where = ("WHERE " + " AND ".join(filters)) if filters else ""
|
|
cur.execute(
|
|
f"""
|
|
SELECT filename, page_num, chunk_index, text, nationality, corpus,
|
|
1 - (embedding <=> %s::vector) AS similarity
|
|
FROM document_chunks
|
|
{where}
|
|
ORDER BY embedding <=> %s::vector
|
|
LIMIT %s;
|
|
""",
|
|
[query_embedding] + params + [query_embedding, top_k]
|
|
)
|
|
rows = cur.fetchall()
|
|
cur.close()
|
|
conn.close()
|
|
return rows |