diff --git a/llm/__init__.py b/llm/__init__.py new file mode 100644 index 0000000..57bfdfd --- /dev/null +++ b/llm/__init__.py @@ -0,0 +1,11 @@ +from .client import LLM +from .models import LLMResponse, LLMUsage +from .settings import Settings, load_settings + +__all__ = [ + "LLM", + "LLMResponse", + "LLMUsage", + "Settings", + "load_settings", +] diff --git a/llm/chat.py b/llm/chat.py new file mode 100644 index 0000000..59acfde --- /dev/null +++ b/llm/chat.py @@ -0,0 +1,87 @@ +from openai import OpenAIError + +from .client import LLM +from .console import print_usage, read_input +from .settings import load_settings + + +def trim_history(messages, keep_count): + """Keep the system message and the latest completed exchanges.""" + if keep_count < 0: + raise ValueError("MEMORY_KEEP_COUNT must be 0 or greater") + + if keep_count == 0: + return messages[:1] + + conversation = messages[1:] + conversation_start = -(keep_count * 2) + return messages[:1] + conversation[conversation_start:] + + +def run_chat(): + try: + settings = load_settings() + except ValueError as error: + print(f"Configuration error: {error}") + return 1 + + model = LLM(settings) + + while True: + system_rule = read_input("System rule: ") + + if system_rule is None: + return 0 + + if system_rule: + break + + print("Please enter a system rule.") + + messages = [ + { + "role": "system", + "content": system_rule, + } + ] + + while True: + user_input = read_input("\nYou: ") + + if user_input is None or user_input.lower() in {"exit", "quit"}: + return 0 + + if not user_input: + print("Please enter a message.") + continue + + messages.append( + { + "role": "user", + "content": user_input, + } + ) + + try: + response = model.complete( + messages, + on_text=lambda text: print(text, end="", flush=True), + ) + except OpenAIError as error: + messages.pop() + print(f"\nRequest failed: {error}") + continue + + 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": response.content, + } + ) + messages = trim_history(messages, settings.memory_keep_count) diff --git a/llm.py b/llm/client.py similarity index 88% rename from llm.py rename to llm/client.py index 49d6899..c1e06e3 100644 --- a/llm.py +++ b/llm/client.py @@ -1,7 +1,7 @@ -from dataclasses import dataclass - from openai import OpenAI +from .models import LLMResponse, LLMUsage + class LLM: def __init__(self, settings, client=None): @@ -91,20 +91,4 @@ def complete(self, messages, on_text=None): return self.build_response(answer, finish_reason, usage) -@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__ = ["LLM", "LLMResponse", "LLMUsage"] diff --git a/utils.py b/llm/console.py similarity index 100% rename from utils.py rename to llm/console.py diff --git a/llm/models.py b/llm/models.py new file mode 100644 index 0000000..85f23ab --- /dev/null +++ b/llm/models.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/settings.py b/llm/settings.py similarity index 100% rename from settings.py rename to llm/settings.py diff --git a/main.py b/main.py index fe65043..981f781 100644 --- a/main.py +++ b/main.py @@ -1,91 +1,4 @@ -from openai import OpenAIError - -from llm import LLM -from settings import load_settings -from utils import print_usage, read_input - - -def trim_history(messages, keep_count): - """Keep the system message and the latest completed exchanges.""" - if keep_count < 0: - raise ValueError("MEMORY_KEEP_COUNT must be 0 or greater") - - if keep_count == 0: - return messages[:1] - - conversation = messages[1:] - conversation_start = -(keep_count * 2) - return messages[:1] + conversation[conversation_start:] - - -def main(): - try: - settings = load_settings() - except ValueError as error: - print(f"Configuration error: {error}") - return 1 - - model = LLM(settings) - - while True: - system_rule = read_input("System rule: ") - - if system_rule is None: - return 0 - - if system_rule: - break - - print("Please enter a system rule.") - - messages = [ - { - "role": "system", - "content": system_rule, - } - ] - - while True: - user_input = read_input("\nYou: ") - - if user_input is None or user_input.lower() in {"exit", "quit"}: - return 0 - - if not user_input: - print("Please enter a message.") - continue - - messages.append( - { - "role": "user", - "content": user_input, - } - ) - - try: - response = model.complete( - messages, - on_text=lambda text: print(text, end="", flush=True), - ) - except OpenAIError as error: - messages.pop() - print(f"\nRequest failed: {error}") - continue - - 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": response.content, - } - ) - messages = trim_history(messages, settings.memory_keep_count) - +from llm.chat import run_chat if __name__ == "__main__": - raise SystemExit(main()) + raise SystemExit(run_chat()) diff --git a/tests/test_main.py b/tests/test_main.py index 11b5e2b..bdb78c3 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -1,6 +1,6 @@ import unittest -from main import trim_history +from llm.chat import trim_history class TrimHistoryTests(unittest.TestCase): diff --git a/tests/test_settings.py b/tests/test_settings.py index 7f81118..d1d18ce 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -1,6 +1,6 @@ import unittest -from settings import Settings +from llm.settings import Settings class SettingsTests(unittest.TestCase):