@@ -45,11 +45,13 @@ class OpenAILLMClient(LLMClient):
4545
4646 def __init__ (
4747 self ,
48+ provider : str ,
4849 api_key : str | None = None ,
4950 base_url : str | None = None ,
5051 model : str | None = None ,
5152 temperature : float = 0.1 ,
5253 ):
54+ self .provider = provider
5355 super ().__init__ (model , temperature )
5456 self .api_key = api_key or os .getenv ("OPENAI_API_KEY" )
5557 self .base_url = base_url or os .getenv ("OPENAI_BASE_URL" )
@@ -94,25 +96,79 @@ def generate(
9496 self .logger .error (f"Error generating completion: { e } )
9597 raise
9698
97- def health_check (self ) -> Dict :
99+ async def health_check (self ) -> Dict :
98100 """Check OpenAI-compatible service health."""
99101 try :
100- # For OpenAI API, check if we have valid configuration
101- # Avoid calling /models endpoint as it can be unreliable
102- if self .api_key and self .api_key != "dummy" and self .model :
102+ if not (self .model and self .base_url ):
103103 return {
104- "status" : "✅ Connected " ,
104+ "status" : "⚠️ Configuration incomplete (missing model or base_url) " ,
105105 "base_url" : self .base_url ,
106106 "default_model" : self .model ,
107107 "api_key_configured" : bool (self .api_key and self .api_key != "dummy" ),
108108 }
109- else :
110- return {
111- "status" : "⚠️ Configuration incomplete" ,
109+
110+ if self .provider == "ollama" :
111+ import aiohttp
112+ ollama_health_url = self .base_url .replace ("/v1" , "" ) if self .base_url .endswith ("/v1" ) else self .base_url
113+
114+ # Initialize response with main LLM status
115+ response_data = {
116+ "status" : "❌ Unknown" ,
112117 "base_url" : self .base_url ,
113118 "default_model" : self .model ,
114- "api_key_configured" : bool (self .api_key and self .api_key != "dummy" ),
119+ "api_key_configured" : False ,
120+ "embedder_model" : os .getenv ("OLLAMA_EMBEDDER_MODEL" ),
121+ "embedder_status" : "❌ Not Checked"
115122 }
123+
124+ try :
125+ async with aiohttp .ClientSession () as session :
126+ # Check main Ollama server health
127+ async with session .get (f"{ ollama_health_url } , timeout = aiohttp .ClientTimeout (total = 5 )) as response :
128+ if response .status == 200 :
129+ response_data ["status" ] = "✅ Connected"
130+ else :
131+ response_data ["status" ] = f"⚠️ Ollama Unhealthy: HTTP { response .status }
132+
133+ # Check embedder model availability
134+ embedder_model_name = os .getenv ("OLLAMA_EMBEDDER_MODEL" )
135+ if embedder_model_name :
136+ try :
137+ # Use /api/show to check if model exists
138+ async with session .post (f"{ ollama_health_url } , json = {"name" : embedder_model_name }, timeout = aiohttp .ClientTimeout (total = 5 )) as embedder_response :
139+ if embedder_response .status == 200 :
140+ response_data ["embedder_status" ] = "✅ Available"
141+ else :
142+ response_data ["embedder_status" ] = "⚠️ Embedder Model Unhealthy"
143+ except aiohttp .ClientError :
144+ response_data ["embedder_status" ] = "❌ Embedder Model Connection Failed"
145+ except asyncio .TimeoutError :
146+ response_data ["embedder_status" ] = "❌ Embedder Model Timeout"
147+ else :
148+ response_data ["embedder_status" ] = "⚠️ Embedder Model Not Configured"
149+
150+ except aiohttp .ClientError :
151+ response_data ["status" ] = "❌ Ollama Connection Failed"
152+ except asyncio .TimeoutError :
153+ response_data ["status" ] = "❌ Ollama Connection Timeout (5s)"
154+
155+ return response_data
156+ else :
157+ # For other OpenAI-compatible APIs, check configuration
158+ if self .api_key and self .api_key != "dummy" :
159+ return {
160+ "status" : "✅ Connected" ,
161+ "base_url" : self .base_url ,
162+ "default_model" : self .model ,
163+ "api_key_configured" : bool (self .api_key and self .api_key != "dummy" ),
164+ }
165+ else :
166+ return {
167+ "status" : "⚠️ Configuration incomplete (missing API key)" ,
168+ "base_url" : self .base_url ,
169+ "default_model" : self .model ,
170+ "api_key_configured" : bool (self .api_key and self .api_key != "dummy" ),
171+ }
116172 except Exception as e :
117173 self .logger .error (f"Health check failed: { e } )
118174 return {
@@ -135,12 +191,20 @@ def create_client() -> LLMClient:
135191 """Create an LLM client based on LLM_PROVIDER environment variable."""
136192 provider = os .getenv ("LLM_PROVIDER" , "openai" ).lower ()
137193
138- if provider in [ "openai" , "ollama" ] :
194+ if provider == "openai" :
139195 return OpenAILLMClient (
196+ provider = "openai" ,
140197 api_key = os .getenv ("OPENAI_API_KEY" ),
141198 base_url = os .getenv ("OPENAI_BASE_URL" ),
142199 model = os .getenv ("OPENAI_MODEL" ),
143200 )
201+ elif provider == "ollama" :
202+ return OpenAILLMClient (
203+ provider = "ollama" ,
204+ api_key = "dummy" , # Ollama doesn't require an API key
205+ base_url = os .getenv ("OLLAMA_BASE_URL" ),
206+ model = os .getenv ("OLLAMA_MODEL" ),
207+ )
144208 else :
145209 raise ValueError (f"Unsupported LLM provider: { provider } )
146210
@@ -181,5 +245,4 @@ async def async_generate(
181245async def async_health_check () -> Dict :
182246 """Async wrapper for LLM health check."""
183247 client = get_llm_client ()
184- loop = asyncio .get_running_loop ()
185- return await loop .run_in_executor (None , client .health_check )
248+ return await client .health_check ()
0 commit comments