import dashscope
from dashscope import Generation
from typing import AsyncGenerator, Optional, Generator


class LLMService:
    def __init__(self):
        from app.core.config import settings
        self.model = settings.llm_model
        dashscope.api_key = settings.dashscope_api_key

    def generate_stream(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
    ) -> Generator[str, None, None]:
        """流式生成文本"""
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})

        responses = Generation.call(
            model=self.model,
            messages=messages,
            result_format="message",
            stream=True,
        )

        for response in responses:
            if response.status_code == 200:
                content = response.output.choices[0].message.content
                if content:
                    yield content

    async def generate(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
    ) -> str:
        """一次性生成文本"""
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})

        response = Generation.call(
            model=self.model,
            messages=messages,
            result_format="message",
        )

        if response.status_code == 200:
            return response.output.choices[0].message.content
        raise Exception(f"LLM generation failed: {response.code}")


llm_service = LLMService()