Update rag_routerv2.py

rag_routerv2.py

@@ -12,9 +12,9 @@ import csv
 
 # LlamaIndex imports
 from llama_index.core import Settings, SimpleDirectoryReader, VectorStoreIndex
-from llama_index.core.schema import TextNode
 from llama_index.vector_stores.lancedb import LanceDBVectorStore
 from llama_index.embeddings.fastembed import FastEmbedEmbedding
+from llama_index.core.schema import TextNode
 from llama_index.core import StorageContext, load_index_from_storage
 import json
 import os
@@ -39,6 +39,7 @@ class CreateTableResponse(BaseModel):
     table_id: str
     message: str
     status: str
+    table_name: str
 
 class QueryTableResponse(BaseModel):
     results: Dict[str, Any]
@@ -49,7 +50,8 @@ class QueryTableResponse(BaseModel):
 async def create_embedding_table(
     user_id: str,
     files: List[UploadFile] = File(...),
-    table_id: Optional[str] = None
+    table_id: Optional[str] = None,
+    table_name: Optional[str] = None
 ) -> CreateTableResponse:
     """Create a table and load embeddings from uploaded files using LlamaIndex."""
     allowed_extensions = {".pdf", ".docx", ".csv", ".txt", ".md"}
@@ -65,10 +67,10 @@ async def create_embedding_table(
     
     if table_id is None:
         table_id = str(uuid.uuid4())
-    table_name = table_id #f"{user_id}__table__{table_id}"
+    #table_name = table_id #f"{user_id}__table__{table_id}"
     
     # Create a directory for the uploaded files
-    directory_path = f"./data/{table_name}"
+    directory_path = f"./data/{table_id}"
     os.makedirs(directory_path, exist_ok=True)
 
     # Save each uploaded file to the data directory
@@ -91,8 +93,8 @@ async def create_embedding_table(
         # Update the tables dictionary
         if user_id not in tables:
             tables[user_id] = []
-        if table_name not in tables[user_id]:
-            tables[user_id].append(table_name)
+        if table_id not in [table['table_id'] for table in tables[user_id]]:
+            tables[user_id].append({"table_id": table_id, "table_name": table_name})
 
         # Write the updated tables back to the JSON file
         with open(tables_file_path, 'w') as f:
@@ -104,7 +106,7 @@ async def create_embedding_table(
         # Setup LanceDB vector store
         vector_store = LanceDBVectorStore(
             uri="./lancedb/dev",
-            table_name=table_name,
+            table_name=table_id,
             mode="overwrite",
             query_type="hybrid"
         )
@@ -117,12 +119,13 @@ async def create_embedding_table(
             documents,
             vector_store=vector_store
         )
-        index.storage_context.persist(persist_dir=f"./lancedb/index/{table_name}")
+        index.storage_context.persist(persist_dir=f"./lancedb/index/{table_id}")
 
         return CreateTableResponse(
             table_id=table_id,
             message=f"Table created and documents indexed successfully",
-            status="success"
+            status="success",
+            table_name="" if table_name is None else table_name
         )
 
     except Exception as e:
@@ -181,6 +184,10 @@ async def get_tables(user_id: str):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Failed to retrieve tables: {str(e)}")
 
+@router.get("/health")
+async def health_check():
+    return {"status": "healthy"}
+
 @router.on_event("startup")
 async def startup():
     print("RAG Router started")
@@ -205,16 +212,19 @@ async def startup():
     index = VectorStoreIndex(nodes, vector_store=vector_store)
     index.storage_context.persist(persist_dir=f"./lancedb/index/{table_name}")
 
-
     # Create tables dictionary
     tables = {}
     user_id = "digiyatra"
-    
-    tables[user_id] = [table_name]
+
+    tables[user_id] = [
+        {
+            "table_id": table_name,
+            "table_name": table_name
+        }
+    ]
     with open(tables_file_path, 'w') as f:
         json.dump(tables, f)
 
-
-@router.get("/health")
-async def health_check():
-    return {"status": "healthy"}
+@router.on_event("shutdown")
+async def shutdown():
+    print("RAG Router shutdown")