app.py

app.py

@@ -1,150 +1,172 @@
-"""
-FastAPI Application for Multimodal RAG System
-US Army Medical Research Papers Q&A
-"""
-
-import os
-import logging
-from typing import List, Dict, Optional, Union
-from contextlib import asynccontextmanager
-
-from fastapi import FastAPI, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import FileResponse
-from fastapi.staticfiles import StaticFiles
-from pydantic import BaseModel, Field
-
-# Import from query_index (standalone)
-from query_index import MultimodalRAGSystem
-
-# Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
-logger = logging.getLogger(__name__)
-
-# Global variables
-rag_system: Optional[MultimodalRAGSystem] = None
-
-# Lifecycle management
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    """Initialize and cleanup RAG system"""
-    global rag_system
-    
-    logger.info("Starting RAG system initialization...")
-    try:
-        rag_system = MultimodalRAGSystem()
-        logger.info("RAG system initialized successfully!")
-    except Exception as e:
-        logger.error(f"Error during initialization: {str(e)}")
-        rag_system = None
-    
-    yield
-    
-    logger.info("Shutting down RAG system...")
-    rag_system = None
-
-# Create FastAPI app
-app = FastAPI(
-    title="Multimodal RAG API",
-    description="Q&A system for US Army medical research papers (Text + Images)",
-    version="2.0.0",
-    lifespan=lifespan
-)
-
-# CORS middleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Mount static files
-app.mount("/static", StaticFiles(directory="static"), name="static")
-
-# Mount extracted images
-# This allows the frontend to load images via /extracted_images/filename.jpg
-if os.path.exists("extracted_images"):
-    app.mount("/extracted_images", StaticFiles(directory="extracted_images"), name="images")
-
-# Mount PDF documents
-if os.path.exists("WHEC_Documents"):
-    app.mount("/documents", StaticFiles(directory="WHEC_Documents"), name="documents")
-
-# Pydantic models
-class QueryRequest(BaseModel):
-    question: str = Field(..., min_length=1, max_length=1000, description="Question to ask")
-
-class ImageSource(BaseModel):
-    path: Optional[str]
-    filename: Optional[str]
-    score: Optional[float]
-    page: Optional[Union[str, int]] # could be int or str depending on metadata
-    file: Optional[str]
-    link: Optional[str] = None
-
-class TextSource(BaseModel):
-    text: str
-    score: float
-    page: Optional[Union[str, int]]
-    file: Optional[str]
-    link: Optional[str] = None
-
-class QueryResponse(BaseModel):
-    answer: str
-    images: List[ImageSource]
-    texts: List[TextSource]
-    question: str
-
-class HealthResponse(BaseModel):
-    status: str
-    rag_initialized: bool
-
-# API Endpoints
-
-@app.get("/", tags=["Root"])
-async def root():
-    """Serve the frontend application"""
-    return FileResponse('static/index.html')
-
-@app.get("/health", response_model=HealthResponse, tags=["Health"])
-async def health_check():
-    """Health check endpoint"""
-    return HealthResponse(
-        status="healthy",
-        rag_initialized=rag_system is not None
-    )
-
-@app.post("/query", response_model=QueryResponse, tags=["Query"])
-async def query_rag(request: QueryRequest):
-    """
-    Query the RAG system
-    """
-    if not rag_system:
-        raise HTTPException(
-            status_code=503,
-            detail="RAG system not initialized. Check logs for errors."
-        )
-    
-    try:
-        # Get answer
-        result = rag_system.ask(request.question)
-        
-        return QueryResponse(
-            answer=result['answer'],
-            images=result['images'],
-            texts=result['texts'],
-            question=request.question
-        )
-        
-    except Exception as e:
-        logger.error(f"Error processing query: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Error processing query: {str(e)}")
-
-if __name__ == "__main__":
-    import uvicorn
+"""
+FastAPI Application for Multimodal RAG System
+US Army Medical Research Papers Q&A
+"""
+
+import os
+import logging
+from typing import List, Dict, Optional, Union
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import FileResponse
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel, Field
+
+# Import from query_index (standalone)
+from query_index import MultimodalRAGSystem
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+# Global variables
+rag_system: Optional[MultimodalRAGSystem] = None
+
+# Store last question-answer pair for simple follow-up
+last_qa_context: Optional[str] = None
+
+# Lifecycle management
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Initialize and cleanup RAG system"""
+    global rag_system
+    
+    logger.info("Starting RAG system initialization...")
+    try:
+        rag_system = MultimodalRAGSystem()
+        logger.info("RAG system initialized successfully!")
+    except Exception as e:
+        logger.error(f"Error during initialization: {str(e)}")
+        rag_system = None
+    
+    yield
+    
+    logger.info("Shutting down RAG system...")
+    rag_system = None
+
+# Create FastAPI app
+app = FastAPI(
+    title="Multimodal RAG API",
+    description="Q&A system for US Army medical research papers (Text + Images)",
+    version="2.0.0",
+    lifespan=lifespan
+)
+
+# CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Mount static files
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+# Mount extracted images
+# This allows the frontend to load images via /extracted_images/filename.jpg
+if os.path.exists("extracted_images"):
+    app.mount("/extracted_images", StaticFiles(directory="extracted_images"), name="images")
+
+# Mount PDF documents
+if os.path.exists("WHEC_Documents"):
+    app.mount("/documents", StaticFiles(directory="WHEC_Documents"), name="documents")
+
+# Pydantic models
+class QueryRequest(BaseModel):
+    question: str = Field(..., min_length=1, max_length=1000, description="Question to ask")
+
+class ImageSource(BaseModel):
+    path: Optional[str]
+    filename: Optional[str]
+    score: Optional[float]
+    page: Optional[Union[str, int]] # could be int or str depending on metadata
+    file: Optional[str]
+    link: Optional[str] = None
+
+class TextSource(BaseModel):
+    text: str
+    score: float
+    page: Optional[Union[str, int]]
+    file: Optional[str]
+    link: Optional[str] = None
+
+class QueryResponse(BaseModel):
+    answer: str
+    images: List[ImageSource]
+    texts: List[TextSource]
+    question: str
+
+class HealthResponse(BaseModel):
+    status: str
+    rag_initialized: bool
+
+# API Endpoints
+
+@app.get("/", tags=["Root"])
+async def root():
+    """Serve the frontend application"""
+    return FileResponse('static/index.html')
+
+@app.get("/health", response_model=HealthResponse, tags=["Health"])
+async def health_check():
+    """Health check endpoint"""
+    return HealthResponse(
+        status="healthy",
+        rag_initialized=rag_system is not None
+    )
+
+@app.post("/query", response_model=QueryResponse, tags=["Query"])
+async def query_rag(request: QueryRequest):
+    """
+    Query the RAG system
+    """
+    global last_qa_context
+    
+    if not rag_system:
+        raise HTTPException(
+            status_code=503,
+            detail="RAG system not initialized. Check logs for errors."
+        )
+    
+    try:
+        # Build prompt using previous Q/A if available
+        if last_qa_context:
+            prompt = (
+                f"Previous question and answer:\n"
+                f"{last_qa_context}\n\n"
+                f"Follow up question:\n"
+                    f"{request.question}"
+            )
+        else:
+            prompt = request.question
+
+        # Query RAG system
+        result = rag_system.ask(prompt)
+
+        # Save current Q/A as context for next turn
+        last_qa_context = (
+            f"Question: {request.question}\n"
+            f"Answer: {result['answer']}"
+        )
+        
+        return QueryResponse(
+            answer=result['answer'],
+            images=result['images'],
+            texts=result['texts'],
+            question=request.question
+        )
+        
+    except Exception as e:
+        logger.error(f"Error processing query: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error processing query: {str(e)}")
+
+if __name__ == "__main__":
+    import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)