merged

kagisearch · May 13, 2024 · 11093ad · 11093ad
2 parents 53a923f + 26676df
commit 11093ad
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 7 deletions.
diff --git a/llms/providers/anthropic.py b/llms/providers/anthropic.py
@@ -147,7 +147,7 @@ def _prepare_message_inputs(
             "messages": messages,
             "system": system_message,
             "temperature": temperature,
-            "max_tokens": max_tokens,
+            "max_tokens_to_sample": max_tokens,
             "stop_sequences": stop_sequences,
         }
         return model_inputs

diff --git a/llms/providers/bedrock_anthropic.py b/llms/providers/bedrock_anthropic.py
@@ -3,7 +3,7 @@
 import os
 from typing import Union
 
-import anthropic_bedrock
+from anthropic import AnthropicBedrock, AsyncAnthropicBedrock
 
 from .anthropic import AnthropicProvider
 
@@ -17,6 +17,8 @@ class BedrockAnthropicProvider(AnthropicProvider):
             "token_limit": 100_000,
         },
         "anthropic.claude-v2": {"prompt": 11.02, "completion": 32.68, "token_limit": 100_000},
+        "anthropic.claude-3-haiku-20240307-v1:0": {"prompt": 0.25, "completion": 1.25, "token_limit": 200_000, "output_limit": 4_096},
+        "anthropic.claude-3-sonnet-20240229-v1:0": {"prompt": 3.00, "completion": 15, "token_limit": 200_000, "output_limit": 4_096},
     }
 
     def __init__(
@@ -39,7 +41,7 @@ def __init__(
 
         if client_kwargs is None:
             client_kwargs = {}
-        self.client = anthropic_bedrock.AnthropicBedrock(
+        self.client = AnthropicBedrock(
             aws_access_key=aws_access_key,
             aws_secret_key=aws_secret_key,
             aws_region=aws_region,
@@ -48,9 +50,13 @@ def __init__(
 
         if async_client_kwargs is None:
             async_client_kwargs = {}
-        self.async_client = anthropic_bedrock.AsyncAnthropicBedrock(
+        self.async_client = AsyncAnthropicBedrock(
             aws_access_key=aws_access_key,
             aws_secret_key=aws_secret_key,
             aws_region=aws_region,
             **async_client_kwargs,
         )
+
+    @property
+    def support_message_api(self):
+        return "claude-3" in self.model
diff --git a/setup.py b/setup.py
@@ -9,7 +9,7 @@
 
 setup(
     name="pyllms",
-    version="0.5.0",
+    version="0.5.1",
     description="Minimal Python library to connect to LLMs (OpenAI, Anthropic, Google Palm2/Vertex, Mistral, Ollama, AI21, Cohere, Aleph-Alpha, HuggingfaceHub), with a built-in model performance benchmark.",
     long_description=long_description,
     long_description_content_type="text/markdown",
@@ -20,14 +20,13 @@
         "openai>=1",
         "tiktoken",
         "anthropic>=0.18",
-        "anthropic_bedrock",
         "ai21",
         "cohere",
         "aleph-alpha-client",
         "huggingface_hub",
         "google-cloud-aiplatform",
         "prettytable",
-        "protobuf~=3.20.3",
+        "protobuf>=3.20.3",
         "grpcio~=1.54.2",
         "google-generativeai",
         "mistralai",