Update services.py

services.py

@@ -1,22 +1,11 @@
 # /services.py
+""" Manages interactions with all external LLM and search APIs. """
 
-"""
-Manages interactions with external services like LLM providers and web search APIs.
-This module has been refactored to support multiple LLM providers:
-- Hugging Face
-- Groq
-- Fireworks AI
-- OpenAI
-- Google Gemini
-- DeepSeek (Direct API via OpenAI client)
-"""
 import os
 import logging
 from typing import Dict, Any, Generator, List
 
 from dotenv import load_dotenv
-
-# Import all necessary clients
 from huggingface_hub import InferenceClient
 from tavily import TavilyClient
 from groq import Groq
@@ -24,9 +13,6 @@ import fireworks.client as Fireworks
 import openai
 import google.generativeai as genai
 
-# <--- FIX: REMOVED the incorrect 'from deepseek import ...' line ---
-
-# --- Setup Logging & Environment ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 load_dotenv()
 
@@ -39,24 +25,17 @@ OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
 DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
 
-# --- Type Definitions ---
 Messages = List[Dict[str, Any]]
 
 class LLMService:
     """A multi-provider wrapper for LLM Inference APIs."""
-
     def __init__(self):
-        # Initialize clients only if their API keys are available
         self.hf_client = InferenceClient(token=HF_TOKEN) if HF_TOKEN else None
         self.groq_client = Groq(api_key=GROQ_API_KEY) if GROQ_API_KEY else None
         self.openai_client = openai.OpenAI(api_key=OPENAI_API_KEY) if OPENAI_API_KEY else None
         
-        # <--- FIX: Correctly instantiate the DeepSeek client using the OpenAI library ---
         if DEEPSEEK_API_KEY:
-            self.deepseek_client = openai.OpenAI(
-                api_key=DEEPSEEK_API_KEY,
-                base_url="https://api.deepseek.com/v1"
-            )
+            self.deepseek_client = openai.OpenAI(api_key=DEEPSEEK_API_KEY, base_url="https://api.deepseek.com/v1")
         else:
             self.deepseek_client = None
 
@@ -73,84 +52,60 @@ class LLMService:
             self.gemini_model = None
 
     def _prepare_messages_for_gemini(self, messages: Messages) -> List[Dict[str, Any]]:
-        # This function remains the same
         gemini_messages = []
         for msg in messages:
+            if msg['role'] == 'system': continue # Gemini doesn't use a system role in this way
             role = 'model' if msg['role'] == 'assistant' else 'user'
             gemini_messages.append({'role': role, 'parts': [msg['content']]})
         return gemini_messages
 
-    def generate_code_stream(
-        self, model_id: str, messages: Messages, max_tokens: int = 8192
-    ) -> Generator[str, None, None]:
-        # This function remains the same, as the dispatcher logic is already correct
+    def generate_code_stream(self, model_id: str, messages: Messages, max_tokens: int = 8192) -> Generator[str, None, None]:
         provider, model_name = model_id.split('/', 1)
         logging.info(f"Dispatching to provider: {provider} for model: {model_name}")
 
         try:
             if provider in ['openai', 'groq', 'deepseek', 'fireworks']:
-                client_map = {
-                    'openai': self.openai_client,
-                    'groq': self.groq_client,
-                    'deepseek': self.deepseek_client,
-                    'fireworks': self.fireworks_client.ChatCompletion if self.fireworks_client else None,
-                }
+                client_map = {'openai': self.openai_client, 'groq': self.groq_client, 'deepseek': self.deepseek_client, 'fireworks': self.fireworks_client.ChatCompletion if self.fireworks_client else None}
                 client = client_map.get(provider)
-                if not client:
-                    raise ValueError(f"{provider.capitalize()} API key not configured.")
-                
-                if provider == 'fireworks':
-                     stream = client.create(model=model_name, messages=messages, stream=True, max_tokens=max_tokens)
-                else:
-                     stream = client.chat.completions.create(model=model_name, messages=messages, stream=True, max_tokens=max_tokens)
+                if not client: raise ValueError(f"{provider.capitalize()} API key not configured.")
                 
