From 31f05f1e7630cca47d14a08a14e12d49d7be3b63 Mon Sep 17 00:00:00 2001
From: Abhishek Arya <aarya@google.com>
Date: Tue, 23 Jul 2024 22:44:00 -0700
Subject: [PATCH 1/4] Add support for Claude models.

---
 experimental/manual/prompter.py |  4 ++-
 llm_toolkit/models.py           | 62 +++++++++++++++++++++++++++++++++
 llm_toolkit/prompts.py          |  4 +++
 requirements.in                 |  1 +
 requirements.txt                | 35 ++++++++++++++++---
 5 files changed, 101 insertions(+), 5 deletions(-)

diff --git a/experimental/manual/prompter.py b/experimental/manual/prompter.py
index 58379d5441..7e455967bc 100644
--- a/experimental/manual/prompter.py
+++ b/experimental/manual/prompter.py
@@ -64,7 +64,9 @@ def setup_model() -> models.LLM:
 def construct_prompt() -> prompts.Prompt:
   with open(args.prompt, 'r') as prompt_file:
     content = prompt_file.read()
-  return model.prompt_type()(initial=content)
+  prompt = model.prompt_type()()
+  prompt.add_problem(content)
+  return prompt
 
 
 if __name__ == "__main__":
diff --git a/llm_toolkit/models.py b/llm_toolkit/models.py
index a88a689076..46923bc0a0 100644
--- a/llm_toolkit/models.py
+++ b/llm_toolkit/models.py
@@ -27,6 +27,7 @@
 from abc import abstractmethod
 from typing import Any, Callable, Optional, Type
 
+import anthropic
 import openai
 import tiktoken
 import vertexai
@@ -263,6 +264,67 @@ class GPT4o(GPT):
   name = 'gpt-4o'
 
 
+class Claude(LLM):
+  """Anthropic's Claude model encapsulator."""
+
+  _max_output_tokens = 4096
+  context_window = 200000
+
+  name = "claude-3-haiku@20240307"
+
+  # ================================ Prompt ================================ #
+  def estimate_token_num(self, text) -> int:
+    """Estimates the number of tokens in |text|."""
+    client = anthropic.Client()
+    return client.count_tokens(text)
+
+  def prompt_type(self) -> type[prompts.Prompt]:
+    """Returns the expected prompt type."""
+    return prompts.ClaudePrompt
+
+  # ============================== Generation ============================== #
+  def query_llm(self,
+                prompt: prompts.Prompt,
+                response_dir: str,
+                log_output: bool = False) -> None:
+    """Queries Claude's API and stores response in |response_dir|."""
+    if self.ai_binary:
+      logger.info(f'Claude does not use local AI binary: {self.ai_binary}')
+    if self.temperature_list:
+      logger.info(
+          f'Claude does not allow temperature list: {self.temperature_list}')
+
+    vertex_ai_locations = os.getenv('VERTEX_AI_LOCATIONS',
+                                    'europe-west1').split(',')
+    project_id = os.getenv('GOOGLE_CLOUD_PROJECT', 'oss-fuzz')
+    region = random.sample(vertex_ai_locations, 1)[0]
+    client = anthropic.AnthropicVertex(region=region, project_id=project_id)
+
+    completion = self.with_retry_on_error(
+        lambda: client.messages.create(max_tokens=self._max_output_tokens,
+                                       messages=prompt.get(),
+                                       model=self.name,
+                                       temperature=self.temperature),
+        anthropic.AnthropicError)
+    if log_output:
+      logger.info(completion)
+    for index, choice in enumerate(completion.content):
+      content = choice.text
+      self._save_output(index, content, response_dir)
+
+
+class ClaudeOpusV3(Claude):
+  """Claude Opus 3."""
+
+  name = 'claude-3-opus@20240229'
+
+
+class ClaudeSonnetV3D5(Claude):
+  """Claude Sonnet 3.5."""
+
+  name = 'claude-3-5-sonnet@20240620'
+
+
 class GoogleModel(LLM):
   """Generic Google model."""
 
diff --git a/llm_toolkit/prompts.py b/llm_toolkit/prompts.py
index 3a57af3ed4..1f5a8b1142 100644
--- a/llm_toolkit/prompts.py
+++ b/llm_toolkit/prompts.py
@@ -133,3 +133,7 @@ def save(self, location: str) -> None:
     """Saves the prompt to a filelocation."""
     with open(location, 'w+') as prompt_file:
       json.dump(self._prompt, prompt_file)
