Generalizes GPT3CompletionModel to work with other providers, adds Anthropic

.github/workflows/run-tests.yml

@@ -40,6 +40,7 @@ jobs:
           poetry install
       - name: Run pytest
         env:
-          # set a placeholder OpenAI key (required by tests)
+          # set a placeholder API keys (required by tests)
           OPENAI_API_KEY: ABCD1234
+          ANTHROPIC_API_KEY: ABCD1234
         run: poetry run pytest

README.md

@@ -5,9 +5,15 @@ Check out the [manuscript about this tool](https://greenelab.github.io/manubot-g
 
 ## Supported Large Language Models (LLMs)
 
-We currently support OpenAI models only, and are working to add support for other models.
-[Our evaluations](https://github.com/pivlab/manubot-ai-editor-evals) show that `gpt-4-turbo` is in general the best model for revising academic manuscripts.
-Therefore, this is the default option.
+We internally use [LangChain](https://www.langchain.com/) to invoke models, which allows our tool to theoretically
+support whichever model providers LangChain supports. That said, we currently support OpenAI and Anthropic models only,
+and are working to add support for other model providers.
+
+When using OpenAI models, [our evaluations](https://github.com/pivlab/manubot-ai-editor-evals) show that `gpt-4-turbo`
+is in general the best model for revising academic manuscripts. Therefore, this is the default option for OpenAI.
+
+We are still evaluating the models for other providers as we add them, and will update this section accordingly
+as we complete our evaluations.
 
 ## Using in a Manubot manuscript
 
@@ -16,12 +22,17 @@ See their official docs for more info on [configuring GitHub Actions](https://do
 
 ### Setup
 
+First, you should decide which model provider you'll use. You can find details on how to set up each provider below:
+- **OpenAI:** you'll want to [make an OpenAI account](https://openai.com/api/) and [create an API key](https://platform.openai.com/api-keys).
+- **Anthropic:** you'll want to [make an Anthropic account](https://www.anthropic.com/api) and [create an API key](https://console.anthropic.com/settings/keys).
+
 Start with a manuscript repo [forked from Manubot rootstock](https://github.com/manubot/rootstock), then follow these steps:
 
 1. In your forks's "▶️ Actions" tab, enable GitHub Actions.
 1. In your fork's "⚙️ Settings" tab, give GitHub Actions workflows read/write permissions and allow them to create pull requests.
-1. If you haven't already, [make an OpenAI account](https://openai.com/api/) and [create an API key](https://platform.openai.com/api-keys).
-1. In your fork's "⚙️ Settings" tab, make a new Actions repository secret with the name `OPENAI_API_KEY` and paste in your API key as the secret.
+1. If you haven't already, follow the directions above to create an account and get an API key for your chosen model provider.
+1. In your fork's "⚙️ Settings" tab, make a new Actions repository secret with the name `<PROVIDER>_API_KEY` and paste in your API key as the secret. Replace `<PROVIDER>`
+with your provider's name in uppercase and without any special symbols, e.g. `OPENAI` for OpenAI and `ANTHROPIC` for Anthropic.
 
 ### Configuring prompts
 
@@ -62,10 +73,11 @@ First, install Manubot in a Python environment, e.g.:
 pip install --upgrade manubot[ai-rev]
 ```
 
-You also need to export an environment variable with your OpenAI API key, e.g.:
+You also need to export an environment variable with your model provider's API key, e.g.:
 
 ```bash
 export OPENAI_API_KEY=ABCD1234
+# export ANTHROPIC_API_KEY=ABCD1234 # if you were using anthropic
 ```
 
 You can also provide other environment variables that will change the behavior of the editor (such as revising certain files only).
@@ -82,7 +94,7 @@ manubot ai-revision --content-directory content/ --config-directory ci/
 The editor will revise each paragraph of your manuscript and write back the revised files in the same directory.
 Finally, (_assuming you are tracking changes to your manuscript with git_) you can review each change and either keep it (commit it) or reject it (revert it).
 
-Using the OpenAI API can sometimes incur costs.
+Using model providers' APIs can sometimes incur costs.
 If you're worried about this or otherwise want to test things out before hitting the real API, you can run a local "dry run" by with a "fake" model:
 
 ```bash
@@ -126,9 +138,11 @@ me = ManuscriptEditor(
 )
 
 # create a model to revise the manuscript
+# (if using another provider, e.g. anthropic, replace model_provider="openai" with model_provider="anthropic")
 model = GPT3CompletionModel(
     title=me.title,
     keywords=me.keywords,
+    model_provider="openai",
 )
 
 # create a temporary directory to store the revised manuscript

libs/manubot_ai_editor/env_vars.py

@@ -15,6 +15,9 @@
 # OpenAI API key to use
 OPENAI_API_KEY = "OPENAI_API_KEY"
 
+# Anthropic API key to use
+ANTHROPIC_API_KEY = "ANTHROPIC_API_KEY"
+
 # Language model to use. For example, "text-davinci-003", "gpt-3.5-turbo", "gpt-3.5-turbo-0301", etc
 # The tool currently supports the "chat/completions" and "completions" endpoints, and you can check
 # compatible models here: https://platform.openai.com/docs/models/model-endpoint-compatibility

libs/manubot_ai_editor/models.py

@@ -6,6 +6,7 @@
 import json
 
 from langchain_openai import OpenAI, ChatOpenAI
+from langchain_anthropic import ChatAnthropic, Anthropic
 from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
 
 from manubot_ai_editor import env_vars
@@ -119,6 +120,22 @@ def get_prompt(self, paragraph_text, section_name, resolved_prompt: str = None):
         return paragraph_text
 
 
+# specifies metadata for each model provider, e.g. OpenAI or Anthropic,
+# that are used in the GPT3CompletionModel to invoke the provider's API
+MODEL_PROVIDERS = {
+    "openai": {
+        "default_model_engine": "gpt-3.5-turbo",
+        "api_key_env_var": env_vars.OPENAI_API_KEY,
+        "clients": {"chat": ChatOpenAI, "completions": OpenAI},
+    },
+    "anthropic": {
+        "default_model_engine": "claude-3-haiku-20240307",
+        "api_key_env_var": env_vars.ANTHROPIC_API_KEY,
+        "clients": {"chat": ChatAnthropic, "completions": Anthropic},
+    },
+}
+
+
 class GPT3CompletionModel(ManuscriptRevisionModel):
     """
     Revises paragraphs using completion or chat completion models. Most of the parameters
@@ -130,8 +147,9 @@ def __init__(
         self,
         title: str,
         keywords: list[str],
-        openai_api_key: str = None,
-        model_engine: str = "gpt-3.5-turbo",
+        model_provider: str = "openai",
+        api_key: str = None,
+        model_engine: str = None,
         temperature: float = 0.5,
         presence_penalty: float = None,
         frequency_penalty: float = None,
@@ -141,18 +159,31 @@ def __init__(
     ):
         super().__init__()
 
+        # first, get metadata about the model provider
+        try:
+            provider_meta = MODEL_PROVIDERS[model_provider]
+        except KeyError:
+            raise ValueError(
+                f"Model provider '{model_provider}' not found; it must be one of {', '.join(MODEL_PROVIDERS.keys())}"
+            )
+
+        if model_engine is None:
+            model_engine = provider_meta["default_model_engine"]
+
         # make sure the OpenAI API key is set
-        if openai_api_key is None:
-            # attempt to get the OpenAI API key from the environment, since one
+        if api_key is None:
+            provider_key_env_var = provider_meta["api_key_env_var"]
+
+            # attempt to get the API key from the environment, since one
             # wasn't specified as an argument
-            openai_api_key = os.environ.get(env_vars.OPENAI_API_KEY, None)
+            api_key = os.environ.get(provider_key_env_var, None)
 
             # if it's *still* not set, bail
-            if openai_api_key is None or openai_api_key.strip() == "":
+            if api_key is None or api_key.strip() == "":
                 raise ValueError(
-                    f"OpenAI API key not found. Please provide it as parameter "
+                    f"API key for provider {model_provider} not found. Please provide it as parameter "
                     f"or set it as an the environment variable "
-                    f"{env_vars.OPENAI_API_KEY}"
+                    f"{provider_key_env_var}"
                 )
 
         if env_vars.LANGUAGE_MODEL in os.environ:
@@ -226,10 +257,16 @@ def __init__(
         # adjust options if chat endpoint was selected
         self.endpoint = "chat"
 
-        if model_engine.startswith(
-            ("text-davinci-", "text-curie-", "text-babbage-", "text-ada-")
-        ):
-            self.endpoint = "completions"
+        # switch to 'completions' endpoint for specific models, depending on the
+        # provider
+        if model_provider == "openai":
+            if model_engine.startswith(
+                ("text-davinci-", "text-curie-", "text-babbage-", "text-ada-")
+            ):
+                self.endpoint = "completions"
+        elif model_provider == "anthropic":
+            if model_engine.startswith(("claude-2",)):
+                self.endpoint = "completions"
 
         print(f"Language model: {model_engine}")
         print(f"Model endpoint used: {self.endpoint}")
@@ -252,15 +289,12 @@ def __init__(
 
         self.several_spaces_pattern = re.compile(r"\s+")
 
-        if self.endpoint == "chat":
-            client_cls = ChatOpenAI
-        else:
-            client_cls = OpenAI
+        client_cls = provider_meta["clients"][self.endpoint]
 
         # construct the OpenAI client after all the rest of
         # the settings above have been processed
         self.client = client_cls(
-            api_key=openai_api_key,
+            api_key=api_key,
             **self.model_parameters,
         )
 

pyproject.toml

@@ -28,6 +28,7 @@ packages = [ { include = "manubot_ai_editor", from = "libs" } ]
 python = ">=3.10,<4.0"
 langchain-core = "^0.3.6"
 langchain-openai = "^0.2.0"
+langchain-anthropic = "^0.3.0"
 pyyaml = "*"
 
 [tool.poetry.group.dev.dependencies]

tests/config_loader_fixtures/prompt_gpt3_e2e/ai-revision-prompts.yaml

@@ -1,6 +1,6 @@
 prompts:
   front_matter: Revise the following paragraph to include the keyword "testify" somewhere in the text.
-  abstract: Revise the following paragraph to include the keyword "bottle" somewhere in the text.
+  abstract: Revise the following paragraph to include the keyword "violin" somewhere in the text.
   introduction: Revise the following paragraph to include the keyword "wound" somewhere in the text.
   results: Revise the following paragraph to include the keyword "classroom" somewhere in the text.
   results_framework: Revise the following paragraph to include the keyword "secretary" somewhere in the text.