Spaces:
Sleeping
Sleeping
File size: 3,268 Bytes
4722db8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from app.config import get_settings
from app.core.embeddings import EmbeddingGenerator
from app.core.vector_store import VectorStore
from app.core.llm import OllamaLLM
from app.services.document_processor import DocumentProcessor
from app.services.rag_chain import RAGChain
from app.models.schemas import QueryRequest, QueryResponse, IngestResponse
import tempfile
import os
# Initialize FastAPI app
app = FastAPI(
title="RAG Portfolio Project",
description="Production-grade Retrieval-Augmented Generation system",
version="1.0.0"
)
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Initialize components
settings = get_settings()
embedding_generator = EmbeddingGenerator(settings.embedding_model)
vector_store = VectorStore(
host=settings.qdrant_host,
port=settings.qdrant_port,
collection_name=settings.qdrant_collection_name,
vector_size=embedding_generator.dimension
)
llm = OllamaLLM(settings.ollama_base_url, settings.ollama_model)
document_processor = DocumentProcessor()
rag_chain = RAGChain(embedding_generator, vector_store, llm)
@app.get("/")
async def root():
return {
"message": "RAG Portfolio Project API",
"status": "running",
"docs": "/docs"
}
@app.get("/health")
async def health_check():
return {"status": "healthy", "ollama_connected": True}
@app.post("/ingest/file", response_model=IngestResponse)
async def ingest_file(file: UploadFile = File(...)):
"""Upload and ingest a document into the RAG system"""
try:
# Save uploaded file temporarily
with tempfile.NamedTemporaryFile(delete=False, suffix=file.filename) as tmp:
content = await file.read()
tmp.write(content)
tmp_path = tmp.name
# Process document
chunks = document_processor.process_document(tmp_path)
# Ingest into RAG system
result = rag_chain.ingest_documents(chunks)
# Clean up
os.unlink(tmp_path)
return IngestResponse(**result, message=f"Successfully ingested {file.filename}")
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/query", response_model=QueryResponse)
async def query(request: QueryRequest):
"""Query the RAG system"""
try:
result = rag_chain.query(request.question, request.top_k)
return QueryResponse(**result)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.delete("/reset")
async def reset_collection():
"""Reset the vector collection (delete all documents)"""
try:
vector_store.client.delete_collection(settings.qdrant_collection_name)
vector_store._ensure_collection()
return {"status": "success", "message": "Collection reset successfully"}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host=settings.app_host, port=settings.app_port)
|