api_wrapper.py

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import asyncio
import os
from app import BorgesGraphRAG

# Create FastAPI app for direct API access
api_app = FastAPI(title="Borges GraphRAG API", version="1.0.0")

# Add CORS middleware
api_app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class QueryRequest(BaseModel):
    query: str
    mode: str = "local"
    book_id: str = None

class QueryResponse(BaseModel):
    success: bool
    answer: str = None
    searchPath: dict = None
    error: str = None
    book_id: str = None
    mode: str = None

# Initialize GraphRAG
borges_rag = BorgesGraphRAG()

# Load available books
available_books = []
for item in os.listdir('.'):
    if os.path.isdir(item) and not item.startswith('.'):
        graph_file = os.path.join(item, 'graph_chunk_entity_relation.graphml')
        if os.path.exists(graph_file):
            available_books.append(item)

if available_books:
    default_book = available_books[0]
    borges_rag.load_book_data(default_book)
    print(f"✅ Livre chargé: {default_book}")

@api_app.get("/")
async def root():
    return {
        "status": "active",
        "service": "Borges GraphRAG API",
        "books_loaded": len(available_books),
        "default_book": available_books[0] if available_books else None
    }

@api_app.post("/query", response_model=QueryResponse)
async def query_graphrag(request: QueryRequest):
    """Direct GraphRAG query endpoint"""
    try:
        if not request.query.strip():
            raise HTTPException(status_code=400, detail="Query cannot be empty")

        # Query the GraphRAG system
        result = await borges_rag.query_book(request.query, request.mode)

        return QueryResponse(**result)

    except Exception as e:
        return QueryResponse(
            success=False,
            error=str(e)
        )

@api_app.get("/books")
async def list_books():
    """List available books"""
    return {
        "books": available_books,
        "count": len(available_books)
    }

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