Initial commit — kb-app RAG server

app/db.py

@@ -0,0 +1,150 @@
+import psycopg
+from psycopg.rows import dict_row
+from pgvector.psycopg import register_vector
+from contextlib import contextmanager
+from .config import CONFIG
+
+
+@contextmanager
+def get_conn():
+    conn = psycopg.connect(
+        host=CONFIG.pg_host,
+        port=CONFIG.pg_port,
+        dbname=CONFIG.pg_db,
+        user=CONFIG.pg_user,
+        password=CONFIG.pg_password,
+        row_factory=dict_row,
+    )
+    register_vector(conn)
+    try:
+        yield conn
+        conn.commit()
+    except Exception:
+        conn.rollback()
+        raise
+    finally:
+        conn.close()
+
+
+def upsert_document(source_path, file_hash, file_type, category, title, last_modified, metadata):
+    with get_conn() as conn:
+        cur = conn.cursor()
+        cur.execute("SELECT id, file_hash FROM kb.documents WHERE source_path = %s", (source_path,))
+        existing = cur.fetchone()
+        
+        if existing and existing["file_hash"] == file_hash:
+            return existing["id"], False
+        
+        if existing:
+            cur.execute("""
+                UPDATE kb.documents 
+                SET file_hash = %s, file_type = %s, category = %s, title = %s,
+                    last_modified = %s, indexed_at = NOW(), metadata = %s
+                WHERE id = %s RETURNING id
+            """, (file_hash, file_type, category, title, last_modified,
+                  psycopg.types.json.Jsonb(metadata), existing["id"]))
+            doc_id = cur.fetchone()["id"]
+            cur.execute("DELETE FROM kb.chunks WHERE document_id = %s", (doc_id,))
+        else:
+            cur.execute("""
+                INSERT INTO kb.documents 
+                (source_path, file_hash, file_type, category, title, last_modified, metadata)
+                VALUES (%s, %s, %s, %s, %s, %s, %s) RETURNING id
+            """, (source_path, file_hash, file_type, category, title, last_modified,
+                  psycopg.types.json.Jsonb(metadata)))
+            doc_id = cur.fetchone()["id"]
+        
+        return doc_id, True
+
+
+def insert_chunks(document_id, chunks):
+    with get_conn() as conn:
+        cur = conn.cursor()
+        for c in chunks:
+            cur.execute("""
+                INSERT INTO kb.chunks 
+                (document_id, chunk_index, content, embedding, token_count, metadata)
+                VALUES (%s, %s, %s, %s, %s, %s)
+            """, (document_id, c["chunk_index"], c["content"], c["embedding"],
+                  c["token_count"], psycopg.types.json.Jsonb(c.get("metadata", {}))))
+
+
+def search(query_embedding, k=10, category=None):
+    with get_conn() as conn:
+        cur = conn.cursor()
+        if category:
+            cur.execute("""
+                SELECT 
+                    c.id as chunk_id, c.content, c.chunk_index, 
+                    c.metadata as chunk_metadata, c.document_id,
+                    d.source_path, d.category, d.title, d.metadata as doc_metadata,
+                    1 - (c.embedding <=> %s::vector) as similarity
+                FROM kb.chunks c
+                JOIN kb.documents d ON c.document_id = d.id
+                WHERE d.category = %s
+                ORDER BY c.embedding <=> %s::vector
+                LIMIT %s
+            """, (query_embedding, category, query_embedding, k))
+        else:
+            cur.execute("""
+                SELECT 
+                    c.id as chunk_id, c.content, c.chunk_index,
+                    c.metadata as chunk_metadata, c.document_id,
+                    d.source_path, d.category, d.title, d.metadata as doc_metadata,
+                    1 - (c.embedding <=> %s::vector) as similarity
+                FROM kb.chunks c
+                JOIN kb.documents d ON c.document_id = d.id
+                ORDER BY c.embedding <=> %s::vector
+                LIMIT %s
+            """, (query_embedding, query_embedding, k))
+        return cur.fetchall()
+
+
+def search_with_neighbors(query_embedding, k=10, category=None, neighbor_window=1):
+    """Search and include neighboring chunks for more context."""
+    primary = search(query_embedding, k=k, category=category)
+    
+    if not primary or neighbor_window <= 0:
+        return primary
+    
+    with get_conn() as conn:
+        cur = conn.cursor()
+        for r in primary:
+            doc_id = r["document_id"]
+            chunk_idx = r["chunk_index"]
+            
+            cur.execute("""
+                SELECT chunk_index, content, metadata
+                FROM kb.chunks
+                WHERE document_id = %s 
+                  AND chunk_index BETWEEN %s AND %s
+                ORDER BY chunk_index
+            """, (doc_id, chunk_idx - neighbor_window, chunk_idx + neighbor_window))
+            
+            neighbors = cur.fetchall()
+            # Combine into one expanded passage
+            r["expanded_content"] = "\n\n[...]\n\n".join(n["content"] for n in neighbors)
+            r["neighbor_count"] = len(neighbors) - 1
+        
+        return primary
+
+
+def get_stats():
+    with get_conn() as conn:
+        cur = conn.cursor()
+        cur.execute("SELECT COUNT(*) as total_docs FROM kb.documents")
+        docs = cur.fetchone()
+        cur.execute("SELECT COUNT(*) as total_chunks FROM kb.chunks")
+        chunks = cur.fetchone()
+        cur.execute("""
+            SELECT category, COUNT(*) as count 
+            FROM kb.documents 
+            GROUP BY category 
+            ORDER BY count DESC
+        """)
+        by_cat = cur.fetchall()
+        return {
+            "total_documents": docs["total_docs"],
+            "total_chunks": chunks["total_chunks"],
+            "by_category": by_cat,
+        }