fastest_api.py

"""
This module defines the FastAPI application for the RAG (Retrieval-Augmented Generation) system.
It provides endpoints for document upload, querying, database operations, and content retrieval.
"""

import os
import tempfile

from fastapi import FastAPI, File, UploadFile
from pydantic import BaseModel

from rag_system import RAGSystem

app = FastAPI()
# Initialize the RAG system with PostgreSQL connection details
rag_system = RAGSystem(postgres_host='postgres', postgres_port=5432)

class QueryModel(BaseModel):
    text: str

@app.post("/upsert")
async def upsert_document(file: UploadFile = File(...)):
    """
    Endpoint to upload and process a document.
    
    Args:
        file (UploadFile): The PDF file to be uploaded and processed.
    
    Returns:
        dict: A message indicating the success of the operation.
    """
    with tempfile.NamedTemporaryFile(delete=False) as tmp:
        tmp.write(await file.read())
        tmp_path = tmp.name

    try:
        result = rag_system.upsert_document(tmp_path)
        return result
    finally:
        os.unlink(tmp_path)

@app.post("/query")
async def query(query: QueryModel):
    """
    Endpoint to query the RAG system.
    
    Args:
        query (QueryModel): The query text.
    
    Returns:
        dict: The answer to the query and related information.
    """
    return rag_system.query(query.text)

@app.post("/clear_db")
async def clear_db():
    """
    Endpoint to clear the database. Utility endpoint for testing.
    
    Returns:
        dict: A message indicating the success of the operation.
    """
    return rag_system.clear_db()

@app.get("/print_db")
async def print_db():
    """
    Endpoint to retrieve and print the contents of the database.
    Utility endpoint for testing.
    
    Returns:
        dict: The contents of the database.
    """
    return rag_system.print_db_contents()

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)