+
+
+class ClaudePrompt(OpenAIPrompt):
+  """Claude style structured prompt."""
diff --git a/requirements.in b/requirements.in
index c1b84b5756..7a2b784af1 100644
--- a/requirements.in
+++ b/requirements.in
@@ -1,3 +1,4 @@
+anthropic==0.31.2
 chardet==5.2.0
 cxxfilt==0.3.0
 GitPython==3.1.43
diff --git a/requirements.txt b/requirements.txt
index bf1b1ccb32..67c74fd485 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,8 +6,11 @@
 #
 annotated-types==0.7.0
     # via pydantic
+anthropic==0.31.2
+    # via -r requirements.in
 anyio==4.4.0
     # via
+    #   anthropic
     #   httpx
     #   openai
 astroid==3.2.2
@@ -28,9 +31,15 @@ cxxfilt==0.3.0
 dill==0.3.8
     # via pylint
 distro==1.9.0
-    # via openai
+    # via
+    #   anthropic
+    #   openai
 docstring-parser==0.16
     # via google-cloud-aiplatform
+filelock==3.15.4
+    # via huggingface-hub
+fsspec==2024.6.1
+    # via huggingface-hub
 gitdb==4.0.11
     # via gitpython
 gitpython==3.1.43
@@ -90,7 +99,11 @@ h11==0.14.0
 httpcore==1.0.5
     # via httpx
 httpx==0.27.0
-    # via openai
+    # via
+    #   anthropic
+    #   openai
+huggingface-hub==0.24.1
+    # via tokenizers
 idna==3.7
     # via
     #   anyio
@@ -102,6 +115,8 @@ isort==5.13.2
     # via pylint
 jinja2==3.1.4
     # via -r requirements.in
+jiter==0.5.0
+    # via anthropic
 markupsafe==2.1.5
     # via jinja2
 mccabe==0.7.0
@@ -118,6 +133,7 @@ packaging==24.1
     # via
     #   google-cloud-aiplatform
     #   google-cloud-bigquery
+    #   huggingface-hub
 pandas==2.2.2
     # via -r requirements.in
 platformdirs==4.2.2
@@ -146,6 +162,7 @@ pyasn1-modules==0.4.0
     # via google-auth
 pydantic==2.8.2
     # via
+    #   anthropic
     #   google-cloud-aiplatform
     #   openai
 pydantic-core==2.20.1
@@ -161,7 +178,9 @@ python-dateutil==2.9.0.post0
 pytz==2024.1
     # via pandas
 pyyaml==6.0.1
-    # via -r requirements.in
+    # via
+    #   -r requirements.in
+    #   huggingface-hub
 regex==2024.5.15
     # via tiktoken
 requests==2.32.0
@@ -170,6 +189,7 @@ requests==2.32.0
     #   google-api-core
     #   google-cloud-bigquery
     #   google-cloud-storage
+    #   huggingface-hub
     #   tiktoken
 rsa==4.9
     # via google-auth
@@ -181,19 +201,26 @@ smmap==5.0.1
     # via gitdb
 sniffio==1.3.1
     # via
+    #   anthropic
     #   anyio
     #   httpx
     #   openai
 tiktoken==0.7.0
     # via -r requirements.in
+tokenizers==0.19.1
+    # via anthropic
 tomli==2.0.1
     # via yapf
 tomlkit==0.12.5
     # via pylint
 tqdm==4.66.4
-    # via openai
+    # via
+    #   huggingface-hub
+    #   openai
 typing-extensions==4.12.2
     # via
+    #   anthropic
+    #   huggingface-hub
     #   openai
     #   pydantic
     #   pydantic-core

From 42efe6d591424277d74a0a283ae8a6c23e81dbf5 Mon Sep 17 00:00:00 2001
From: Abhishek Arya <aarya@google.com>
Date: Wed, 24 Jul 2024 06:03:30 -0700
Subject: [PATCH 2/4] Address review comments.

---
 llm_toolkit/models.py | 25 ++++++++++++++++++-------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/llm_toolkit/models.py b/llm_toolkit/models.py
index 46923bc0a0..bf36c0a7b4 100644
--- a/llm_toolkit/models.py
+++ b/llm_toolkit/models.py
@@ -231,7 +231,7 @@ def query_llm(self,
                 log_output: bool = False) -> None:
     """Queries OpenAI's API and stores response in |response_dir|."""
     if self.ai_binary:
-      logger.info(f'OpenAI does not use local AI binary: {self.ai_binary}')
+      raise ValueError(f'OpenAI does not use local AI binary: {self.ai_binary}')
     if self.temperature_list:
       logger.info(
           f'OpenAI does not allow temperature list: {self.temperature_list}')
@@ -268,10 +268,9 @@ class Claude(LLM):
   """Anthropic's Claude model encapsulator."""
 
   _max_output_tokens = 4096
