workspace working

LightRAG-main/lightrag/api/routers/query_routes.py

@@ -153,21 +153,41 @@ def create_query_routes(rag_or_manager, api_key: Optional[str] = None, top_k: in
     # Accept either a LightRAG instance or a WorkspaceManager
     from lightrag.api.workspace_manager import WorkspaceManager
     from lightrag import LightRAG
-    
-    if isinstance(rag_or_manager, WorkspaceManager):
-        # Get default RAG instance from workspace manager
-        rag = rag_or_manager.get_rag()
-    elif isinstance(rag_or_manager, LightRAG):
-        rag = rag_or_manager
-    else:
-        raise TypeError(f"Expected LightRAG or WorkspaceManager, got {type(rag_or_manager)}")
+    from fastapi import Request, Depends, HTTPException
+    from lightrag.base import StoragesStatus
+    from lightrag.utils import logger
     
     combined_auth = get_combined_auth_dependency(api_key)
+    
+    # Define dependency to get workspace-specific RAG instance
+    async def get_workspace_rag(request: Request):
+        if isinstance(rag_or_manager, WorkspaceManager):
+            workspace = request.headers.get("X-Workspace", "").strip()
+            rag = rag_or_manager.get_rag(workspace)
+            # Ensure storages are initialized for this workspace
+            try:
+                logger.info(f"Workspace '{workspace}': storages status = {rag._storages_status}")
+                if rag._storages_status != StoragesStatus.FINALIZED:
+                    logger.info(f"Initializing storages for workspace '{workspace}'")
+                    await rag.initialize_storages()
+                    logger.info(f"Storages initialized, status now = {rag._storages_status}")
+            except Exception as e:
+                logger.error(f"Failed to initialize storages for workspace '{workspace}': {e}")
+                raise HTTPException(status_code=500, detail=f"Storage initialization failed: {e}")
+            return rag
+        elif isinstance(rag_or_manager, LightRAG):
+            # Single RAG instance mode - ignore workspace header
+            return rag_or_manager
+        else:
+            raise TypeError(f"Expected LightRAG or WorkspaceManager, got {type(rag_or_manager)}")
 
     @router.post(
         "/query", response_model=QueryResponse, dependencies=[Depends(combined_auth)]
     )
-    async def query_text(request: QueryRequest):
+    async def query_text(
+        request: QueryRequest,
+        rag: LightRAG = Depends(get_workspace_rag)
+    ):
         """
         Handle a POST request at the /query endpoint to process user queries using RAG capabilities.
 
@@ -200,7 +220,10 @@ def create_query_routes(rag_or_manager, api_key: Optional[str] = None, top_k: in
             raise HTTPException(status_code=500, detail=str(e))
 
     @router.post("/query/stream", dependencies=[Depends(combined_auth)])
-    async def query_text_stream(request: QueryRequest):
+    async def query_text_stream(
+        request: QueryRequest,
+        rag: LightRAG = Depends(get_workspace_rag)
+    ):
         """
         This endpoint performs a retrieval-augmented generation (RAG) query and streams the response.
 
@@ -253,7 +276,10 @@ def create_query_routes(rag_or_manager, api_key: Optional[str] = None, top_k: in
         response_model=QueryDataResponse,
         dependencies=[Depends(combined_auth)],
     )
-    async def query_data(request: QueryRequest):
+    async def query_data(
+        request: QueryRequest,
+        rag: LightRAG = Depends(get_workspace_rag)
+    ):
         """
         Retrieve structured data without LLM generation.