Nihal2000 commited on
Commit
30025c3
·
verified ·
1 Parent(s): 1992efa

Update services/llm_service.py

Browse files
Files changed (1) hide show
  1. services/llm_service.py +43 -91
services/llm_service.py CHANGED
@@ -3,7 +3,6 @@ import logging
3
  import asyncio
4
  from typing import List, Dict, Any, Optional
5
 
6
- import anthropic
7
  import openai
8
  import config
9
 
@@ -13,35 +12,29 @@ class LLMService:
13
  def __init__(self):
14
  self.config = config.config
15
 
16
- self.anthropic_client = None
17
  self.mistral_client = None
18
- self.openai_async_client = None
19
 
20
  self._initialize_clients()
21
 
22
  def _initialize_clients(self):
23
  """Initialize LLM clients"""
24
  try:
25
- if self.config.ANTHROPIC_API_KEY:
26
- self.anthropic_client = anthropic.Anthropic(
27
- api_key=self.config.ANTHROPIC_API_KEY
 
28
  )
29
- logger.info("Anthropic client initialized")
30
 
31
  if self.config.MISTRAL_API_KEY:
32
  self.mistral_client = Mistral( # Standard sync client
33
  api_key=self.config.MISTRAL_API_KEY
34
  )
35
  logger.info("Mistral client initialized")
36
-
37
- if self.config.OPENAI_API_KEY:
38
- self.openai_async_client = openai.AsyncOpenAI(
39
- api_key=self.config.OPENAI_API_KEY
40
- )
41
- logger.info("OpenAI client initialized")
42
 
43
  # Check if at least one client is initialized
44
- if not any([self.openai_async_client, self.mistral_client, self.anthropic_client]):
45
  logger.warning("No LLM clients could be initialized based on current config. Check API keys.")
46
  else:
47
  logger.info("LLM clients initialized successfully (at least one).")
@@ -56,99 +49,68 @@ class LLMService:
56
  selected_model_name_for_call: str = ""
57
 
58
  if model == "auto":
59
- # New Priority: 1. OpenAI, 2. Mistral, 3. Anthropic
60
- if self.openai_async_client and self.config.OPENAI_MODEL:
61
- selected_model_name_for_call = self.config.OPENAI_MODEL
62
- logger.debug(f"Auto-selected OpenAI model: {selected_model_name_for_call}")
63
- return await self._generate_with_openai(prompt, selected_model_name_for_call, max_tokens, temperature)
64
  elif self.mistral_client and self.config.MISTRAL_MODEL:
65
  selected_model_name_for_call = self.config.MISTRAL_MODEL
66
  logger.debug(f"Auto-selected Mistral model: {selected_model_name_for_call}")
67
  return await self._generate_with_mistral(prompt, selected_model_name_for_call, max_tokens, temperature)
68
- elif self.anthropic_client and self.config.ANTHROPIC_MODEL:
69
- selected_model_name_for_call = self.config.ANTHROPIC_MODEL
70
- logger.debug(f"Auto-selected Anthropic model: {selected_model_name_for_call}")
71
- return await self._generate_with_claude(prompt, selected_model_name_for_call, max_tokens, temperature)
72
  else:
73
  logger.error("No LLM clients available for 'auto' mode or default models not configured.")
74
  raise ValueError("No LLM clients available for 'auto' mode or default models not configured.")
75
 
76
- elif model.startswith("gpt-") or model.lower().startswith("openai/"):
77
- if not self.openai_async_client:
78
- raise ValueError("OpenAI client not available. Check API key or model prefix.")
79
  actual_model = model.split('/')[-1] if '/' in model else model
80
- return await self._generate_with_openai(prompt, actual_model, max_tokens, temperature)
81
 
82
  elif model.startswith("mistral"):
83
  if not self.mistral_client:
84
  raise ValueError("Mistral client not available. Check API key or model prefix.")
85
  return await self._generate_with_mistral(prompt, model, max_tokens, temperature)
