Update rag_routerv2.py

rag_routerv2.py

@@ -76,63 +76,68 @@ class QueryTableResponse(BaseModel):
 
 @router.post("/create_table", response_model=CreateTableResponse)
 async def create_embedding_table(
-   user_id: str,
-   files: List[UploadFile] = File(...),
-   table_id: Optional[str] = None,
-   table_name: Optional[str] = None
+  user_id: str,
+  files: List[UploadFile] = File(...),
+  table_id: Optional[str] = None,
+  table_name: Optional[str] = None
 ) -> CreateTableResponse:
-   allowed_extensions = {".pdf", ".docx", ".csv", ".txt", ".md"}
-   for file in files:
-       if not file.filename:
-           raise HTTPException(status_code=400, detail="Invalid filename")
-       if os.path.splitext(file.filename)[1].lower() not in allowed_extensions:
-           raise HTTPException(status_code=400, detail="Unsupported file type")
-
-   table_id = table_id or str(uuid.uuid4())
-   table_name = table_name or f"knowledge-base-{str(uuid.uuid4())[:4]}"
-   
-   directory_path = f"./data/{table_id}"
-   os.makedirs(directory_path, exist_ok=True)
-
-   for file in files:
-       file_path = os.path.join(directory_path, file.filename)
-       with open(file_path, "wb") as buffer:
-           shutil.copyfileobj(file.file, buffer)
-
-   try:
-       vector_store = LanceDBVectorStore(
-           uri="./lancedb/dev",
-           table_name=table_id,
-           mode="overwrite",
-           query_type="hybrid"
-       )
-
-       documents = SimpleDirectoryReader(directory_path).load_data()
-       index = VectorStoreIndex.from_documents(documents, vector_store=vector_store)
-       index.storage_context.persist(persist_dir=f"./lancedb/index/{table_id}")
-
-       db = get_db()
-       db.execute(
-           'INSERT INTO tables (user_id, table_id, table_name) VALUES (?, ?, ?)',
-           (user_id, table_id, table_name)
-       )
-       
-       for file in files:
-           db.execute(
-               'INSERT INTO table_files (table_id, filename, file_path) VALUES (?, ?, ?)',
-               (table_id, file.filename, f"./data/{table_id}/{file.filename}")
-           )
-       db.commit()
-
-       return CreateTableResponse(
-           table_id=table_id,
-           message="Success",
-           status="success", 
-           table_name=table_name
-       )
-
-   except Exception as e:
-       raise HTTPException(status_code=500, detail=str(e))
+  try:
+      db = get_db()
+      table_id = table_id or str(uuid.uuid4())
+      table_name = table_name or f"knowledge-base-{str(uuid.uuid4())[:4]}"
+      
+      # Check if table exists
+      existing = db.execute(
+          'SELECT id FROM tables WHERE user_id = ? AND table_id = ?', 
+          (user_id, table_id)
+      ).fetchone()
+
+      directory_path = f"./data/{table_id}"
+      os.makedirs(directory_path, exist_ok=True)
+
+      for file in files:
+          if not file.filename:
+              raise HTTPException(status_code=400, detail="Invalid filename")
+          if os.path.splitext(file.filename)[1].lower() not in {".pdf", ".docx", ".csv", ".txt", ".md"}:
+              raise HTTPException(status_code=400, detail="Unsupported file type")
+              
+          file_path = os.path.join(directory_path, file.filename)
+          with open(file_path, "wb") as buffer:
+              shutil.copyfileobj(file.file, buffer)
+
+      vector_store = LanceDBVectorStore(
+          uri="./lancedb/dev",
+          table_name=table_id,
+          mode="overwrite",
+          query_type="hybrid"
+      )
+
+      documents = SimpleDirectoryReader(directory_path).load_data()
+      index = VectorStoreIndex.from_documents(documents, vector_store=vector_store)
+      index.storage_context.persist(persist_dir=f"./lancedb/index/{table_id}")
+
+      if not existing:
+          db.execute(
+              'INSERT INTO tables (user_id, table_id, table_name) VALUES (?, ?, ?)',
+              (user_id, table_id, table_name)
+          )
+      
+      for file in files:
+          db.execute(
+              'INSERT OR REPLACE INTO table_files (table_id, filename, file_path) VALUES (?, ?, ?)',
+              (table_id, file.filename, f"./data/{table_id}/{file.filename}")
+          )
+      db.commit()
+
+      return CreateTableResponse(
+          table_id=table_id,
+          message="Success",
+          status="success",
+          table_name=table_name
+      )
+
+  except Exception as e:
+      raise HTTPException(status_code=500, detail=str(e))
 
 
 @router.post("/query_table/{table_id}", response_model=QueryTableResponse)