File size: 4,137 Bytes
1687ea3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
# /services.py
"""
Manages interactions with external services like LLM providers and web search APIs.
This module uses a class-based approach to encapsulate API clients and their
logic, making it easy to manage connections and mock services for testing.
"""
import os
import logging
from typing import Dict, Any, Generator, List
from dotenv import load_dotenv
from huggingface_hub import InferenceClient
from tavily import TavilyClient
# --- Setup Logging ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
# --- Load Environment Variables ---
load_dotenv()
HF_TOKEN = os.getenv("HF_TOKEN")
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
if not HF_TOKEN:
raise ValueError("HF_TOKEN environment variable is not set. Please get a token from https://huggingface.co/settings/tokens")
# --- Type Definitions ---
Messages = List[Dict[str, Any]]
class LLMService:
"""A wrapper for the Hugging Face Inference API."""
def __init__(self, api_key: str = HF_TOKEN):
if not api_key:
raise ValueError("Hugging Face API key is required.")
self.api_key = api_key
def get_client(self, model_id: str, provider: str = "auto") -> InferenceClient:
"""Initializes and returns an InferenceClient."""
return InferenceClient(provider=provider, api_key=self.api_key, bill_to="huggingface")
def generate_code_stream(
self, model_id: str, messages: Messages, provider: str = "auto", max_tokens: int = 10000
) -> Generator[str, None, None]:
"""
Streams code generation from the specified model.
Yields content chunks as they are received.
"""
client = self.get_client(model_id, provider)
try:
stream = client.chat.completions.create(
model=model_id,
messages=messages,
stream=True,
max_tokens=max_tokens,
)
for chunk in stream:
if chunk.choices and chunk.choices[0].delta.content:
yield chunk.choices[0].delta.content
except Exception as e:
logging.error(f"LLM API Error for model {model_id}: {e}")
yield f"Error: Could not get a response from the model. Details: {str(e)}"
# Re-raise or handle as appropriate for your application flow
# For this app, we yield an error message to the user.
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 ---
# These instances can be imported and used throughout the application.
llm_service = LLMService()
search_service = SearchService() |