from typing import AsyncIterator
import litellm
from app.core.config import DEFAULT_MODEL

litellm.drop_params = True


async def stream_chat(
    messages: list[dict],
    model: str | None = None,
) -> AsyncIterator[str]:
    response = await litellm.acompletion(
        model=model or DEFAULT_MODEL,
        messages=messages,
        stream=True,
    )
    async for chunk in response:
        delta = chunk.choices[0].delta.content
        if delta:
            yield delta


async def complete(messages: list[dict], model: str | None = None) -> str:
    response = await litellm.acompletion(
        model=model or DEFAULT_MODEL,
        messages=messages,
        stream=False,
    )
    return response.choices[0].message.content or ""