+  _vertex_ai_model = ''
   context_window = 200000
 
-  name = "claude-3-haiku@20240307"
-
   # ================================ Prompt ================================ #
   def estimate_token_num(self, text) -> int:
     """Estimates the number of tokens in |text|."""
@@ -282,6 +281,9 @@ def prompt_type(self) -> type[prompts.Prompt]:
     """Returns the expected prompt type."""
     return prompts.ClaudePrompt
 
+  def get_model(self) -> Any:
+    return self._vertex_ai_model
+
   # ============================== Generation ============================== #
   def query_llm(self,
                 prompt: prompts.Prompt,
@@ -289,7 +291,7 @@ def query_llm(self,
                 log_output: bool = False) -> None:
     """Queries Claude's API and stores response in |response_dir|."""
     if self.ai_binary:
-      logger.info(f'Claude does not use local AI binary: {self.ai_binary}')
+      raise ValueError(f'Claude does not use local AI binary: {self.ai_binary}')
     if self.temperature_list:
       logger.info(
           f'Claude does not allow temperature list: {self.temperature_list}')
@@ -303,7 +305,7 @@ def query_llm(self,
     completion = self.with_retry_on_error(
         lambda: client.messages.create(max_tokens=self._max_output_tokens,
                                        messages=prompt.get(),
-                                       model=self.name,
+                                       model=self.get_model(),
                                        temperature=self.temperature),
         anthropic.AnthropicError)
     if log_output:
@@ -313,16 +315,25 @@ def query_llm(self,
       self._save_output(index, content, response_dir)
 
 
+class ClaudeHaikuV3(Claude):
+  """Claude Haiku 3."""
+
+  name = 'claude-3-haiku'
+  _vertex_ai_model = 'claude-3-haiku@20240307'
+
+
 class ClaudeOpusV3(Claude):
   """Claude Opus 3."""
 
-  name = 'claude-3-opus@20240229'
+  name = 'claude-3-opus'
+  _vertex_ai_model = 'claude-3-opus@20240229'
 
 
 class ClaudeSonnetV3D5(Claude):
   """Claude Sonnet 3.5."""
 
-  name = 'claude-3-5-sonnet@20240620'
+  name = 'claude-3-5-sonnet'
+  _vertex_ai_model = 'claude-3-5-sonnet@20240620'
 
 
 class GoogleModel(LLM):

From 9f41f6169544c45873415cd41e9156e49e05d92d Mon Sep 17 00:00:00 2001
From: Abhishek Arya <aarya@google.com>
Date: Wed, 24 Jul 2024 06:07:38 -0700
Subject: [PATCH 3/4] Fix return type.

---
 llm_toolkit/models.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llm_toolkit/models.py b/llm_toolkit/models.py
index bf36c0a7b4..6d2a673e02 100644
--- a/llm_toolkit/models.py
+++ b/llm_toolkit/models.py
@@ -281,7 +281,7 @@ def prompt_type(self) -> type[prompts.Prompt]:
     """Returns the expected prompt type."""
     return prompts.ClaudePrompt
 
-  def get_model(self) -> Any:
+  def get_model(self) -> str:
     return self._vertex_ai_model
 
   # ============================== Generation ============================== #

From 19b4ac179fa00f2fe20bcda1c15630b04068caf1 Mon Sep 17 00:00:00 2001
From: Abhishek Arya <aarya@google.com>
Date: Wed, 24 Jul 2024 07:50:32 -0700
Subject: [PATCH 4/4] Add vertex_ai prefix since these are invoked via VertexAI

---
 llm_toolkit/models.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/llm_toolkit/models.py b/llm_toolkit/models.py
index 6d2a673e02..aaaba21f0f 100644
--- a/llm_toolkit/models.py
+++ b/llm_toolkit/models.py
@@ -318,21 +318,21 @@ def query_llm(self,
 class ClaudeHaikuV3(Claude):
   """Claude Haiku 3."""
 
-  name = 'claude-3-haiku'
+  name = 'vertex_ai_claude-3-haiku'
   _vertex_ai_model = 'claude-3-haiku@20240307'
 
 
 class ClaudeOpusV3(Claude):
   """Claude Opus 3."""
 
-  name = 'claude-3-opus'
+  name = 'vertex_ai_claude-3-opus'
   _vertex_ai_model = 'claude-3-opus@20240229'
 
 
 class ClaudeSonnetV3D5(Claude):
   """Claude Sonnet 3.5."""
 
-  name = 'claude-3-5-sonnet'
+  name = 'vertex_ai_claude-3-5-sonnet'
   _vertex_ai_model = 'claude-3-5-sonnet@20240620'