OCR-Viewer (#1)

Co-authored-by: nathan <nathan.kehler@gmail.com> Reviewed-on: #1
2026-05-19 22:03:28 +00:00
parent 727f4fcc57
commit 330b703539
104 changed files with 144649 additions and 2 deletions
--- a/GoogleDocumentOCR/db.py
+++ b/GoogleDocumentOCR/db.py
@@ -0,0 +1,108 @@
+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