Add bedrock support

aisploit/embedding/__init__.py

@@ -1,7 +1,9 @@
+from .bedrock import BedrockEmbeddings
 from .ollama import OllamaEmbeddings
 from .openai import OpenAIEmbeddings
 
 __all__ = [
+    "BedrockEmbeddings",
     "OllamaEmbeddings",
     "OpenAIEmbeddings",
 ]

aisploit/embedding/bedrock.py

@@ -0,0 +1,19 @@
+from langchain_community.embeddings import (
+    BedrockEmbeddings as LangchainBedrockEmbeddings,
+)
+
+
+from ..core import BaseEmbeddings
+
+
+class BedrockEmbeddings(LangchainBedrockEmbeddings, BaseEmbeddings):
+    def __init__(
+        self,
+        *,
+        model_id: str = "amazon.titan-embed-text-v1",
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            model_id=model_id,
+            **kwargs,
+        )

aisploit/embedding/ollama.py

@@ -13,4 +13,5 @@ def __init__(
     ) -> None:
         super().__init__(
             model=model,
+            **kwargs,
         )

aisploit/model/__init__.py

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

aisploit/model/bedrock_chat.py

@@ -0,0 +1,19 @@
+from langchain_community.chat_models import BedrockChat as LangchainBedrockChat
+
+from ..core import BaseChatModel
+
+
+class BedrockChat(LangchainBedrockChat, BaseChatModel):
+    def __init__(
+        self,
+        *,
+        model_id: str,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            model_id=model_id,
+            **kwargs,
+        )
+
+    def supports_functions(self) -> bool:
+        return False