86
-
87
- elif model.startswith("claude"):
88
- if not self.anthropic_client:
89
- raise ValueError("Anthropic client not available. Check API key or model prefix.")
90
- return await self._generate_with_claude(prompt, model, max_tokens, temperature)
91
 
92
  else:
93
- raise ValueError(f"Unsupported model: {model}. Must start with 'gpt-', 'openai/', 'claude', 'mistral', or be 'auto'.")
94
 
95
  except Exception as e:
96
  logger.error(f"Error generating text with model '{model}': {str(e)}")
97
  raise
98
 
99
- async def _generate_with_openai(self, prompt: str, model_name: str, max_tokens: int, temperature: float) -> str:
100
- """Generate text using OpenAI (Async)"""
101
- if not self.openai_async_client:
102
- raise RuntimeError("OpenAI async client not initialized.")
103
  try:
104
- logger.debug(f"Generating with OpenAI model: {model_name}, max_tokens: {max_tokens}, temp: {temperature}, prompt: '{prompt[:50]}...'")
105
- response = await self.openai_async_client.chat.completions.create(
106
- model=model_name,
107
- messages=[{"role": "user", "content": prompt}],
108
- max_tokens=max_tokens,
109
- temperature=temperature
 
 
 
 
 
110
  )
111
  if response.choices and response.choices[0].message:
112
  content = response.choices[0].message.content
113
  if content is not None:
114
  return content.strip()
115
  else:
116
- logger.warning(f"OpenAI response message content is None for model {model_name}.")
117
  return ""
118
  else:
119
- logger.warning(f"OpenAI response did not contain expected choices or message for model {model_name}.")
120
  return ""
121
  except Exception as e:
122
- logger.error(f"Error with OpenAI generation (model: {model_name}): {str(e)}")
123
  raise
124
 
125
- async def _generate_with_claude(self, prompt: str, model_name: str, max_tokens: int, temperature: float) -> str:
126
- """Generate text using Anthropic/Claude (Sync via run_in_executor)"""
127
- if not self.anthropic_client:
128
- raise RuntimeError("Anthropic client not initialized.")
129
- try:
130
- logger.debug(f"Generating with Anthropic model: {model_name}, max_tokens: {max_tokens}, temp: {temperature}, prompt: '{prompt[:50]}...'")
131
- loop = asyncio.get_event_loop()
132
- response = await loop.run_in_executor(
133
- None,
134
- lambda: self.anthropic_client.messages.create(
135
- model=model_name,
136
- max_tokens=max_tokens,
137
- temperature=temperature,
138
- messages=[
139
- {"role": "user", "content": prompt}
140
- ]
141
- )
142
- )
143
- if response.content and response.content[0].text:
144
- return response.content[0].text.strip()
145
- else:
146
- logger.warning(f"Anthropic response did not contain expected content for model {model_name}.")
147
- return ""
148
- except Exception as e:
149
- logger.error(f"Error with Anthropic (Claude) generation (model: {model_name}): {str(e)}")
150
- raise
151
-
152
  async def _generate_with_mistral(self, prompt: str, model_name: str, max_tokens: int, temperature: float) -> str:
153
  """Generate text using Mistral (Sync via run_in_executor)"""
154
  if not self.mistral_client:
@@ -348,9 +310,8 @@ Answer:"""
348
  async def check_availability(self) -> Dict[str, bool]:
349
  """Check which LLM services are available by making a tiny test call."""
350
  availability = {
351
- "openai": False,
352
- "mistral": False,
353
- "anthropic": False
354
  }
355
  test_prompt = "Hello"
356
  test_max_tokens = 5
@@ -358,14 +319,14 @@ Answer:"""
358
 
359
  logger.info("Checking LLM availability...")
360
 
361
- if self.openai_async_client and self.config.OPENAI_MODEL:
362
  try:
