Added Anthropic models. Removed old style API class and tests (#10)

Author: christo-olivier

.vscode/launch.json

@@ -0,0 +1,15 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python Debugger: Current File",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "${file}",
+            "console": "integratedTerminal"
+        }
+    ]
+}

README.md

@@ -6,10 +6,11 @@
 ### Modelsmith is a Python library that allows you to get structured responses in the form of Pydantic models and Python types from Google Vertex AI and OpenAI models.
 
 Currently it allows you to use the following classes of model:
-- __OpenAIModel__ (most commonly used with `gpt-3.5-turbo`, `gpt-4` and `gpt-4o`)
-- __VertexAIChatModel__ (most commonly used with `chat-bison`)
-- __VertexAITextGenerationModel__ (most commonly used with `text-bison`)
-- __VertexAIGenerativeModel__ (most commonly used with `gemini-pro`)
+- __AnthropicModel__ (used with Anthropic's set of models such as `claude-3-haiku`, `claude-3-sonnet`, `claude-3-opus` and `claude-3_5-sonnet`)
+- __OpenAIModel__ (used with OpenAI's set of models such as `gpt-3.5-turbo`, `gpt-4` and `gpt-4o`)
+- __VertexAIChatModel__ (used with Google Vertex AI's chat models such as `chat-bison`)
+- __VertexAITextGenerationModel__ (used with Google Vertex AI's text generation models such as `text-bison`)
+- __VertexAIGenerativeModel__ (used with Google Vertex AI's generative models such as `gemini-pro`)
 
 Modelsmith provides a unified interface over all of these. It has been designed to be extensible and can adapt to other models in the future.
 
@@ -126,10 +127,10 @@ response = forge.generate("I have lived in Irvine, CA and Dallas TX")
 print(response)  # [City(city='Irvine', state='CA'), City(city='Dallas', state='TX')]
 ```
 
-If we want to use an OpenAI model the same applies. Simply select the appropriate model class, specify which OpenAI model to use (in this case `gpt-4o`), and pass it to the `Forge` instance.
+If we want to use an Anthropic model the same applies. Simply select the appropriate model class, specify which Anthropic model to use (in this case `claude-3-haiku-20240307`), and pass it to the `Forge` instance.
 
 ```python
-from modelsmith import Forge, OpenAIModel  # import the correct class
+from modelsmith import Forge, AnthropicModel  # import the correct class
 from pydantic import BaseModel, Field
 
 
@@ -138,9 +139,9 @@ class City(BaseModel):
     state: str = Field(description="2-letter abbreviation of the state")
 
 
-# text-bison instead of gemini-pro
+# Anthropic's claude-3-haiku-20240307 instead of gemini-pro
 forge = Forge(
-    model=OpenAIModel("gpt-4o"),
+    model=AnthropicModel("claude-3-haiku-20240307"),
     response_model=list[City],
 )
 

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "modelsmith"
-version = "0.5.0"
+version = "0.6.0"
 description = "Get Pydantic models and Python types as LLM responses from Google Vertex AI and OpenAI models."
 authors = ["Christo Olivier <mail@christoolivier.com>"]
 maintainers = ["Christo Olivier <mail@christoolivier.com>"]
@@ -18,6 +18,7 @@ google-cloud-aiplatform = "^1.43.0"
 pydantic = "^2.6.4"
 jinja2 = "^3.1.3"
 openai = "^1.35.7"
+anthropic = "^0.30.1"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.1.1"

pyproject.toml

@@ -1,6 +1,7 @@
 from .exceptions import ModelNotDerivedError, PatternNotFound
 from .forge import Forge
 from .language_models import (
+    AnthropicModel,
     OpenAIModel,
     VertexAIChatModel,
     VertexAIGenerativeModel,
@@ -9,6 +10,7 @@
 from .prompt import Prompt
 
 __all__ = [
+    "AnthropicModel",
     "Forge",
     "ModelNotDerivedError",
     "OpenAIModel",

src/modelsmith/__init__.py

@@ -3,11 +3,6 @@
 
 from pydantic import ValidationError
 from tenacity import RetryError, Retrying, stop_after_attempt
-from vertexai.generative_models import GenerativeModel
-from vertexai.language_models import (
-    ChatModel,
-    TextGenerationModel,
-)
 
 from modelsmith.enums import ResponseModelType
 from modelsmith.exceptions import (
@@ -16,11 +11,11 @@
     PatternNotFound,
 )
 from modelsmith.language_models import (
+    AnthropicModel,
     OpenAIModel,
     VertexAIChatModel,
     VertexAIGenerativeModel,
     VertexAITextGenerationModel,
-    _LanguageModelWrapper,
 )
 from modelsmith.prompt import Prompt
 from modelsmith.response_model import ResponseModel
@@ -40,9 +35,7 @@ class Forge(Generic[T]):
     def __init__(
         self,
         *,
-        model: ChatModel
-        | GenerativeModel
-        | TextGenerationModel
+        model: AnthropicModel
         | VertexAIChatModel
         | VertexAIGenerativeModel
         | VertexAITextGenerationModel
@@ -70,13 +63,7 @@ def __init__(
                                  return a response in the form of the response_model. If
                                  False, return None.
         """
-        # check if a Vertex AI model is being passed directly or if the new
-        # Wrapper classes are being used.
-        self.model = (
-            _LanguageModelWrapper(model)
-            if isinstance(model, (ChatModel, GenerativeModel, TextGenerationModel))
-            else model
-        )
+        self.model = model
         self.response_model = ResponseModel(response_model)
         self.prompt = Prompt(prompt)
 

src/modelsmith/forge.py

@@ -1,7 +1,7 @@
-import warnings
 from abc import ABC, abstractmethod
-from typing import Any, Callable
+from typing import Any
 
+from anthropic import Anthropic
 from openai import OpenAI
 from vertexai.generative_models import GenerationResponse, GenerativeModel
 from vertexai.language_models import (
@@ -10,11 +10,7 @@
     TextGenerationResponse,
 )
 
-DEPRECATION_WARNING = (
-    "Using VertexAI classes directly is deprecated and will be removed in "
-    "a future release. Please import the classes from the `language_model` "
-    "module instead."
-)
+DEFAULT_MAX_TOKENS = 1024
 
 
 class BaseLanguageModel(ABC):
@@ -30,6 +26,42 @@ def send(self, input: str, model_settings: dict[str, Any] | None = None) -> str:
         pass
 
 
+class AnthropicModel(BaseLanguageModel):
+    """
+    Class that wraps the Anthropic API to handle sending inputs and receiving outputs.
+    """
+
+    def __init__(self, model_name: str, api_key: str | None = None) -> None:
+        """
+        Create a new synchronous anthropic client instance.
+        """
+        self.model_name = model_name
+        self._client = Anthropic(api_key=api_key)
+
+    def send(self, input: str, model_settings: dict[str, Any] | None = None) -> str:
+        """
+        Send the input to the LLM using the correct method from the underlying model.
+        Return the text response from the LLM.
+
+        :param input: The input string to send to the LLM.
+        :param model_settings: The dictionary containing the model's settings.
+        :return: The response from the LLM.
+        """
+        # If `max_tokens` not provided in model settings then set it to the default
+        # of 1024
+        model_settings = model_settings or {}
+        if "max_tokens" not in model_settings:
+            model_settings["max_tokens"] = DEFAULT_MAX_TOKENS
+
+        response = self._client.messages.create(
+            model=self.model_name,
+            messages=[{"role": "user", "content": input}],
+            **(model_settings or {}),
+        )
+
+        return response.content[0].text
+
+
 class OpenAIModel(BaseLanguageModel):
     """
     Class that wraps the OpenAI API to handle sending inputs and receiving outputs.
@@ -153,55 +185,3 @@ def send(
         """
         response = self.model.predict(input, **(model_settings or {}))
         return response.text
-
-
-class _LanguageModelWrapper:
-    """
-    Class that wraps the LLM model to handle sending inputs and receiving outputs.
-    """
-
-    def __init__(
-        self, model: ChatModel | GenerativeModel | TextGenerationModel
-    ) -> None:
-        self.model = model
-        self._llm_send_method = self._init_llm_send_method()
-
-    def _init_llm_send_method(self) -> Callable:
-        """
-        Set the method to use to send the user input to the LLM.
-        """
-        warnings.warn(
-            DEPRECATION_WARNING,
-            DeprecationWarning,
-            stacklevel=2,
-        )
-
-        if isinstance(self.model, TextGenerationModel):
-            return self.model.predict
-
-        if isinstance(self.model, GenerativeModel):
-            return self.model.generate_content
-
-        if isinstance(self.model, ChatModel):
-            chat_session = self.model.start_chat()
-            return chat_session.send_message
-
-        raise TypeError(
-            "The model type must be ChatModel, TextGenerationModel or GenerativeModel"
-        )
-
-    def send(self, input: str, model_settings: dict[str, Any] | None = None) -> str:
-        """
-        Send the input to the LLM using the correct method from the underlying model.
-        Return the response from the LLM.
-
-        :param input: The input string to send to the LLM.
-        :param model_settings: The dictionary containing the model's settings.
-        :return: The response from the LLM.
-        """
-        # For a GenerativeModel the function signature differs from the other models
-        if isinstance(self.model, GenerativeModel):
-            response = self._llm_send_method(input, generation_config=model_settings)
-        else:
-            response = self._llm_send_method(input, **(model_settings or {}))
-        return response.text

src/modelsmith/language_models.py

@@ -1,26 +1,38 @@
 import pytest
 from modelsmith.language_models import (
+    AnthropicModel,
     OpenAIModel,
     VertexAIChatModel,
     VertexAIGenerativeModel,
     VertexAITextGenerationModel,
 )
-from vertexai.generative_models import GenerativeModel
-from vertexai.language_models import ChatModel, TextGenerationModel
 
 MODEL_INSTANCE_PARAMS = [
     pytest.param(
-        TextGenerationModel.from_pretrained("text-bison"), id="text_model_old_style"
+        AnthropicModel("claude-3-haiku-20240307"), id="anthropic_claude_3_haiku_model"
     ),
-    pytest.param(VertexAITextGenerationModel("text-bison"), id="text_model"),
-    pytest.param(ChatModel.from_pretrained("chat-bison"), id="chat_model_old_style"),
-    pytest.param(VertexAIChatModel("chat-bison"), id="chat_model"),
-    pytest.param(VertexAIGenerativeModel("gemini-1.0-pro"), id="gemini_1_0_pro"),
-    pytest.param(VertexAIGenerativeModel("gemini-1.5-pro"), id="gemini_1_5_pro"),
-    pytest.param(GenerativeModel("gemini-1.5-flash"), id="gemini_1_5_flash_old_style"),
-    pytest.param(VertexAIGenerativeModel("gemini-1.5-flash"), id="gemini_1_5_flash"),
-    pytest.param(OpenAIModel("gpt-3.5-turbo"), id="gpt-3.5-turbo"),
-    pytest.param(OpenAIModel("gpt-4o"), id="gpt-4o"),
+    pytest.param(
+        AnthropicModel("claude-3-opus-20240229"), id="anthropic_claude_3_opus_model"
+    ),
+    pytest.param(
+        AnthropicModel("claude-3-5-sonnet-20240620"),
+        id="anthropic_claude_3.5_sonnet_model",
+    ),
+    pytest.param(
+        VertexAITextGenerationModel("text-bison"), id="vertexai_text_bison_model"
+    ),
+    pytest.param(VertexAIChatModel("chat-bison"), id="vertexai_chat_bison_model"),
+    pytest.param(
+        VertexAIGenerativeModel("gemini-1.0-pro"), id="vertexai_gemini_1_0_pro"
+    ),
+    pytest.param(
+        VertexAIGenerativeModel("gemini-1.5-pro"), id="vertexai_gemini_1_5_pro"
+    ),
+    pytest.param(
+        VertexAIGenerativeModel("gemini-1.5-flash"), id="vertexai_gemini_1_5_flash"
+    ),
+    pytest.param(OpenAIModel("gpt-3.5-turbo"), id="openai_gpt-3.5-turbo"),
+    pytest.param(OpenAIModel("gpt-4o"), id="openai_gpt-4o"),
 ]
 
 MODEL_SETTINGS_PARAMS = [