Add more bedrock models

nus-apr · Dec 6, 2024 · 8c4b839 · 8c4b839
1 parent 58c48b0
commit 8c4b839
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -306,9 +306,14 @@ The current list of supported models:
 |                | Claude 3 Haiku         | --model claude-3-haiku-20240307 |
 | Meta           | Llama 3 70B            | --model llama3:70b |
 |                | Llama 3 8B             | --model llama3     |
-| AWS            | Claude 3 Opus          | --model bedrock/anthropic.claude-3-opus-20240229-v1:0 |
+| AWS Bedrock    | Claude 3 Opus          | --model bedrock/anthropic.claude-3-opus-20240229-v1:0 |
 |                | Claude 3 Sonnet        | --model bedrock/anthropic.claude-3-sonnet-20240229-v1:0 |
 |                | Claude 3 Haiku         | --model bedrock/anthropic.claude-3-haiku-20240307-v1:0 |
+|                | Claude 3.5 Sonnet      | --model bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0 |
+|                | Nova Pro               | --model bedrock/us.amazon.nova-pro-v1:0 |
+|                | Nova Lite              | --model bedrock/us.amazon.nova-lite-v1:0 |
+|                | Nova Micro             | --model bedrock/us.amazon.nova-micro-v1:0 |
+| LiteLLM        | Any LiteLLM model      | --model litellm-generic-<MODEL_NAME_HERE> |
 | Groq           | Llama 3 8B             | --model groq/llama3-8b-8192 |
 |                | Llama 3 70B            | --model groq/llama3-70b-8192 |
 |                | Llama 2 70B            | --model groq/llama2-70b-4096 |

diff --git a/app/model/bedrock.py b/app/model/bedrock.py
@@ -117,9 +117,9 @@ def call(
             output_tokens = int(resp_usage.completion_tokens)
             cost = self.calc_cost(input_tokens, output_tokens)
 
-            self.model_cost.process_cost += cost
-            self.model_cost.process_input_tokens += input_tokens
-            self.model_cost.process_output_tokens += output_tokens
+            common.thread_cost.process_cost += cost
+            common.thread_cost.process_input_tokens += input_tokens
+            common.thread_cost.process_output_tokens += output_tokens
 
             first_resp_choice = response.choices[0]
             assert isinstance(first_resp_choice, Choices)

diff --git a/app/model/register.py b/app/model/register.py
@@ -36,6 +36,10 @@ def register_all_models() -> None:
     common.register_model(bedrock.AnthropicClaude3Opus())
     common.register_model(bedrock.AnthropicClaude3Sonnet())
     common.register_model(bedrock.AnthropicClaude3Haiku())
+    common.register_model(bedrock.AmazonNovaLitev1())
+    common.register_model(bedrock.AmazonNovaProv1())
+    common.register_model(bedrock.AmazonNovaMicrov1())
+    common.register_model(bedrock.AnthropicClaude35Sonnet())
 
     common.register_model(ollama.Llama3_8B())
     common.register_model(ollama.Llama3_70B())