Auto-commit: OCR workflow improvements, performance optimizations, and bug fixes

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

@@ -219,9 +219,22 @@ def parse_query_mode(query: str) -> tuple[str, SearchMode, bool, Optional[str]]:
 
 
 class OllamaAPI:
-    def __init__(self, rag: LightRAG, top_k: int = 60, api_key: Optional[str] = None):
-        self.rag = rag
-        self.ollama_server_infos = rag.ollama_server_infos
+    def __init__(self, rag_or_manager, top_k: int = 60, api_key: Optional[str] = None):
+        # 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
+            self.workspace_manager = rag_or_manager
+            self.rag = rag_or_manager.get_rag()  # Get default workspace RAG
+        elif isinstance(rag_or_manager, LightRAG):
+            self.workspace_manager = None
+            self.rag = rag_or_manager
+        else:
+            raise TypeError(f"Expected LightRAG or WorkspaceManager, got {type(rag_or_manager)}")
+        
+        self.ollama_server_infos = self.rag.ollama_server_infos
         self.top_k = top_k
         self.api_key = api_key
         self.router = APIRouter(tags=["ollama"])