modules/vector_store.py

import os
import psycopg2
from typing import List, Dict, Any

# DB connection parameters
DB_HOST = os.getenv("DB_HOST")
DB_NAME = os.getenv("DB_NAME")
DB_USER = os.getenv("DB_USER")
DB_PASS = os.getenv("DB_PASS")

def get_connection():
    return psycopg2.connect(
        host=DB_HOST,
        database=DB_NAME,
        user=DB_USER,
        password=DB_PASS
    )

def search_relevant_chunks(query_embedding: List[float], limit: int = 5) -> List[str]:
    """
    Searches the pgvector database for the most similar text chunks.
    Assumes a table 'ebook_chunks' with columns 'content' (text) and 'embedding' (vector).
    """
    conn = get_connection()
    cur = conn.cursor()
    
    # Convert list to string format for pgvector: [0.1, 0.2, ...]
    embedding_str = "[" + ",".join(map(str, query_embedding)) + "]"
    
    query = """
    SELECT content
    FROM ebook_chunks
    ORDER BY embedding <=> %s::vector
    LIMIT %s
    """
    
    try:
        cur.execute(query, (embedding_str, limit))
        rows = cur.fetchall()
        return [row[0] for row in rows]
    finally:
        cur.close()
        conn.close()

def get_stats() -> Dict[str, Any]:
    """
    Returns counts of books and chunks in the library.
    """
    conn = get_connection()
    cur = conn.cursor()
    
    try:
        cur.execute("SELECT COUNT(*) FROM ebook_chunks")
        chunk_count = cur.fetchone()[0]
        
        # Assuming metadata table or distinct titles in chunks
        cur.execute("SELECT COUNT(DISTINCT ebook_id) FROM ebook_chunks")
        book_count = cur.fetchone()[0]
        
        return {
            "total_books": book_count,
            "total_chunks": chunk_count
        }
    except:
        return {"error": "Could not retrieve stats."}
    finally:
        cur.close()
        conn.close()