363
- logger.debug(f"Testing OpenAI availability with model {self.config.OPENAI_MODEL}...")
364
- test_response = await self._generate_with_openai(test_prompt, self.config.OPENAI_MODEL, test_max_tokens, test_temp)
365
- availability["openai"] = bool(test_response.strip())
366
  except Exception as e:
367
- logger.warning(f"OpenAI availability check failed for model {self.config.OPENAI_MODEL}: {e}")
368
- logger.info(f"OpenAI available: {availability['openai']}")
369
 
370
  if self.mistral_client and self.config.MISTRAL_MODEL:
371
  try:
@@ -375,15 +336,6 @@ Answer:"""
375
  except Exception as e:
376
  logger.warning(f"Mistral availability check failed for model {self.config.MISTRAL_MODEL}: {e}")
377
  logger.info(f"Mistral available: {availability['mistral']}")
378
-
379
- if self.anthropic_client and self.config.ANTHROPIC_MODEL:
380
- try:
381
- logger.debug(f"Testing Anthropic availability with model {self.config.ANTHROPIC_MODEL}...")
382
- test_response = await self._generate_with_claude(test_prompt, self.config.ANTHROPIC_MODEL, test_max_tokens, test_temp)
383
- availability["anthropic"] = bool(test_response.strip())
384
- except Exception as e:
385
- logger.warning(f"Anthropic availability check failed for model {self.config.ANTHROPIC_MODEL}: {e}")
386
- logger.info(f"Anthropic available: {availability['anthropic']}")
387
 
388
  logger.info(f"Final LLM Availability: {availability}")
389
  return availability
 
3
  import asyncio
4
  from typing import List, Dict, Any, Optional
5
 
 
6
  import openai
7
  import config
8
 
 
12
  def __init__(self):
13
  self.config = config.config
14
 
15
+ self.nebius_client = None
16
  self.mistral_client = None
 
17
 
18
  self._initialize_clients()
19
 
20
  def _initialize_clients(self):
21
  """Initialize LLM clients"""
22
  try:
23
+ if self.config.NEBIUS_API_KEY:
24
+ self.nebius_client = openai.OpenAI(
25
+ api_key=self.config.NEBIUS_API_KEY,
26
+ base_url=self.config.NEBIUS_BASE_URL
27
  )
28
+ logger.info("NEBIUS client initialized")
29
 
30
  if self.config.MISTRAL_API_KEY:
31
  self.mistral_client = Mistral( # Standard sync client
32
  api_key=self.config.MISTRAL_API_KEY
33
  )
34
  logger.info("Mistral client initialized")
 
 
 
 
 
 
35
 
36
  # Check if at least one client is initialized
37
+ if not any([self.nebius_client, self.mistral_client]):
38
  logger.warning("No LLM clients could be initialized based on current config. Check API keys.")
39
  else:
40
  logger.info("LLM clients initialized successfully (at least one).")
 
49
  selected_model_name_for_call: str = ""
50
 
51
  if model == "auto":
52
+ # New Priority: 1. NEBIUS (OpenAI OSS), 2. Mistral
53
+ if self.nebius_client and self.config.NEBIUS_MODEL:
54
+ selected_model_name_for_call = self.config.NEBIUS_MODEL
55
+ logger.debug(f"Auto-selected NEBIUS model: {selected_model_name_for_call}")
56
+ return await self._generate_with_nebius(prompt, selected_model_name_for_call, max_tokens, temperature)
57
  elif self.mistral_client and self.config.MISTRAL_MODEL:
58
  selected_model_name_for_call = self.config.MISTRAL_MODEL
59
  logger.debug(f"Auto-selected Mistral model: {selected_model_name_for_call}")
60
  return await self._generate_with_mistral(prompt, selected_model_name_for_call, max_tokens, temperature)
 
 
 
 
61
  else:
62
  logger.error("No LLM clients available for 'auto' mode or default models not configured.")
63
  raise ValueError("No LLM clients available for 'auto' mode or default models not configured.")
64
 
65
+ elif model.startswith("gpt-") or model.startswith("openai/") or model.lower().startswith("nebius/"):
66
+ if not self.nebius_client:
67
+ raise ValueError("NEBIUS client not available. Check API key or model prefix.")
68
  actual_model = model.split('/')[-1] if '/' in model else model
69
+ return await self._generate_with_nebius(prompt, actual_model, max_tokens, temperature)
70
 
71
  elif model.startswith("mistral"):
72
  if not self.mistral_client:
73
  raise ValueError("Mistral client not available. Check API key or model prefix.")
74
  return await self._generate_with_mistral(prompt, model, max_tokens, temperature)
 
 
 
 
 
75
 
76
  else:
77
+ raise ValueError(f"Unsupported model: {model}. Must start with 'gpt-', 'openai/', 'nebius/', 'mistral', or be 'auto'.")
78
 
79
  except Exception as e:
80
  logger.error(f"Error generating text with model '{model}': {str(e)}")
81
  raise
82
 
83
+ async def _generate_with_nebius(self, prompt: str, model_name: str, max_tokens: int, temperature: float) -> str:
84
+ """Generate text using NEBIUS (OpenAI OSS models via sync client)"""
85
+ if not self.nebius_client:
86
+ raise RuntimeError("NEBIUS client not initialized.")
87
  try:
88
+ logger.debug(f"Generating with NEBIUS model: {model_name}, max_tokens: {max_tokens}, temp: {temperature}, prompt: '{prompt[:50]}...'")
89
+ loop = asyncio.get_event_loop()
90
+
91
+ response = await loop.run_in_executor(
92
+ None,
93
+ lambda: self.nebius_client.chat.completions.create(
94
+ model=model_name,
95
+ messages=[{"role": "user", "content": prompt}],
96
+ max_tokens=max_tokens,
97
+ temperature=temperature
98
+ )
99
  )
100
  if response.choices and response.choices[0].message:
101
  content = response.choices[0].message.content
102
  if content is not None:
103
  return content.strip()
104
  else:
105
+ logger.warning(f"NEBIUS response message content is None for model {model_name}.")
106
  return ""
107
  else:
108
+ logger.warning(f"NEBIUS response did not contain expected choices or message for model {model_name}.")
109
  return ""
110
  except Exception as e:
111
+ logger.error(f"Error with NEBIUS generation (model: {model_name}): {str(e)}")
112
  raise
113
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
114
  async def _generate_with_mistral(self, prompt: str, model_name: str, max_tokens: int, temperature: float) -> str:
115
  """Generate text using Mistral (Sync via run_in_executor)"""
116
  if not self.mistral_client:
 
310
  async def check_availability(self) -> Dict[str, bool]:
311
  """Check which LLM services are available by making a tiny test call."""
312
  availability = {
313
+ "nebius": False,
314
+ "mistral": False
 
315
  }
316
  test_prompt = "Hello"
317
  test_max_tokens = 5
 
319
 
320
  logger.info("Checking LLM availability...")
321
 
322
+ if self.nebius_client and self.config.NEBIUS_MODEL:
323
  try:
324
+ logger.debug(f"Testing NEBIUS availability with model {self.config.NEBIUS_MODEL}...")
325
+ test_response = await self._generate_with_nebius(test_prompt, self.config.NEBIUS_MODEL, test_max_tokens, test_temp)
326
+ availability["nebius"] = bool(test_response.strip())
327
  except Exception as e:
328
+ logger.warning(f"NEBIUS availability check failed for model {self.config.NEBIUS_MODEL}: {e}")
329
+ logger.info(f"NEBIUS available: {availability['nebius']}")
330
 
331
  if self.mistral_client and self.config.MISTRAL_MODEL:
332
  try:
 
336
  except Exception as e:
337
  logger.warning(f"Mistral availability check failed for model {self.config.MISTRAL_MODEL}: {e}")
338
  logger.info(f"Mistral available: {availability['mistral']}")
 
 
 
 
 
 
 
 
 
339
 
340
  logger.info(f"Final LLM Availability: {availability}")
341
  return availability