Add support for Claude models.

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__":

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."""
 

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."""

requirements.in

@@ -1,3 +1,4 @@
+anthropic==0.31.2
 chardet==5.2.0
 cxxfilt==0.3.0
 GitPython==3.1.43

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