Spaces:
Running
Running
Chandima Prabhath
commited on
Commit
·
cf4d43c
1
Parent(s):
8bf445f
Add upload endpoint for document ingestion and processing
Browse files- 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():
|