From 6d40dfaa6a5c3c4492b1dd308e0d722df934e147 Mon Sep 17 00:00:00 2001
From: Michael Jarvis <mjarvis@ep.com>
Date: Sun, 12 Jan 2025 11:21:17 +0000
Subject: [PATCH] Increase default max tokens for Claude and add environment
 variable to configure this

---
 assistants/ai/anthropic.py    | 3 ++-
 assistants/cli/constants.py   | 2 ++
 assistants/config/__init__.py | 4 ++++
 3 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/assistants/ai/anthropic.py b/assistants/ai/anthropic.py
index bd5cf62..837b527 100644
--- a/assistants/ai/anthropic.py
+++ b/assistants/ai/anthropic.py
@@ -12,6 +12,7 @@
 
 from assistants.ai.memory import MemoryMixin
 from assistants.ai.types import MessageData
+from assistants.cli.constants import CLAUDE_CLI_MAX_TOKENS
 from assistants.config import environment
 from assistants.lib.exceptions import ConfigError
 
@@ -34,7 +35,7 @@ class Claude(MemoryMixin):
     def __init__(
         self,
         model: str,
-        max_tokens: int = 1024,
+        max_tokens: int = CLAUDE_CLI_MAX_TOKENS,
         max_memory: int = 50,
         api_key: Optional[str] = environment.ANTHROPIC_API_KEY,
     ) -> None:
diff --git a/assistants/cli/constants.py b/assistants/cli/constants.py
index c450a7a..f4e14e9 100644
--- a/assistants/cli/constants.py
+++ b/assistants/cli/constants.py
@@ -14,3 +14,5 @@
 
 Anything else you type will be sent to the assistant for processing.\
 """
+
+CLAUDE_CLI_MAX_TOKENS = 2048
diff --git a/assistants/config/__init__.py b/assistants/config/__init__.py
index 2455f76..e9927e4 100644
--- a/assistants/config/__init__.py
+++ b/assistants/config/__init__.py
@@ -2,6 +2,8 @@
 from enum import Enum
 from dataclasses import dataclass
 
+from assistants.cli.constants import CLAUDE_CLI_MAX_TOKENS
+
 
 class CustomKeyNames(str, Enum):
     ASSISTANTS_API_KEY_NAME: str = "ASSISTANTS_API_KEY_NAME"
@@ -35,6 +37,7 @@ class Config:
     ASSISTANT_INSTRUCTIONS: str
     ASSISTANT_NAME: str
     TELEGRAM_BOT_TOKEN: str
+    CLAUDE_MAX_TOKENS: str
 
 
 def get_config() -> Config:
@@ -50,6 +53,7 @@ def get_config() -> Config:
         ),
         ASSISTANT_NAME=os.environ.get("ASSISTANT_NAME", "DefaultAssistant"),
         TELEGRAM_BOT_TOKEN=os.environ.get("TG_BOT_TOKEN", None),
+        CLAUDE_MAX_TOKENS=os.environ.get("CLAUDE_MAX_TOKENS", CLAUDE_CLI_MAX_TOKENS),
     )