Chandima Prabhath commited on
Commit
cf4d43c
·
1 Parent(s): 8bf445f

Add upload endpoint for document ingestion and processing

Browse files
Files changed (1) hide show
  1. src/query_service/api.py +33 -1
src/query_service/api.py CHANGED
@@ -1,5 +1,5 @@
1
  # src/query_service/api.py
2
- from fastapi import FastAPI, HTTPException
3
  from fastapi.middleware.cors import CORSMiddleware # Import CORSMiddleware
4
  from pydantic import BaseModel
5
  from src.retrieval_handler.retriever import RetrievalHandler
@@ -9,6 +9,8 @@ from src.vector_store_manager.chroma_manager import ChromaManager
9
  import logging
10
  from typing import Literal, Optional, Dict, Any, List # Import List
11
  from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
 
 
12
 
13
  logger = logging.getLogger(__name__)
14
 
@@ -156,6 +158,36 @@ async def generate_chat_title(request: TitleRequest):
156
  logger.error(f"Title generation error: {e}")
157
  return {"title": "New Chat"}
158
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
159
  # You can add more endpoints here, e.g., /health for health checks
160
  # @app.get("/health")
161
  # async def health_check():
 
1
  # src/query_service/api.py
2
+ from fastapi import FastAPI, HTTPException, UploadFile, File
3
  from fastapi.middleware.cors import CORSMiddleware # Import CORSMiddleware
4
  from pydantic import BaseModel
5
  from src.retrieval_handler.retriever import RetrievalHandler
 
9
  import logging
10
  from typing import Literal, Optional, Dict, Any, List # Import List
11
  from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
12
+ import shutil
13
+ import uuid
14
 
15
  logger = logging.getLogger(__name__)
16
 
 
158
  logger.error(f"Title generation error: {e}")
159
  return {"title": "New Chat"}
160
 
161
+ @app.post("/upload-docs")
162
+ async def upload_docs(files: list[UploadFile] = File(...)):
163
+ """
164
+ Upload new documents and trigger ingestion for them.
165
+ """
166
+ import os
167
+ from src.ingestion_orchestrator.orchestrator import IngestionOrchestrator
168
+
169
+ # Create a unique folder in /tmp
170
+ upload_id = str(uuid.uuid4())
171
+ upload_dir = f"/tmp/ingest_{upload_id}"
172
+ os.makedirs(upload_dir, exist_ok=True)
173
+
174
+ saved_files = []
175
+ for file in files:
176
+ file_path = os.path.join(upload_dir, file.filename)
177
+ with open(file_path, "wb") as buffer:
178
+ shutil.copyfileobj(file.file, buffer)
179
+ saved_files.append(file.filename)
180
+
181
+ # Run the ingestion pipeline for the uploaded folder
182
+ try:
183
+ orchestrator = IngestionOrchestrator()
184
+ orchestrator.run_ingestion_pipeline(docs_folder=upload_dir)
185
+ logger.info(f"Ingested files: {saved_files}")
186
+ return {"status": "success", "files": saved_files}
187
+ except Exception as e:
188
+ logger.error(f"Ingestion failed: {e}")
189
+ raise HTTPException(status_code=500, detail="Ingestion failed.")
190
+
191
  # You can add more endpoints here, e.g., /health for health checks
192
  # @app.get("/health")
193
  # async def health_check():