workspace working

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

@@ -72,21 +72,41 @@ def create_search_routes(rag_or_manager, api_key: Optional[str] = None, top_k: i
     # 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
+    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(
         "/search", response_model=SearchResponse, dependencies=[Depends(combined_auth)]
     )
-    async def search_documents(request: SearchRequest):
+    async def search_documents(
+        request: SearchRequest,
+        rag: LightRAG = Depends(get_workspace_rag)
+    ):
         """
         Handle a POST request at the /search endpoint to search documents using RAG capabilities.
         This is an alias for the /query endpoint with simplified response format.
@@ -169,11 +189,14 @@ def create_search_routes(rag_or_manager, api_key: Optional[str] = None, top_k: i
             raise HTTPException(status_code=500, detail=str(e))
 
     @router.post(
-        "/api/search", 
-        response_model=SearchDataResponse, 
+        "/api/search",
+        response_model=SearchDataResponse,
         dependencies=[Depends(combined_auth)]
     )
-    async def search_data(request: SearchRequest):
+    async def search_data(
+        request: SearchRequest,
+        rag: LightRAG = Depends(get_workspace_rag)
+    ):
         """
         API search endpoint that returns structured data without LLM generation.
         This is an alias for the /query/data endpoint.