# /services.py """ 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) """ 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 import fireworks.client as Fireworks import openai import google.generativeai as genai from deepseek import OpenaiClient as DeepSeekClient # --- Setup Logging & Environment --- logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') load_dotenv() # --- API Keys from .env --- HF_TOKEN = os.getenv("HF_TOKEN") TAVILY_API_KEY = os.getenv("TAVILY_API_KEY") GROQ_API_KEY = os.getenv("GROQ_API_KEY") FIREWORKS_API_KEY = os.getenv("FIREWORKS_API_KEY") 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 self.deepseek_client = DeepSeekClient(api_key=DEEPSEEK_API_KEY) if DEEPSEEK_API_KEY else None if FIREWORKS_API_KEY: Fireworks.api_key = FIREWORKS_API_KEY self.fireworks_client = Fireworks else: self.fireworks_client = None if GEMINI_API_KEY: genai.configure(api_key=GEMINI_API_KEY) self.gemini_model = genai.GenerativeModel('gemini-1.5-pro-latest') else: self.gemini_model = None def _prepare_messages_for_gemini(self, messages: Messages) -> List[Dict[str, Any]]: """Gemini requires a slightly different message format.""" gemini_messages = [] for msg in messages: # Gemini uses 'model' for assistant role 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]: """ Streams code generation, dispatching to the correct provider based on model_id. """ provider, model_name = model_id.split('/', 1) logging.info(f"Dispatching to provider: {provider} for model: {model_name}") try: # --- OpenAI, Groq, DeepSeek, Fireworks (OpenAI-compatible) --- 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 = client_map.get(provider) if not client: raise ValueError(f"{provider.capitalize()} API key not configured.") # Fireworks has a slightly different call signature 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) for chunk in stream: if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content: yield chunk.choices[0].delta.content # --- Google Gemini --- elif provider == 'gemini': if not self.gemini_model: raise ValueError("Gemini API key not configured.") gemini_messages = self._prepare_messages_for_gemini(messages) stream = self.gemini_model.generate_content(gemini_messages, stream=True) for chunk in stream: yield chunk.text # --- Hugging Face --- elif provider == 'huggingface': if not self.hf_client: raise ValueError("Hugging Face API token not configured.") # For HF, model_name is the rest of the ID, e.g., baidu/ERNIE... 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 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)}" class SearchService: """A wrapper for the Tavily Search API.""" 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: try: self.client = TavilyClient(api_key=api_key) except Exception as e: logging.error(f"Failed to initialize Tavily client: {e}") self.client = None def is_available(self) -> bool: """Checks if the search service is configured and available.""" return self.client is not None def search(self, query: str, max_results: int = 5) -> str: """Performs a web search and returns a formatted string of results.""" 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)}" # --- Singleton Instances --- llm_service = LLMService() search_service = SearchService()