import os from fastapi import FastAPI, File, UploadFile, HTTPException, Form, BackgroundTasks, Query, APIRouter from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse from pydantic import BaseModel from typing import List, Optional, Dict, Any import tempfile import uuid import json from datetime import datetime import shutil from pathlib import Path import asyncio # Import Services from core.extractor import process_document from services.ingestion import VectorStoreManager from services.retrieval import VectorStoreRetriever from services.theme import ThemeSynthesizer # Intialize the FastAPI app app = FastAPI( title="Document Research & Theme Identification API", description="API for processing documents, querying content, and identifying themes", version="1.0.0" ) # Add CORS middleware app.add_middleware( CORSMiddleware, allow_origins=["*"], # Allows all origins allow_credentials=True, allow_methods=["*"], allow_headers=["*"] ) # Global variables DATA_DIR = Path("./data") CHROMA_DIR = DATA_DIR / "chroma_db" UPLOAD_DIR = DATA_DIR / "uploads" # Create necessary directories DATA_DIR.mkdir(exist_ok=True) CHROMA_DIR.mkdir(exist_ok=True) UPLOAD_DIR.mkdir(exist_ok=True) # Intialize services vector_store_manager = VectorStoreManager(persist_directory=str(CHROMA_DIR)) vector_store_retriever = VectorStoreRetriever(persist_directory=str(CHROMA_DIR)) # Initialize ThemeSynthesizer with Groq API key groq_api_key = os.getenv("GROQ_API_KEY") if groq_api_key: theme_synthesizer = ThemeSynthesizer(api_key=groq_api_key) else: theme_synthesizer = None print("Warning: GROQ_API_KEY not found. Theme synthesis will not be available.") # Process queue for background tasks processing_queue = {} # Pydantic models for request/response class DocumentResponse(BaseModel): id: str filename: str summary: Optional[str] = None added_st: str status: str metadata: Optional[Dict[str, Any]] = None class QueryRequest(BaseModel): query: str provider: Optional[str] = "groq" model: Optional[str] = "llama3-8b-8192" search_depth: int = 5 include_themes: bool = True theme_threshold: int = 2 class ThemeAnalysis(BaseModel): themes: List[str] status: str agreements: Optional[List[str]] = None contradictions: Optional[List[str]] = None insights: Optional[List[str]] = None analysis: Optional[str] = None class QueryResponse(BaseModel): results: List[Dict[str, Any]] theme_analysis: Optional[ThemeAnalysis] = None query: str total_results: int class LLMProvider(BaseModel): id: str display_name: str models: List[Dict[str, str]] class LLMProvidersResponse(BaseModel): providers: Dict[str, LLMProvider] # Router for document operations document_router = APIRouter(prefix="/api/documents", tags=["Documents"]) @document_router.post("/upload") async def upload_document( background_tasks: BackgroundTasks, file: UploadFile = File(...) ): """Upload a document for processing""" doc_id = str(uuid.uuid4()) # Save file to disk file_path = UPLOAD_DIR / f"{doc_id}_{file.filename}" try: # Create a temporary file with open(file_path, "wb") as buffer: shutil.copyfileobj(file.file, buffer) # Add to processing queue processing_queue[doc_id] = { "id": doc_id, "filename": file.filename, "added_at": datetime.now().isoformat(), "status": "processing" } # Start background processing background_tasks.add_task( process_document_task, str(file_path), file.filename, doc_id ) return { "id": doc_id, "filename": file.filename, "status": "processing", "message": "Document upload successful, processing started" } except Exception as e: if file_path.exists(): file_path.unlink() # Clean up the file if it was created raise HTTPException(status_code=500, detail=f"Error uploading document: {str(e)}") @document_router.get("/", response_model=List[DocumentResponse]) async def get_documents(): """Get all processed documents""" try: # Get documents from vector store vector_docs = vector_store_retriever.search_by_metadata({}, limit=100) # Also check processing queue for documents still processing all_docs = [] # Add documents from the vector store for doc in vector_docs: doc_id = doc.get("metadata", {}).get("doc_id") if not doc_id: continue # Try to load document info from disk doc_info_path = DATA_DIR / f"{doc_id}.json" if doc_info_path.exists(): with open(doc_info_path, "r") as f: doc_info = json.load(f) all_docs.append(DocumentResponse(**doc_info)) else: # Create from vector store data all_docs.append(DocumentResponse( id=doc_id, filename=doc.get("metadata", {}).get("filename", "Unknown"), added_at=doc.get("metadata", {}).get("timestamp", datetime.now().isoformat()), status="completed", metadata=doc.get("metadata") )) # Add documents still in processing queue for doc_id, doc_info in processing_queue.items(): # Skip if already added from vector store if any(d.id == doc_id for d in all_docs): continue all_docs.append(DocumentResponse( id=doc_id, filename=doc_info["filename"], added_at=doc_info["added_at"], status=doc_info["status"], summary=doc_info.get("summary") )) return all_docs except Exception as e: raise HTTPException(status_code=500, detail=f"Error retrieving documents: {str(e)}") @document_router.get("/{doc_id}", response_model=DocumentResponse) async def get_document(doc_id: str): """Get details for a specific document""" try: # Check if document info exists on disk doc_info_path = DATA_DIR / f"{doc_id}.json" if doc_info_path.exists(): with open(doc_info_path, "r") as f: doc_info = json.load(f) return DocumentResponse(**doc_info) # Check processing queue if doc_id in processing_queue: return DocumentResponse( id=doc_id, filename=processing_queue[doc_id]["filename"], added_at=processing_queue[doc_id]["added_at"], status=processing_queue[doc_id]["status"], summary=processing_queue[doc_id].get("summary") ) raise HTTPException(status_code=404, detail=f"Document {doc_id} not found") except HTTPException: raise except Exception as e: raise HTTPException(status_code=500, detail=f"Error retrieving document: {str(e)}") # Router for document extraction extraction_router = APIRouter(prefix="/api", tags=["Extraction"]) @extraction_router.post("/upload-and-extract") async def upload_and_extract(file: UploadFile = File(...)): """Extract text from a document without storing it""" try: temp_dir = tempfile.mkdtemp() try: file_ext = os.path.splitext(file.filename)[1] temp_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}{file_ext}") with open(temp_file_path, "wb") as buffer: buffer.write(await file.read()) # Extract text result = process_document(file, temp_file_path) return JSONResponse(content=result) finally: # Always clean up temp directory shutil.rmtree(temp_dir, ignore_errors=True) except Exception as e: raise HTTPException(status_code=500, detail=f"Error extracting text: {str(e)}") # Router for querying query_router = APIRouter(prefix="/api", tags=["Query"]) @query_router.post("/query", response_model=QueryResponse) async def query_documents(request: QueryRequest): """Query documents and identify themes""" try: # Search vector store search_results = vector_store_retriever.search( query=request.query, k=request.search_depth ) if not search_results: return QueryResponse( results=[], theme_analysis=None, query=request.query, total_results=0 ) # Process each search result with the LLM document_responses = [] for result in search_results: doc_id = result.get("metadata", {}).get("doc_id") doc_name = result.get("metadata", {}).get("filename", "Unknown") text_chunk = result.get("text", "") if not doc_id or not text_chunk: continue # Get response from LLM if theme_synthesizer: response = theme_synthesizer.process_query( query=request.query, context=text_chunk, provider=request.provider, model=request.model ) else: response = "LLM processing not available. Theme synthesizer is not initialized." document_responses.append({ "doc_id": doc_id, "doc_name": doc_name, "response": response, "chunk_id": result.get("id", "unknown_chunk") }) # Identify themes if requested and if we have enough documents theme_analysis = None if request.include_themes and len(document_responses) >= request.theme_threshold and theme_synthesizer: theme_analysis_result = theme_synthesizer.synthesize_themes( query=request.query, document_responses=document_responses, provider=request.provider, model=request.model ) theme_analysis = ThemeAnalysis( themes=theme_analysis_result.get("themes", []), status="completed", agreements=theme_analysis_result.get("agreements", []), contradictions=theme_analysis_result.get("contradictions", []), insights=theme_analysis_result.get("insights", []), analysis=theme_analysis_result.get("analysis", "") ) elif request.include_themes: theme_analysis = ThemeAnalysis( themes=[], status="skipped", analysis="Skipped theme analysis: not enough documents or theme synthesizer not available" ) return QueryResponse( results=document_responses, theme_analysis=theme_analysis, query=request.query, total_results=len(document_responses) ) except Exception as e: raise HTTPException(status_code=500, detail=f"Error processing query: {str(e)}") # Provider API provider_router = APIRouter(prefix="/api", tags = ["LLM"]) @provider_router.get("/providers", response_model=LLMProvidersResponse) async def get_available_providers(): """Get available LLM providers and models""" # Get API keys from environment groq_api_key = os.getenv("GROQ_API_KEY", "") providers = {} if groq_api_key: providers["groq"] = LLMProvider( id="groq", display_name="Groq", models=[ {"id": "llama3-70b-8192", "name": "Llama 3 70B", "description": "Largest Llama 3 model"}, {"id": "llama3-8b-8192", "name": "Llama 3 8B", "description": "Smaller, faster Llama 3 model"}, {"id": "mixtral-8x7b-32768", "name": "Mixtral 8x7B", "description": "Mixtral, competitor to Llama"} ] ) return {"providers": providers} # Health check @app.get("/", tags=["Health"]) async def root(): return {"status": "ok", "message": "Document Research & Theme Identification API is running"} # Background task for document processing async def process_document_task(file_path: str, filename: str, doc_id: str): try: # Extract text from the document # Extract text from the document with open(file_path, "rb") as f: file_content = f.read() # Create a temporary UploadFile-like object mock_file = type('MockFile', (), {'filename': filename}) result = process_document(mock_file, file_path) if not result.get("success", False): raise Exception(f"Failed to extract text: {result.get('error', 'Unknown error')}") text = result.get("text", "") # Create metadata metadata = { "filename": filename, "doc_id": doc_id, "timestamp": datetime.now().isoformat() } # Get a summary if there's enough content summary = "No meaningful content to summarize." if len(text) > 100 and theme_synthesizer: try: summary = theme_synthesizer.analyze_single_document(text) except Exception as e: print(f"Error generating summary: {e}") summary = f"Error generating summary: {e}" # Store document info document_info = { "id": doc_id, "filename": filename, "summary": summary, "added_at": datetime.now().isoformat(), "status": "completed", "metadata": metadata } # Save document info to disk doc_info_path = DATA_DIR / f"{doc_id}.json" with open(doc_info_path, "w") as f: json.dump(document_info, f) # Save the full text separately (could be large) doc_text_path = DATA_DIR / f"{doc_id}.txt" with open(doc_text_path, "w") as f: f.write(text) # Add to vector store vector_store_manager.add_document(text, metadata, doc_id) # Update processing status processing_queue[doc_id]["status"] = "completed" processing_queue[doc_id]["summary"] = summary except Exception as e: # Update processing status with error processing_queue[doc_id]["status"] = "failed" processing_queue[doc_id]["error"] = str(e) print(f"Error processing document {doc_id}: {e}") finally: # Clean up the uploaded file try: Path(file_path).unlink(missing_ok=True) except Exception as e: print(f"Error removing temporary file {file_path}: {e}") # Include routers app.include_router(document_router) app.include_router(extraction_router) app.include_router(query_router) app.include_router(provider_router) if __name__ == "__main__": import uvicorn uvicorn.run("api:app", host="0.0.0.0", port=8000, reload=True)