gmakstutis/graphiti
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Override default max tokens for Anthropic and Groq clients (#143)

Browse source 
* Override default max tokens for Anthropic and Groq clients

* Override default max tokens for Anthropic and Groq clients

* Override default max tokens for Anthropic and Groq clients
This commit is contained in:
Daniel Chalef 2024-09-22 11:33:54 -07:00 • committed by GitHub
parent d8c49c1c0a
commit 14d5ce0b36
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 10 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
graphiti_core/llm_client/anthropic_client.py
View file

@ -30,13 +30,17 @@ from .errors import RateLimitError
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
DEFAULT_MODEL = 'claude-3-5-sonnet-20240620' DEFAULT_MODEL = 'claude-3-5-sonnet-20240620'
DEFAULT_MAX_TOKENS = 8192
class AnthropicClient(LLMClient): class AnthropicClient(LLMClient):
def __init__(self, config: LLMConfig | None = None, cache: bool = False): def __init__(self, config: LLMConfig | None = None, cache: bool = False):
if config is None: 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) super().__init__(config, cache)
self.client = AsyncAnthropic( self.client = AsyncAnthropic(
api_key=config.api_key, api_key=config.api_key,
# we'll use tenacity to retry # we'll use tenacity to retry

6
graphiti_core/llm_client/groq_client.py
View file

@ -31,13 +31,17 @@ from .errors import RateLimitError
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
DEFAULT_MODEL = 'llama-3.1-70b-versatile' DEFAULT_MODEL = 'llama-3.1-70b-versatile'
DEFAULT_MAX_TOKENS = 2048
class GroqClient(LLMClient): class GroqClient(LLMClient):
def __init__(self, config: LLMConfig | None = None, cache: bool = False): def __init__(self, config: LLMConfig | None = None, cache: bool = False):
if config is None: 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) super().__init__(config, cache)
self.client = AsyncGroq(api_key=config.api_key) self.client = AsyncGroq(api_key=config.api_key)
def get_embedder(self) -> typing.Any: def get_embedder(self) -> typing.Any: