From 9172585af88ea2959f075c095e5aff0335d4b332 Mon Sep 17 00:00:00 2001
From: Milad Khoshdel <miladkhoshdel@gmail.com>
Date: Sat, 8 Aug 2026 15:35:08 +0330
Subject: [PATCH] refactor: model LLM responses with dataclasses

- add immutable response and usage dataclasses
- build structured responses from streamed output
- extract usage formatting into a helper
- add dataclass tests
---
 llm.py            | 20 ++++++++++++++++++++
 main.py           | 48 ++++++++++++++++++++++++++++++-----------------
 tests/test_llm.py | 40 +++++++++++++++++++++++++++++++++++++++
 utils.py          | 12 +++++++++++-
 4 files changed, 102 insertions(+), 18 deletions(-)
 create mode 100644 llm.py
 create mode 100644 tests/test_llm.py

diff --git a/llm.py b/llm.py
new file mode 100644
index 0000000..85f23ab
--- /dev/null
+++ b/llm.py
@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class LLMUsage:
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    reasoning_tokens: int | None = None
+    cost: float | None = None
+
+
+@dataclass(frozen=True)
+class LLMResponse:
+    content: str
+    finish_reason: str | None
+    usage: LLMUsage | None = None
+
+
+__all__ = ["LLMResponse", "LLMUsage"]
diff --git a/main.py b/main.py
index a763c1f..64cd2dd 100644
--- a/main.py
+++ b/main.py
@@ -1,7 +1,8 @@
 from openai import OpenAI, OpenAIError
 
+from llm import LLMResponse, LLMUsage
 from settings import load_settings
-from utils import read_input
+from utils import print_usage, read_input
 
 
 def trim_history(messages, keep_count):
@@ -108,31 +109,44 @@ def main():
             print(f"\nRequest failed: {error}")
             continue
 
-        answer = "".join(answer)
+        response_usage = None
 
-        if finish_reason == "length":
-            print("Warning: the answer may be incomplete.")
+        if usage is not None:
 
-        if settings.show_usage and usage:
-            print("\n-----")
-            print("Prompt tokens:", usage.prompt_tokens)
-            print("Completion tokens:", usage.completion_tokens)
-            print("Total tokens:", usage.total_tokens)
+            completion_details = getattr(
+                usage,
+                "completion_tokens_details",
+                None,
+            )
 
-            reasoning_tokens = (
-                usage.completion_tokens_details.reasoning_tokens
-                if usage.completion_tokens_details
-                else None
+            response_usage = LLMUsage(
+                prompt_tokens=usage.prompt_tokens,
+                completion_tokens=usage.completion_tokens,
+                total_tokens=usage.total_tokens,
+                reasoning_tokens=getattr(
+                    completion_details,
+                    "reasoning_tokens",
+                    None,
+                ),
+                cost=getattr(usage, "cost", None),
             )
 
-            print("Reasoning tokens:", reasoning_tokens)
-            print("Cost:", getattr(usage, "cost", None))
-            print("-----\n")
+        response = LLMResponse(
+            content="".join(answer),
+            finish_reason=finish_reason,
+            usage=response_usage,
+        )
+
+        if response.finish_reason == "length":
+            print("Warning: the answer may be incomplete.")
+
+        if settings.show_usage and response.usage:
+            print_usage(response.usage)
 
         messages.append(
             {
                 "role": "assistant",
-                "content": answer,
+                "content": response.content,
             }
         )
         messages = trim_history(messages, settings.memory_keep_count)
diff --git a/tests/test_llm.py b/tests/test_llm.py
new file mode 100644
index 0000000..0cf1679
--- /dev/null
+++ b/tests/test_llm.py
@@ -0,0 +1,40 @@
+import unittest
+from dataclasses import FrozenInstanceError
+
+from llm import LLMResponse, LLMUsage
+
+
+class LLMDataClassTests(unittest.TestCase):
+    def test_response_contains_content_finish_reason_and_usage(self):
+        usage = LLMUsage(
+            prompt_tokens=10,
+            completion_tokens=5,
+            total_tokens=15,
+            reasoning_tokens=2,
+            cost=0.001,
+        )
+
+        response = LLMResponse(
+            content="Hello",
+            finish_reason="stop",
+            usage=usage,
+        )
+
+        self.assertEqual(response.content, "Hello")
+        self.assertEqual(response.finish_reason, "stop")
+        self.assertEqual(response.usage, usage)
+
+    def test_usage_is_optional(self):
+        response = LLMResponse(content="Hello", finish_reason="stop")
+
+        self.assertIsNone(response.usage)
+
+    def test_response_is_immutable(self):
+        response = LLMResponse(content="Hello", finish_reason="stop")
+
+        with self.assertRaises(FrozenInstanceError):
+            response.content = "Changed"
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/utils.py b/utils.py
index 3e2e657..95415db 100644
--- a/utils.py
+++ b/utils.py
@@ -6,4 +6,14 @@ def read_input(prompt):
         return None
 
 
-__all__ = ["read_input"]
+def print_usage(usage):
+    print("\n-----")
+    print(f"Prompt tokens:     {usage.prompt_tokens}")
+    print(f"Completion tokens: {usage.completion_tokens}")
+    print(f"Reasoning tokens:  {usage.reasoning_tokens}")
+    print(f"Total tokens:      {usage.total_tokens}")
+    print(f"Cost:              {usage.cost}")
+    print("-----\n")
+
+
+__all__ = ["print_usage", "read_input"]