python -m venv .venv
source .venv/bin/activate
python -m pip install "graphrag>=3.1,<4" fastapi "uvicorn[standard]" pandas pyarrow

mkdir rag
graphrag init --root ./rag

graphrag index --root ./rag

import asyncio
from contextlib import asynccontextmanager
from pathlib import Path

import graphrag.api as graphrag_api
import pandas as pd
from fastapi import FastAPI, HTTPException, status
from graphrag.config.load_config import load_config
from pydantic import BaseModel, Field


PROJECT_ROOT = Path("./rag").resolve()
OUTPUT_DIR = PROJECT_ROOT / "output"


def load_index_tables() -> dict[str, pd.DataFrame]:
    required = {
        "entities": OUTPUT_DIR / "entities.parquet",
        "communities": OUTPUT_DIR / "communities.parquet",
        "community_reports": OUTPUT_DIR / "community_reports.parquet",
    }
    missing = [str(path) for path in required.values() if not path.exists()]
    if missing:
        raise FileNotFoundError(f"Missing GraphRAG index files: {', '.join(missing)}")
    return {name: pd.read_parquet(path) for name, path in required.items()}


class QueryRequest(BaseModel):
    query: str = Field(min_length=1, max_length=4_000)


index_state = {
    "status": "idle",
    "error": None,
}


@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.config = load_config(PROJECT_ROOT)
    app.state.index_lock = asyncio.Lock()
    try:
        app.state.tables = load_index_tables()
    except FileNotFoundError:
        app.state.tables = None
    yield


app = FastAPI(title="GraphRAG API", lifespan=lifespan)


async def rebuild_index() -> None:
    async with app.state.index_lock:
        index_state.update(status="running", error=None)
        try:
            results = await graphrag_api.build_index(config=app.state.config)
            errors = [error for result in results for error in result.errors]
            if errors:
                raise RuntimeError("; ".join(str(error) for error in errors))
            app.state.tables = load_index_tables()
            index_state["status"] = "complete"
        except Exception as error:
            index_state.update(status="failed", error=str(error))


@app.post("/index", status_code=status.HTTP_202_ACCEPTED)
async def start_index() -> dict[str, str]:
    if index_state["status"] == "running":
        raise HTTPException(status_code=409, detail="Indexing is already running")
    index_state.update(status="running", error=None)
    asyncio.create_task(rebuild_index())
    return {"status": "running"}


@app.get("/index/status")
async def get_index_status() -> dict[str, str | None]:
    return index_state


@app.post("/query")
async def query_index(request: QueryRequest) -> dict[str, str]:
    tables = app.state.tables
    if tables is None:
        raise HTTPException(status_code=409, detail="Build an index before querying")

    response, _context = await graphrag_api.global_search(
        config=app.state.config,
        entities=tables["entities"],
        communities=tables["communities"],
        community_reports=tables["community_reports"],
        community_level=2,
        dynamic_community_selection=False,
        response_type="Multiple Paragraphs",
        query=request.query,
    )
    return {"response": str(response)}

uvicorn api:app --host 127.0.0.1 --port 3000

curl -X POST http://127.0.0.1:3000/index

curl http://127.0.0.1:3000/index/status

curl -X POST http://127.0.0.1:3000/query \
  -H "Content-Type: application/json" \
  -d '{"query":"What are the main themes in these documents?"}'