+                stream = client.create(model=model_name, messages=messages, stream=True, max_tokens=max_tokens) if provider == 'fireworks' else client.chat.completions.create(model=model_name, messages=messages, stream=True, max_tokens=max_tokens)
                 for chunk in stream:
-                    if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content:
-                        yield chunk.choices[0].delta.content
-
+                    if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content: yield chunk.choices[0].delta.content
+            
             elif provider == 'gemini':
-                if not self.gemini_model:
-                    raise ValueError("Gemini API key not configured.")
+                if not self.gemini_model: raise ValueError("Gemini API key not configured.")
+                system_prompt = next((msg['content'] for msg in messages if msg['role'] == 'system'), "")
                 gemini_messages = self._prepare_messages_for_gemini(messages)
+                # Prepend system prompt to first user message for Gemini
+                if system_prompt and gemini_messages and gemini_messages[0]['role'] == 'user':
+                    gemini_messages[0]['parts'][0] = f"{system_prompt}\n\n{gemini_messages[0]['parts'][0]}"
                 stream = self.gemini_model.generate_content(gemini_messages, stream=True)
-                for chunk in stream:
-                    yield chunk.text
+                for chunk in stream: yield chunk.text
 
             elif provider == 'huggingface':
-                if not self.hf_client:
-                    raise ValueError("Hugging Face API token not configured.")
+                if not self.hf_client: raise ValueError("Hugging Face API token not configured.")
                 hf_model_id = model_id.split('/', 1)[1]
                 stream = self.hf_client.chat_completion(model=hf_model_id, messages=messages, stream=True, max_tokens=max_tokens)
                 for chunk in stream:
-                    if chunk.choices[0].delta.content:
-                        yield chunk.choices[0].delta.content
+                    if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content: yield chunk.choices[0].delta.content
             else:
                 raise ValueError(f"Unknown provider: {provider}")
-
         except Exception as e:
             logging.error(f"LLM API Error with provider {provider}: {e}")
             yield f"Error from {provider.capitalize()}: {str(e)}"
 
-# The SearchService class remains unchanged
 class SearchService:
     def __init__(self, api_key: str = TAVILY_API_KEY):
-        if not api_key:
-            logging.warning("TAVILY_API_KEY not set. Web search will be disabled.")
-            self.client = None
-        else:
-            self.client = TavilyClient(api_key=api_key)
-    def is_available(self) -> bool:
-        return self.client is not None
+        self.client = TavilyClient(api_key=api_key) if api_key else None
+        if not self.client: logging.warning("TAVILY_API_KEY not set. Web search will be disabled.")
+    def is_available(self) -> bool: return self.client is not None
     def search(self, query: str, max_results: int = 5) -> str:
         if not self.is_available(): return "Web search is not available."
         try:
             response = self.client.search(query, search_depth="advanced", max_results=min(max(1, max_results), 10))
-            results = [f"Title: {res.get('title', 'N/A')}\nURL: {res.get('url', 'N/A')}\nContent: {res.get('content', 'N/A')}" for res in response.get('results', [])]
-            return "Web Search Results:\n\n" + "\n---\n".join(results) if results else "No search results found."
-        except Exception as e:
-            logging.error(f"Tavily search error: {e}")
-            return f"Search error: {str(e)}"
+            return "Web Search Results:\n\n" + "\n---\n".join([f"Title: {res.get('title', 'N/A')}\nURL: {res.get('url', 'N/A')}\nContent: {res.get('content', 'N/A')}" for res in response.get('results', [])])
+        except Exception as e: return f"Search error: {str(e)}"
 
-# --- Singleton Instances ---
 llm_service = LLMService()
 search_service = SearchService()