Spaces:
Sleeping
Sleeping
| import json | |
| from pathlib import Path | |
| from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request | |
| from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from pydantic import BaseModel, HttpUrl | |
| from .rag import RAGManager | |
| # Initialize FastAPI app | |
| app = FastAPI( | |
| title="Premium RAG Application", | |
| description="FastAPI + LangChain + OpenRouter owl-alpha RAG Web Interface", | |
| version="1.0.0" | |
| ) | |
| # Enable CORS for local development | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| # Initialize RAG Manager | |
| try: | |
| rag_manager = RAGManager() | |
| except Exception as e: | |
| print(f"Error initializing RAG Manager: {e}") | |
| rag_manager = None | |
| # Paths | |
| BASE_DIR = Path(__file__).resolve().parent | |
| TEMPLATES_DIR = BASE_DIR / "templates" | |
| TEMPLATES_DIR.mkdir(exist_ok=True) | |
| # Pydantic schemas for requests | |
| class URLUploadRequest(BaseModel): | |
| url: str | |
| chunk_size: int = 800 | |
| chunk_overlap: int = 100 | |
| class TextUploadRequest(BaseModel): | |
| text: str | |
| source_name: str | |
| chunk_size: int = 800 | |
| chunk_overlap: int = 100 | |
| class ChatRequest(BaseModel): | |
| query: str | |
| system_prompt: str = None | |
| temperature: float = 0.7 | |
| k: int = 4 | |
| # --- API Endpoints ----------------------------------------------------------- | |
| async def index(): | |
| """Serves the main web UI page.""" | |
| index_path = TEMPLATES_DIR / "index.html" | |
| if not index_path.exists(): | |
| raise HTTPException(status_code=404, detail="index.html templates not found.") | |
| return index_path.read_text(encoding="utf-8") | |
| async def status(): | |
| """Checks model availability, configuration details, and returns list of documents.""" | |
| if rag_manager is None: | |
| return JSONResponse( | |
| status_code=500, | |
| content={ | |
| "ok": False, | |
| "error": "RAG Manager failed to initialize. Please check if OPENROUTER_API_KEY is configured in your .env file." | |
| } | |
| ) | |
| try: | |
| sources = rag_manager.get_indexed_sources() | |
| return { | |
| "ok": True, | |
| "model": rag_manager.model_name, | |
| "embedding_model": rag_manager.embedding_model, | |
| "document_count": len(sources), | |
| "documents": sources | |
| } | |
| except Exception as e: | |
| return JSONResponse( | |
| status_code=500, | |
| content={"ok": False, "error": f"Failed to fetch status: {str(e)}"} | |
| ) | |
| async def upload_file( | |
| file: UploadFile = File(...), | |
| chunk_size: int = Form(800), | |
| chunk_overlap: int = Form(100) | |
| ): | |
| """Handles file uploads (TXT, MD, PDF), splits the text, embeds, and saves to database.""" | |
| if rag_manager is None: | |
| raise HTTPException(status_code=500, detail="RAG Manager not initialized.") | |
| try: | |
| content = await file.read() | |
| result = rag_manager.add_file(content, file.filename, chunk_size, chunk_overlap) | |
| return {"ok": True, "message": f"Successfully indexed {file.filename}.", "data": result} | |
| except Exception as e: | |
| raise HTTPException(status_code=400, detail=f"File upload and indexing failed: {str(e)}") | |
| async def upload_url(request: URLUploadRequest): | |
| """Scrapes a URL, extracts visible text content, embeds, and saves to database.""" | |
| if rag_manager is None: | |
| raise HTTPException(status_code=500, detail="RAG Manager not initialized.") | |
| try: | |
| result = await rag_manager.add_url(request.url, request.chunk_size, request.chunk_overlap) | |
| return {"ok": True, "message": f"Successfully scraped and indexed {request.url}.", "data": result} | |
| except Exception as e: | |
| raise HTTPException(status_code=400, detail=f"URL scraping and indexing failed: {str(e)}") | |
| async def upload_text(request: TextUploadRequest): | |
| """Indexes directly pasted raw text content.""" | |
| if rag_manager is None: | |
| raise HTTPException(status_code=500, detail="RAG Manager not initialized.") | |
| try: | |
| result = rag_manager.add_text_content( | |
| request.text, | |
| request.source_name, | |
| request.chunk_size, | |
| request.chunk_overlap | |
| ) | |
| return {"ok": True, "message": f"Successfully indexed raw text source '{request.source_name}'.", "data": result} | |
| except Exception as e: | |
| raise HTTPException(status_code=400, detail=f"Text indexing failed: {str(e)}") | |
| async def reset_database(): | |
| """Clears all indexed vectors and document caches.""" | |
| if rag_manager is None: | |
| raise HTTPException(status_code=500, detail="RAG Manager not initialized.") | |
| try: | |
| rag_manager.clear_database() | |
| return {"ok": True, "message": "Database and local serialized files cleared successfully."} | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Reset failed: {str(e)}") | |
| async def chat(request: ChatRequest): | |
| """Streams OpenRouter RAG answers with citations back via Server-Sent Events (SSE).""" | |
| if rag_manager is None: | |
| raise HTTPException(status_code=500, detail="RAG Manager not initialized.") | |
| async def event_stream(): | |
| try: | |
| async for event_type, data in rag_manager.stream_chat( | |
| request.query, | |
| request.system_prompt, | |
| request.temperature, | |
| request.k | |
| ): | |
| yield f"data: {json.dumps({'event': event_type, 'data': data}, ensure_ascii=False)}\n\n" | |
| except Exception as e: | |
| yield f"data: {json.dumps({'event': 'error', 'data': str(e)}, ensure_ascii=False)}\n\n" | |
| return StreamingResponse( | |
| event_stream(), | |
| media_type="text/event-stream", | |
| headers={ | |
| "Cache-Control": "no-cache", | |
| "X-Accel-Buffering": "no", # Disables proxy buffering | |
| "Connection": "keep-alive", | |
| } | |
| ) | |