Spaces:
Sleeping
Sleeping
| #!/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 | |
| 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 | |
| async def setup_status(): | |
| return get_setup_status() | |
| async def health(): | |
| return {"status": "ok", "timestamp": time.time()} | |
| 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)) | |
| 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)) | |
| 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)) | |
| 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)) | |
| 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) | |