from collections.abc import AsyncIterator

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

from tests.fixtures.agents.prompt import SYSTEM_PROMPT


class LangchainAgent:
    def ask(self, question: str) -> str:
        return (
            ChatPromptTemplate.from_messages(
                [
                    ("system", SYSTEM_PROMPT),
                    ("human", "{question}"),
                ]
            )
            | ChatOpenAI(model="gpt-4o-mini", temperature=0)
            | StrOutputParser()
        ).invoke(
            {
                "question": question,
            }
        )

    async def stream(self, question: str) -> AsyncIterator[str]:
        async for chunk in (
            ChatPromptTemplate.from_messages(
                [
                    ("system", SYSTEM_PROMPT),
                    ("human", "{question}"),
                ]
            )
            | ChatOpenAI(model="gpt-4o-mini", temperature=0, streaming=True)
            | StrOutputParser()
        ).astream({"question": question}):
            yield chunk
