Add v2 clients

cohere-ai · Nov 29, 2024 · 245949d · 245949d
1 parent 061f005
commit 245949d
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 3 deletions.
diff --git a/src/cohere/aws_client.py b/src/cohere/aws_client.py
@@ -12,7 +12,7 @@
 from .client import Client, ClientEnvironment
 from .core import construct_type
 from .manually_maintained.lazy_aws_deps import lazy_boto3, lazy_botocore
-
+from .client_v2 import ClientV2
 
 class AwsClient(Client):
     def __init__(
@@ -45,6 +45,37 @@ def __init__(
         )
 
 
+class AwsClientV2(ClientV2):
+    def __init__(
+            self,
+            *,
+            aws_access_key: typing.Optional[str] = None,
+            aws_secret_key: typing.Optional[str] = None,
+            aws_session_token: typing.Optional[str] = None,
+            aws_region: typing.Optional[str] = None,
+            timeout: typing.Optional[float] = None,
+            service: typing.Union[typing.Literal["bedrock"], typing.Literal["sagemaker"]],
+    ):
+        Client.__init__(
+            self,
+            base_url="https://api.cohere.com",  # this url is unused for BedrockClient
+            environment=ClientEnvironment.PRODUCTION,
+            client_name="n/a",
+            timeout=timeout,
+            api_key="n/a",
+            httpx_client=httpx.Client(
+                event_hooks=get_event_hooks(
+                    service=service,
+                    aws_access_key=aws_access_key,
+                    aws_secret_key=aws_secret_key,
+                    aws_session_token=aws_session_token,
+                    aws_region=aws_region,
+                ),
+                timeout=timeout,
+            ),
+        )
+
+
 EventHook = typing.Callable[..., typing.Any]
 
 

diff --git a/src/cohere/bedrock_client.py b/src/cohere/bedrock_client.py
@@ -2,7 +2,7 @@
 
 from tokenizers import Tokenizer  # type: ignore
 
-from .aws_client import AwsClient
+from .aws_client import AwsClient, AwsClientV2
 
 
 class BedrockClient(AwsClient):
@@ -24,3 +24,24 @@ def __init__(
             aws_region=aws_region,
             timeout=timeout,
         )
+
+
+class BedrockClientV2(AwsClientV2):
+    def __init__(
+            self,
+            *,
+            aws_access_key: typing.Optional[str] = None,
+            aws_secret_key: typing.Optional[str] = None,
+            aws_session_token: typing.Optional[str] = None,
+            aws_region: typing.Optional[str] = None,
+            timeout: typing.Optional[float] = None,
+    ):
+        AwsClientV2.__init__(
+            self,
+            service="bedrock",
+            aws_access_key=aws_access_key,
+            aws_secret_key=aws_secret_key,
+            aws_session_token=aws_session_token,
+            aws_region=aws_region,
+            timeout=timeout,
+        )
diff --git a/src/cohere/sagemaker_client.py b/src/cohere/sagemaker_client.py
@@ -1,6 +1,6 @@
 import typing
 
-from .aws_client import AwsClient
+from .aws_client import AwsClient, AwsClientV2
 from .manually_maintained.cohere_aws.client import Client
 from .manually_maintained.cohere_aws.mode import Mode
 
@@ -26,4 +26,28 @@ def __init__(
             aws_region=aws_region,
             timeout=timeout,
         )
+        self.sagemaker_finetuning = Client(aws_region=aws_region)
+
+
+class SagemakerClientV2(AwsClientV2):
+    sagemaker_finetuning: Client
+
+    def __init__(
+            self,
+            *,
+            aws_access_key: typing.Optional[str] = None,
+            aws_secret_key: typing.Optional[str] = None,
+            aws_session_token: typing.Optional[str] = None,
+            aws_region: typing.Optional[str] = None,
+            timeout: typing.Optional[float] = None,
+    ):
+        AwsClientV2.__init__(
+            self,
+            service="sagemaker",
+            aws_access_key=aws_access_key,
+            aws_secret_key=aws_secret_key,
+            aws_session_token=aws_session_token,
+            aws_region=aws_region,
+            timeout=timeout,
+        )
         self.sagemaker_finetuning = Client(aws_region=aws_region)