Skip to content

Commit

Permalink
Add anthropic support
Browse files Browse the repository at this point in the history
  • Loading branch information
hupe1980 committed Apr 9, 2024
1 parent e573a72 commit 493c4ba
Show file tree
Hide file tree
Showing 9 changed files with 111 additions and 2 deletions.
6 changes: 6 additions & 0 deletions aisploit/core/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
2 changes: 1 addition & 1 deletion aisploit/core/report.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
2 changes: 2 additions & 0 deletions aisploit/model/__init__.py
Original file line number Diff line number Diff line change
@@ -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",
]
6 changes: 6 additions & 0 deletions aisploit/model/bedrock_chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
42 changes: 42 additions & 0 deletions aisploit/model/chat_anthropic.py
Original file line number Diff line number Diff line change
@@ -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
6 changes: 6 additions & 0 deletions aisploit/model/chat_ollama.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
6 changes: 6 additions & 0 deletions aisploit/model/chat_openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
42 changes: 41 additions & 1 deletion poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down

0 comments on commit 493c4ba

Please sign in to comment.