diff --git a/src/cogs/ollama.py b/src/cogs/ollama.py index ea41170..4122ecb 100644 --- a/src/cogs/ollama.py +++ b/src/cogs/ollama.py @@ -20,6 +20,7 @@ from discord.ui import View, button from yarl import URL from conf import CONFIG +from .quote_quota import TruthPayload async def ollama_stream(iterator: aiohttp.StreamReader) -> typing.AsyncIterator[dict]: @@ -1036,10 +1037,11 @@ class Ollama(commands.Cog): auth=(username, password), ) response.raise_for_status() - truths = response.json() + truths: list[dict] = response.json() + truths: list[TruthPayload] = list(map(lambda t: TruthPayload.model_validate(t), truths)) for truth in truths: - if truth["author"] == "trump": - self.history.add_message(thread_id, "assistant", truth["content"], save=False) + if truth.author == "trump": + self.history.add_message(thread_id, "assistant", truth.content, save=False) break self.history.add_message(thread_id, "user", "Generate a new truth post.")