Add anthropic support

Author: hupe1980

aisploit/core/model.py

@@ -13,6 +13,12 @@ class BaseLLM(Runnable[LanguageModelInput, str]):
 class BaseChatModel(Runnable[LanguageModelInput, BaseMessage]):
     @abstractmethod
     def supports_functions(self) -> bool:
+        """
+        Check if the model supports additional functions beyond basic chat.
+
+        Returns:
+            bool: True if the model supports additional functions, False otherwise.
+        """
         pass
 
 

aisploit/core/report.py

@@ -32,7 +32,7 @@ def __iter__(self):
 
     def __len__(self):
         return len(self._entries)
-    
+
     def __getitem__(self, index: int) -> T:
         """Get an entry from the report by index."""
         return self._entries[index]

aisploit/model/__init__.py

@@ -1,9 +1,11 @@
 from .bedrock_chat import BedrockChat
+from .chat_anthropic import ChatAnthropic
 from .chat_ollama import ChatOllama
 from .chat_openai import ChatOpenAI
 
 __all__ = [
     "BedrockChat",
+    "ChatAnthropic",
     "ChatOllama",
     "ChatOpenAI",
 ]

aisploit/model/bedrock_chat.py

@@ -16,4 +16,10 @@ def __init__(
         )
 
     def supports_functions(self) -> bool:
+        """
+        Check if the model supports additional functions beyond basic chat.
+
+        Returns:
+            bool: True if the model supports additional functions, False otherwise.
+        """
         return False

aisploit/model/chat_anthropic.py

@@ -0,0 +1,42 @@
+from typing import Optional
+from langchain_core.utils.utils import convert_to_secret_str
+from langchain_anthropic import ChatAnthropic as LangchainChatAnthropic
+
+from ..core import BaseChatModel
+
+
+class ChatAnthropic(LangchainChatAnthropic, BaseChatModel):
+    """A chat model based on Anthropic's language generation technology."""
+
+    def __init__(
+        self,
+        *,
+        api_key: Optional[str],
+        model_name: str = "claude-3-opus-20240229",
+        temperature: float = 1.0,
+        **kwargs,
+    ) -> None:
+        """
+        Initialize the ChatAnthropic instance.
+
+        Args:
+            api_key (str or None): The API key for accessing the Anthropic API.
+            model_name (str): The name of the language model to use.
+            temperature (float): The temperature parameter controlling the randomness of the generated text.
+            **kwargs: Additional keyword arguments to be passed to the base class constructor.
+        """
+        super().__init__(
+            anthropic_api_key=convert_to_secret_str(api_key) if api_key else None,
+            model_name=model_name,
+            temperature=temperature,
+            **kwargs,
+        )
+
+    def supports_functions(self) -> bool:
+        """
+        Check if the model supports additional functions beyond basic chat.
+
+        Returns:
+            bool: True if the model supports additional functions, False otherwise.
+        """
+        return False

aisploit/model/chat_ollama.py

@@ -30,4 +30,10 @@ def __init__(
         )
 
     def supports_functions(self) -> bool:
+        """
+        Check if the model supports additional functions beyond basic chat.
+
+        Returns:
+            bool: True if the model supports additional functions, False otherwise.
+        """
         return False

aisploit/model/chat_openai.py

@@ -38,4 +38,10 @@ def __init__(
         )
 
     def supports_functions(self) -> bool:
+        """
+        Check if the model supports additional functions beyond basic chat.
+
+        Returns:
+            bool: True if the model supports additional functions, False otherwise.
+        """
         return True

poetry.lock

@@ -38,6 +38,7 @@ torch = "^2.2.2"
 jinja2 = "^3.1.3"
 ipython = "^8.23.0"
 imapclient = "^3.0.1"
+langchain-anthropic = "^0.1.6"
 
 [tool.poetry.group.dev.dependencies]
 chromadb = "^0.4.23"