rag-server/app/main.py

from __future__ import annotations

from io import BytesIO
from typing import Any
import logging

from fastapi import FastAPI, Header, HTTPException
from minio import Minio
from pgvector.psycopg import register_vector
from pydantic import BaseModel, Field
from psycopg import connect
from psycopg.rows import dict_row
from pypdf import PdfReader

# ✅ LangChain
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_groq import ChatGroq

from app.config import settings
import os
logger = logging.getLogger(__name__)

app = FastAPI(title="Courtrix RAG Service", version="0.2.0")

# ✅ Embedding


embedding_model = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2"

)

# ✅ LLM
llm = ChatGroq(
    groq_api_key="YOUR_GROQ_API_KEY",
    model_name="llama3-70b-8192",
    temperature=0.2
)

# ✅ Text Splitter
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=settings.chunk_size,
    chunk_overlap=settings.chunk_overlap
)

# ✅ MinIO
minio_client = Minio(
    endpoint="minio:9000",
    access_key=settings.minio_access_key,
    secret_key=settings.minio_secret_key,
    secure=False,
)


# ================= MODELS =================

class IngestRequest(BaseModel):
    owner_id: str
    case_id: str
    file_id: str
    bucket: str
    object_key: str
    file_name: str


class HistoryTurn(BaseModel):
    question: str
    answer: str


class AnswerRequest(BaseModel):
    owner_id: str
    case_id: str
    question: str = Field(min_length=2, max_length=4000)
    history: list[HistoryTurn] = Field(default_factory=list)
    top_k: int = Field(default=settings.default_top_k, ge=1, le=12)


# ================= DB =================

def get_db_connection():
    connection = connect(settings.database_url, row_factory=dict_row)
    register_vector(connection)
    return connection


def ensure_rag_table():
    with get_db_connection() as conn:
        with conn.cursor() as cur:
            cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
            cur.execute(
                f"""
                CREATE TABLE IF NOT EXISTS rag_chunks (
                    id BIGSERIAL PRIMARY KEY,
                    owner_id TEXT,
                    case_id TEXT,
                    file_id TEXT,
                    file_name TEXT,
                    page_number INTEGER,
                    chunk_index INTEGER,
                    chunk_text TEXT,
                    embedding VECTOR(384),
                    created_at TIMESTAMPTZ DEFAULT NOW()
                );
                """
            )
        conn.commit()


@app.on_event("startup")
def startup():
    ensure_rag_table()


# ================= HELPERS =================

def download_file_bytes(bucket: str, object_key: str) -> bytes:
    response = minio_client.get_object(bucket, object_key)
    try:
        return response.read()
    finally:
        response.close()
        response.release_conn()


def extract_text(file_bytes: bytes) -> list[dict[str, Any]]:
    reader = PdfReader(BytesIO(file_bytes))
    pages = []

    for i, page in enumerate(reader.pages, start=1):
        text = (page.extract_text() or "").strip().replace("\x00", "")
        if text:
            pages.append({"page_number": i, "text": text})

    return pages


def build_chunks(pages):
    chunks = []

    for page in pages:
        docs = text_splitter.create_documents([page["text"]])
        for i, doc in enumerate(docs):
            chunks.append({
                "page_number": page["page_number"],
                "chunk_index": i,
                "chunk_text": doc.page_content
            })

    return chunks


def embed_texts(texts):
    return embedding_model.embed_documents(texts)


def embed_query(text):
    return embedding_model.embed_query(text)


# ================= INGEST =================

@app.post("/ingest")
def ingest(payload: IngestRequest, x_rag_service_secret: str | None = Header(None)):

    file_bytes = download_file_bytes(payload.bucket, payload.object_key)
    pages = extract_text(file_bytes)
    chunks = build_chunks(pages)

    embeddings = embed_texts([c["chunk_text"] for c in chunks])

    with get_db_connection() as conn:
        with conn.cursor() as cur:

            cur.execute(
                "DELETE FROM rag_chunks WHERE owner_id=%s AND case_id=%s AND file_id=%s",
                (payload.owner_id, payload.case_id, payload.file_id)
            )

            for chunk, emb in zip(chunks, embeddings):
                cur.execute(
                    """
                    INSERT INTO rag_chunks
                    (owner_id, case_id, file_id, file_name,
                     page_number, chunk_index, chunk_text, embedding)
                    VALUES (%s,%s,%s,%s,%s,%s,%s,%s)
                    """,
                    (
                        payload.owner_id,
                        payload.case_id,
                        payload.file_id,
                        payload.file_name,
                        chunk["page_number"],
                        chunk["chunk_index"],
                        chunk["chunk_text"],
                        emb
                    )
                )

        conn.commit()

    return {"indexed_chunks": len(chunks)}


# ================= ANSWER =================

@app.post("/answer")
def answer(payload: AnswerRequest, x_rag_service_secret: str | None = Header(None)):

    query_emb = embed_query(payload.question)

    with get_db_connection() as conn:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT file_name, page_number, chunk_text,
                1-(embedding <=> %s::vector) AS score
                FROM rag_chunks
                WHERE owner_id=%s AND case_id=%s
                ORDER BY embedding <=> %s::vector
                LIMIT %s
                """,
                (
                    query_emb,
                    payload.owner_id,
                    payload.case_id,
                    query_emb,
                    payload.top_k
                )
            )
            rows = cur.fetchall()

    if not rows:
        return {"answer": "لسه مفيش بيانات", "sources": []}

    context = "\n\n".join([r["chunk_text"] for r in rows])

    prompt = f"""
    انت مساعد قانوني. جاوب بالمصري.

    السياق:
    {context}

    السؤال:
    {payload.question}
    """

    response = llm.invoke(prompt)

    return {
        "answer": response.content,
        "sources": rows
    }