use sqlite to replace tables.json

rag_routerv2.py

@@ -9,6 +9,7 @@ import uuid
 import io
 from io import BytesIO
 import csv
+import sqlite3
 
 # LlamaIndex imports
 from llama_index.core import Settings, SimpleDirectoryReader, VectorStoreIndex
@@ -27,13 +28,38 @@ router = APIRouter(
 
 # Configure global LlamaIndex settings
 Settings.embed_model = FastEmbedEmbedding(model_name="BAAI/bge-small-en-v1.5")
-tables_file_path = './data/tables.json'
 
 # Database connection dependency
 @lru_cache()
 def get_db_connection(db_path: str = "./lancedb/dev"):
     return lancedb.connect(db_path)
 
+def get_db():
+   conn = sqlite3.connect('./data/tables.db')
+   conn.row_factory = sqlite3.Row
+   return conn
+
+def init_db():
+   db = get_db()
+   db.execute('''
+       CREATE TABLE IF NOT EXISTS tables (
+           id INTEGER PRIMARY KEY,
+           user_id TEXT NOT NULL,
+           table_id TEXT NOT NULL,
+           table_name TEXT NOT NULL
+       )
+   ''')
+   db.execute('''
+       CREATE TABLE IF NOT EXISTS table_files (
+           id INTEGER PRIMARY KEY,
+           table_id TEXT NOT NULL,
+           filename TEXT NOT NULL,
+           file_path TEXT NOT NULL,
+           FOREIGN KEY (table_id) REFERENCES tables (table_id)
+       )
+   ''')
+   db.commit()
+
 # Pydantic models
 class CreateTableResponse(BaseModel):
     table_id: str
@@ -48,94 +74,65 @@ 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:
-    """Create a table and load embeddings from uploaded files using LlamaIndex."""
-    allowed_extensions = {".pdf", ".docx", ".csv", ".txt", ".md"}
-    for file in files:
-        if file.filename is None:
-            raise HTTPException(status_code=400, detail="File must have a valid name.")
-        file_extension = os.path.splitext(file.filename)[1].lower()
-        if file_extension not in allowed_extensions:
-            raise HTTPException(
-                status_code=400, 
-                detail=f"File type {file_extension} is not allowed. Supported file types are: {', '.join(allowed_extensions)}."
-            )
-    
-    if table_id is None:
-        table_id = str(uuid.uuid4())
-    table_name = f"knowledge-base-{str(uuid.uuid4())[:4]}" if not table_name else table_name
-    
-    #table_name = table_id #f"{user_id}__table__{table_id}"
-    
-    # Create a directory for the uploaded files
-    directory_path = f"./data/{table_id}"
-    os.makedirs(directory_path, exist_ok=True)
-
-    # Save each uploaded file to the data directory
-    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:
-        # Setup LanceDB vector store
-        vector_store = LanceDBVectorStore(
-            uri="./lancedb/dev",
-            table_name=table_id,
-            mode="overwrite",
-            query_type="hybrid"
-        )
-
-        # Load documents using SimpleDirectoryReader
-        documents = SimpleDirectoryReader(directory_path).load_data()
-        
-        # Create the index
-        index = VectorStoreIndex.from_documents(
-            documents,
-            vector_store=vector_store
-        )
-        index.storage_context.persist(persist_dir=f"./lancedb/index/{table_id}")
-
-
-        # Store user_id and table_name in a JSON file
-        try:
-            tables_file_path = './data/tables.json'
-            os.makedirs(os.path.dirname(tables_file_path), exist_ok=True)
-            # Load existing tables or create a new file if it doesn't exist
-            try:
-                with open(tables_file_path, 'r') as f:
-                    tables = json.load(f)
-            except (FileNotFoundError, json.JSONDecodeError):
-                tables = {}
-    
-            # Update the tables dictionary
-            if user_id not in tables:
-                tables[user_id] = []
-            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:
-                json.dump(tables, f)
-        
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Failed to update tables file: {str(e)}")
-        
-        return CreateTableResponse(
-            table_id=table_id,
-            message="Table created and documents indexed successfully",
-            status="success",
-            table_name=table_name
-        )
+   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))
 
 
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Table creation failed: {str(e)}")
-
 @router.post("/query_table/{table_id}", response_model=QueryTableResponse)
 async def query_table(
     table_id: str,
@@ -172,22 +169,28 @@ async def query_table(
 
 @router.get("/get_tables/{user_id}")
 async def get_tables(user_id: str):
-    """Get all tables for a user."""
-    
-    tables_file_path = './data/tables.json'
-    try:
-        # Load existing tables from the JSON file
-        with open(tables_file_path, 'r') as f:
-            tables = json.load(f)
-
-        # Retrieve tables for the specified user
-        user_tables = tables.get(user_id, [])
-        return user_tables
-
-    except (FileNotFoundError, json.JSONDecodeError):
-        return []  # Return an empty list if the file doesn't exist or is invalid
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Failed to retrieve tables: {str(e)}")
+   db = get_db()
+   tables = db.execute('''
+       SELECT t.*, GROUP_CONCAT(tf.filename) as filenames, GROUP_CONCAT(tf.file_path) as file_paths
+       FROM tables t
+       LEFT JOIN table_files tf ON t.table_id = tf.table_id
+       WHERE t.user_id = ?
+       GROUP BY t.table_id
+   ''', (user_id,)).fetchall()
+   
+   result = []
+   for table in tables:
+       table_dict = dict(table)
+       table_dict['files'] = [
+           {'filename': f, 'file_path': p} 
+           for f, p in zip(
+               table_dict.pop('filenames').split(',') if table_dict['filenames'] else [],
+               table_dict.pop('file_paths').split(',') if table_dict['file_paths'] else []
+           )
+       ]
+       result.append(table_dict)
+   
+   return result
 
 @router.get("/health")
 async def health_check():
@@ -195,40 +198,42 @@ async def health_check():
 
 @router.on_event("startup")
 async def startup():
-    print("RAG Router started")
-    from llama_index.core.schema import TextNode
-    table_name = "digiyatra"
-    nodes = []
-    vector_store = LanceDBVectorStore(
-        
-        uri="./lancedb/dev",
-        table_name=table_name,
-        mode="overwrite",
-        query_type="hybrid"
-    )
-    # load digiyatra csv and create node for each row using csv.reader
-    with open('combined_digi_yatra.csv', newline='') as f:
-        reader = csv.reader(f)
-        data = list(reader)
-        for row in data[1:]:
-            node = TextNode(text=str(row), id_=str(uuid.uuid4()))
-            nodes.append(node)
-
-    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_id": table_name,
-            "table_name": table_name
-        }
-    ]
-    with open(tables_file_path, 'w') as f:
-        json.dump(tables, f)
+   init_db()
+   print("RAG Router started")
+   
+   table_name = "digiyatra"
+   user_id = "digiyatra"
+   
+   # Create vector store and index
+   vector_store = LanceDBVectorStore(
+       uri="./lancedb/dev",
+       table_name=table_name,
+       mode="overwrite",
+       query_type="hybrid"
+   )
+
+   # Load CSV and create nodes
+   with open('combined_digi_yatra.csv', newline='') as f:
+       nodes = [
+           TextNode(text=str(row), id_=str(uuid.uuid4()))
+           for row in list(csv.reader(f))[1:]
+       ]
+
+   # Create and persist index
+   index = VectorStoreIndex(nodes, vector_store=vector_store)
+   index.storage_context.persist(persist_dir=f"./lancedb/index/{table_name}")
+
+   # Store in SQLite
+   db = get_db()
+   db.execute(
+       'INSERT INTO tables (user_id, table_id, table_name) VALUES (?, ?, ?)',
+       (user_id, table_name, table_name)
+   )
+   db.execute(
+       'INSERT INTO table_files (table_id, filename, file_path) VALUES (?, ?, ?)',
+       (table_name, 'combined_digi_yatra.csv', 'combined_digi_yatra.csv')
+   )
+   db.commit()
 
 @router.on_event("shutdown")
 async def shutdown():