"""Provider wrappers for question generation. Both providers expose complete(prompt) so dataset code can share one interface. """ import os class DeepSeekSource: """DeepSeek-based question source.""" def __init__(self, model_name, temperature): """Initialize DeepSeek client from env key.""" from openai import OpenAI api_key = os.environ.get("DEEPSEEK_API_KEY") if not api_key: raise SystemExit( "DEEPSEEK_API_KEY is not set. Add it to .env. Get a key with " "free credits at https://platform.deepseek.com (no card needed)." ) self.client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com") self.model_name = model_name self.temperature = temperature self.label = "deepseek:" + model_name def complete(self, prompt): """Send one prompt and return raw text.""" response = self.client.chat.completions.create( model=self.model_name, messages=[{"role": "user", "content": prompt}], temperature=self.temperature, response_format={"type": "json_object"}, ) return response.choices[0].message.content class GeminiSource: """Gemini-based question source.""" def __init__(self, model_name, temperature): """Initialize Gemini client from env key.""" from google import genai api_key = os.environ.get("GEMINI_API_KEY") if not api_key: raise SystemExit("GEMINI_API_KEY is not set. Add it to .env.") self.client = genai.Client(api_key=api_key) self.model_name = model_name self.temperature = temperature self.label = "gemini:" + model_name def complete(self, prompt): """Send one prompt and return raw text.""" from google.genai import types response = self.client.models.generate_content( model=self.model_name, contents=prompt, config=types.GenerateContentConfig( response_mime_type="application/json", temperature=self.temperature, ), ) return response.text def build_source(provider, model_name, temperature): """Return provider-specific source instance.""" if provider == "deepseek": return DeepSeekSource(model_name, temperature) if provider == "gemini": return GeminiSource(model_name, temperature) raise ValueError(f"Unknown provider: {provider}")