import os
from openai import OpenAI

class Qwen_3_5_9B:
    def __init__(self, token=None):
        openrouter_token = os.getenv("OPENROUTER_API_KEY")
        if not openrouter_token:
            print("Warning: OPENROUTER_API_KEY environment variable is not set")
            # Fallback to the token passed in if available
            openrouter_token = token
            
        self.client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=openrouter_token,
        )
        self.model_id = "qwen/qwen3.5-9b" 

    def generate_stream(self, prompt, max_tokens=100000, temperature=0.1):
        for message in self.client.chat.completions.create(
            model=self.model_id,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=max_tokens,
            temperature=temperature,
            stream=True,
        ):
            if message.choices and len(message.choices) > 0:
                content = message.choices[0].delta.content
                if content:
                    yield content

    def generate(self, prompt, max_tokens=100000, temperature=0.1):
        return "".join(self.generate_stream(prompt, max_tokens=max_tokens, temperature=temperature))