kb-app/app/db.py

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,
        }