Update main.py

main.py

@@ -13,7 +13,7 @@ import os
 import time
 import asyncio
 from typing import List, Optional, Dict
-from fastapi import FastAPI, UploadFile, File, HTTPException, Depends, BackgroundTasks,Header
+from fastapi import FastAPI, UploadFile, File, HTTPException, Depends, BackgroundTasks, Header
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from threading import Lock
@@ -124,6 +124,7 @@ class UploadResponse(BaseModel):
 class QueryRequest(BaseModel):
     question: str
     top_k: Optional[int] = 5
+    doc_id: Optional[str] = None
 
 
 class QueryResponse(BaseModel):
@@ -271,7 +272,8 @@ async def query_documents(
         result = await asyncio.to_thread(
             engine.ask,
             query=request.question,
-            top_k=request.top_k
+            top_k=request.top_k,
+            doc_id=request.doc_id
         )
         
         print(f"✅ Query success for user {user_id}")