Override default max tokens for Anthropic and Groq clients

getzep · Sep 22, 2024 · ebd7d4b · ebd7d4b
1 parent ecfc45f
commit ebd7d4b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 10 deletions.
diff --git a/graphiti_core/llm_client/anthropic_client.py b/graphiti_core/llm_client/anthropic_client.py
@@ -36,13 +36,11 @@
 class AnthropicClient(LLMClient):
     def __init__(self, config: LLMConfig | None = None, cache: bool = False):
         if config is None:
-            config = LLMConfig()
+            config = LLMConfig(max_tokens=DEFAULT_MAX_TOKENS)
+        elif config.max_tokens is None:
+            config.max_tokens = DEFAULT_MAX_TOKENS
         super().__init__(config, cache)
 
-        # Override the default max tokens for Anthropic
-        if self.max_tokens is None:
-            self.max_tokens = DEFAULT_MAX_TOKENS
-
         self.client = AsyncAnthropic(
             api_key=config.api_key,
             # we'll use tenacity to retry

diff --git a/graphiti_core/llm_client/groq_client.py b/graphiti_core/llm_client/groq_client.py
@@ -37,13 +37,11 @@
 class GroqClient(LLMClient):
     def __init__(self, config: LLMConfig | None = None, cache: bool = False):
         if config is None:
-            config = LLMConfig()
+            config = LLMConfig(max_tokens=DEFAULT_MAX_TOKENS)
+        elif config.max_tokens is None:
+            config.max_tokens = DEFAULT_MAX_TOKENS
         super().__init__(config, cache)
 
-        # Override the default max tokens for Groq
-        if self.max_tokens is None:
-            self.max_tokens = DEFAULT_MAX_TOKENS
-
         self.client = AsyncGroq(api_key=config.api_key)
 
     def get_embedder(self) -> typing.Any: