diff options
Diffstat (limited to 'agent.py')
| -rw-r--r-- | agent.py | 45 |
1 files changed, 45 insertions, 0 deletions
diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..51661fd --- /dev/null +++ b/agent.py @@ -0,0 +1,45 @@ +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, + }) |
