Spaces:
Sleeping
Sleeping
File size: 4,126 Bytes
9514a77 |
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 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
#!/usr/bin/env python3
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional
import logging
import time
import json
from pathlib import Path
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
STATUS_FILE = Path('/tmp/setup_status.json')
READY_FLAG = Path('/tmp/faiss_ready')
def get_setup_status():
if not STATUS_FILE.exists():
return {'status': 'initializing', 'message': 'Setup não iniciado', 'progress': 0}
try:
with open(STATUS_FILE) as f:
return json.load(f)
except:
return {'status': 'unknown', 'message': 'Erro ao ler status', 'progress': 0}
def is_ready():
return READY_FLAG.exists()
query_engine = None
def get_query_engine():
global query_engine
if query_engine is None:
if not is_ready():
raise HTTPException(status_code=503, detail="RAG em construção. Tente em alguns minutos.")
logger.info("Carregando QueryEngine...")
from query_engine import QueryEngine
query_engine = QueryEngine()
logger.info("✅ QueryEngine carregado!")
return query_engine
app = FastAPI(title="Para.AI RAG Cluster (LangChain)", version="1.0.0")
class EmbeddingSearchRequest(BaseModel):
query: str
top_k: int = 10
return_embeddings: bool = False
class KeywordSearchRequest(BaseModel):
keywords: List[str]
operator: str = "AND"
top_k: int = 20
class IDSearchRequest(BaseModel):
ids: List[str]
return_embeddings: bool = False
@app.get("/")
async def root():
setup_status = get_setup_status()
ready = is_ready()
response = {"status": "online", "rag_ready": ready, "setup": setup_status, "backend": "LangChain + FAISS (CPU)"}
if ready and query_engine:
response["cluster_id"] = query_engine.config.get('cluster_id')
response["chunk_range"] = [query_engine.config.get('chunk_start'), query_engine.config.get('chunk_end')]
return response
@app.get("/setup/status")
async def setup_status():
return get_setup_status()
@app.get("/health")
async def health():
return {"status": "ok", "timestamp": time.time()}
@app.post("/search/embedding")
async def search_embedding(request: EmbeddingSearchRequest):
engine = get_query_engine()
try:
start = time.time()
results = engine.search_by_embedding(request.query, request.top_k, request.return_embeddings)
results['query_time_ms'] = round((time.time() - start) * 1000, 2)
return results
except Exception as e:
logger.error(f"Erro: {e}")
raise HTTPException(status_code=500, detail=str(e))
@app.post("/search/keywords")
async def search_keywords(request: KeywordSearchRequest):
engine = get_query_engine()
try:
start = time.time()
results = engine.search_by_keywords(request.keywords, request.operator, request.top_k)
results['query_time_ms'] = round((time.time() - start) * 1000, 2)
return results
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/search/by_id")
async def search_by_id(request: IDSearchRequest):
engine = get_query_engine()
try:
start = time.time()
results = engine.search_by_ids(request.ids, request.return_embeddings)
results['query_time_ms'] = round((time.time() - start) * 1000, 2)
return results
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/cluster/info")
async def cluster_info():
engine = get_query_engine()
try:
info = engine.get_cluster_info()
info['uptime_seconds'] = round(time.time() - app.state.start_time, 2)
return info
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.on_event("startup")
async def startup_event():
app.state.start_time = time.time()
logger.info("="*80)
logger.info("🚀 Para.AI RAG (LangChain + FAISS) ONLINE")
logger.info("="*80)
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)
|