from collections.abc import AsyncIterator
from typing import Any

from tests.fixtures.agents.prompt import SYSTEM_PROMPT
from tests.fixtures.datasource import datasource

from .litellm import run_completion, run_completion_async


class LitellmAgent:
    def ask(self, question: str) -> list[Any]:
        response = run_completion(
            [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": question},
            ],
            datasource,
        )

        return response.choices

    async def stream(self, question: str) -> AsyncIterator[str]:
        async for part in run_completion_async(
            [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": question},
            ],
            datasource,
        ):
            yield part
