app/database/documents.py

from __future__ import annotations
from .client import get_service_client


async def create_document(name: str, uploaded_by: str | None = None) -> dict:
    client = get_service_client()
    resp = client.table("documents").insert(
        {"name": name, "uploaded_by": uploaded_by}
    ).execute()
    return resp.data[0]


async def update_document_status(doc_id: str, status: str) -> None:
    client = get_service_client()
    client.table("documents").update({"status": status}).eq("id", doc_id).execute()


async def get_all_documents() -> list[dict]:
    client = get_service_client()
    resp = client.table("documents").select("*").order("created_at", desc=True).execute()
    return resp.data


async def get_document_count() -> int:
    client = get_service_client()
    resp = client.table("documents").select("id", count="exact").execute()
    return resp.count or 0


async def insert_chunks(chunks: list[dict]) -> None:
    client = get_service_client()
    batch_size = 50
    for i in range(0, len(chunks), batch_size):
        batch = chunks[i : i + batch_size]
        client.table("document_chunks").insert(batch).execute()


async def search_chunks(embedding: list[float], top_k: int = 5, threshold: float = 0.5) -> list[dict]:
    client = get_service_client()
    resp = client.rpc(
        "match_document_chunks",
        {
            "query_embedding": embedding,
            "match_count": top_k,
            "match_threshold": threshold,
        },
    ).execute()
    return resp.data