scripts/generators.py

"""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}")