from collections.abc import Iterator

from context import ContextBuilder
from identity import AgentIdentity, UserProfile
from llm import LLM

class Agent:
    def __init__(
        self,
        llm: LLM,
        agent_identity: AgentIdentity,
        user_profile: UserProfile
    ) -> None:
        self.llm = llm

        self.context = ContextBuilder(
            agent_identity,
            user_profile,
        )

        self.messages: list[dict[str, str]] = [
            {
                "role": "system",
                "content": self.context.build_system_prompt(),
            }
        ]

    def chat(self, user_input: str) -> Iterator[str]:
        self.messages.append({
            "role": "user",
            "content": user_input,
        })

        answer_parts: list[str] = []

        for chunk in self.llm.chat(self.messages):
            yield chunk
            answer_parts.append(chunk)

        answer = "".join(answer_parts)

        self.messages.append({
            "role": "assistant",
            "content": answer,
        })