Hitan2004 commited on
Commit
aec0e75
·
1 Parent(s): aac6cac

Refactor upload endpoint to remove background tasks

Browse files
Files changed (1) hide show
  1. main.py +6 -6
main.py CHANGED
@@ -1,7 +1,7 @@
1
  import os
2
  import shutil
3
  from contextlib import asynccontextmanager
4
- from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
5
  from pydantic import BaseModel
6
  from langchain_core.messages import HumanMessage, AIMessage
7
 
@@ -73,7 +73,8 @@ async def query(req: QueryRequest):
73
 
74
 
75
  @app.post("/upload")
76
- async def upload(background_tasks: BackgroundTasks, file: UploadFile = File(...)):
 
77
  allowed = {".txt", ".pdf"}
78
  ext = os.path.splitext(file.filename or "")[1].lower()
79
  if ext not in allowed:
@@ -84,15 +85,14 @@ async def upload(background_tasks: BackgroundTasks, file: UploadFile = File(...)
84
  with open(dest, "wb") as f:
85
  shutil.copyfileobj(file.file, f)
86
 
87
- background_tasks.add_task(_reindex)
88
  return {"status": "uploaded", "filename": file.filename,
89
- "message": "Indexing started in background."}
90
 
91
 
92
  def _reindex():
93
  try:
94
- from retriever import _model
95
- run_ingestion(model=_model)
96
  reload_indexes()
97
  print("Re-indexing complete.")
98
  except Exception as e:
 
1
  import os
2
  import shutil
3
  from contextlib import asynccontextmanager
4
+ from fastapi import FastAPI, UploadFile, File, HTTPException
5
  from pydantic import BaseModel
6
  from langchain_core.messages import HumanMessage, AIMessage
7
 
 
73
 
74
 
75
  @app.post("/upload")
76
+
77
+ async def upload(file: UploadFile = File(...)):
78
  allowed = {".txt", ".pdf"}
79
  ext = os.path.splitext(file.filename or "")[1].lower()
80
  if ext not in allowed:
 
85
  with open(dest, "wb") as f:
86
  shutil.copyfileobj(file.file, f)
87
 
88
+ _reindex()
89
  return {"status": "uploaded", "filename": file.filename,
90
+ "message": "Indexing complete."}
91
 
92
 
93
  def _reindex():
94
  try:
95
+ run_ingestion()
 
96
  reload_indexes()
97
  print("Re-indexing complete.")
98
  except Exception as e: