Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/model_manager/providers/ollama.py: 23%

1"""Ollama model manager implementation.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6from typing import Any 

7 

8ollama: Any = None 

9try: 

10 import ollama as ollama_module 

11 

12 ollama = ollama_module 

13except ImportError: 

14 ollama = None 

15 

16from lexigram.ai.llm.model_manager.base import AbstractModelManager 

17from lexigram.ai.llm.model_manager.types import ModelLoadResult 

18from lexigram.logging import ( 

19 get_logger, 

20) 

21 

22logger = get_logger(__name__) 

23 

24 

25class OllamaModelManager(AbstractModelManager): 

26 """Model manager for Ollama.""" 

27 

28 def __init__(self, base_url: str = "http://localhost:11434"): 

29 super().__init__(base_url, "ollama-model-manager") 

30 self.client = ollama.Client(host=self.base_url) if ollama is not None else None 

31 

32 async def list_models(self) -> list[dict[str, Any]]: 

33 """List available models in Ollama.""" 

34 if self.client is not None: 

35 response = await asyncio.to_thread(self.client.list) 

36 return [{"name": m["name"], "details": m} for m in response["models"]] 

37 

38 client = await self._get_client() 

39 response = await client.get("/api/tags") 

40 response.raise_for_status() 

41 data = response.json() 

42 if asyncio.iscoroutine(data): 

43 data = await data 

44 models = data.get("models", []) if isinstance(data, dict) else [] 

45 return [{"name": m.get("name", ""), "details": m} for m in models] 

46 

47 async def load_model(self, model_name: str, **kwargs: Any) -> ModelLoadResult: 

48 """Load a model in Ollama (pull if not available).""" 

49 try: 

50 # Check if model exists 

51 if self.client is not None: 

52 models = await asyncio.to_thread(self.client.list) 

53 model_names = [m["name"] for m in models["models"]] 

54 

55 if model_name not in model_names: 

56 logger.info("Model %s not found, pulling...", model_name) 

57 await asyncio.to_thread(self.client.pull, model_name) 

58 logger.info("Successfully pulled model %s", model_name) 

59 

60 # Load model by making a simple generate request 

61 await asyncio.to_thread( 

62 self.client.generate, 

63 model=model_name, 

64 prompt="test", 

65 options={"num_predict": 1}, 

66 ) 

67 else: 

68 client = await self._get_client() 

69 pull_resp = await client.post("/api/pull", json={"name": model_name}) 

70 pull_resp.raise_for_status() 

71 gen_resp = await client.post( 

72 "/api/generate", 

73 json={"model": model_name, "prompt": "test", "stream": False}, 

74 ) 

75 gen_resp.raise_for_status() 

76 logger.info("Successfully loaded model %s", model_name) 

77 return ModelLoadResult(success=True, model_name=model_name) 

78 

79 except TimeoutError: 

80 return ModelLoadResult( 

81 success=False, 

82 model_name=model_name, 

83 error=f"Timeout loading model {model_name}", 

84 retryable=True, 

85 ) 

86 except (ConnectionError, OSError) as e: 

87 return ModelLoadResult( 

88 success=False, 

89 model_name=model_name, 

90 error=f"Network error loading model {model_name}: {e}", 

91 retryable=True, 

92 ) 

93 except Exception as e: 

94 logger.exception("Unexpected error loading model %s", model_name) 

95 return ModelLoadResult( 

96 success=False, 

97 model_name=model_name, 

98 error=f"Unexpected error: {e}", 

99 retryable=False, 

100 ) 

101 

102 async def unload_model(self, model_name: str) -> bool: 

103 """Unload a model from Ollama memory.""" 

104 # Ollama doesn't have explicit unload API 

105 # Models stay loaded until server restart or memory pressure 

106 # For our "1 model at a time" policy, we just log that it will be unloaded when another is loaded 

107 logger.info( 

108 "Ollama model %s will be unloaded automatically when memory is needed or another model is loaded", 

109 model_name, 

110 ) 

111 return True 

112 

113 async def switch_model(self, model_name: str, **kwargs: Any) -> ModelLoadResult: 

114 """Switch to a different model (unload others first for 1-at-a-time policy).""" 

115 # Get currently loaded models 

116 loaded = await self.get_loaded_models() 

117 # Unload all except the target 

118 for m in loaded: 

119 if m != model_name: 

120 await self.unload_model(m) 

121 

122 return await self.load_model(model_name, **kwargs) 

123 

124 async def get_loaded_models(self) -> list[str]: 

125 """Get currently loaded models (Ollama tracks this via ps).""" 

126 if self.client is not None: 

127 response = await asyncio.to_thread(self.client.ps) 

128 # If the client returns a coroutine for ps(), await it 

129 if asyncio.iscoroutine(response): 

130 response = await response 

131 models = response["models"] 

132 return [m["name"] for m in models] 

133 

134 client = await self._get_client() 

135 response = await client.get("/api/ps") 

136 response.raise_for_status() 

137 data = response.json() 

138 if asyncio.iscoroutine(data): 

139 data = await data 

140 models = data.get("models", []) if isinstance(data, dict) else [] 

141 return [m.get("name", "") for m in models]