revert: restore AWS client files to main — no AWS changes in this PR

Author: fede-kamel

src/cohere/aws_client.py

@@ -4,15 +4,14 @@
 import typing
 
 import httpx
-from httpx import URL, ByteStream
+from httpx import URL, SyncByteStream, ByteStream
 
 from . import GenerateStreamedResponse, Generation, \
     NonStreamedChatResponse, EmbedResponse, StreamedChatResponse, RerankResponse, ApiMeta, ApiMetaTokens, \
     ApiMetaBilledUnits
 from .client import Client, ClientEnvironment
 from .core import construct_type
 from .manually_maintained.lazy_aws_deps import lazy_boto3, lazy_botocore
-from .manually_maintained.streaming import Streamer
 from .client_v2 import ClientV2
 
 class AwsClient(Client):
@@ -113,6 +112,16 @@ def get_event_hooks(
                               })
 
 
+class Streamer(SyncByteStream):
+    lines: typing.Iterator[bytes]
+
+    def __init__(self, lines: typing.Iterator[bytes]):
+        self.lines = lines
+
+    def __iter__(self) -> typing.Iterator[bytes]:
+        return self.lines
+
+
 response_mapping: typing.Dict[str, typing.Any] = {
     "chat": NonStreamedChatResponse,
     "embed": EmbedResponse,
@@ -230,7 +239,6 @@ def _event_hook(request: httpx.Request) -> None:
         )
         request.url = URL(url)
         request.headers["host"] = request.url.host
-        headers["host"] = request.url.host
 
         if endpoint == "rerank":
             body["api_version"] = get_api_version(version=api_version)
@@ -282,4 +290,4 @@ def get_api_version(*, version: str):
         "v2": 2,
     }
 
-    return int_version.get(version, 1)
+    return int_version.get(version, 1)

src/cohere/manually_maintained/cohere_aws/client.py

@@ -20,29 +20,19 @@ class Client:
     def __init__(
            self,
             aws_region: typing.Optional[str] = None,
-            mode: Mode = Mode.SAGEMAKER,
         ):
         """
         By default we assume region configured in AWS CLI (`aws configure get region`). You can change the region with
         `aws configure set region us-west-2` or override it with `region_name` parameter.
         """
-        self.mode = mode
+        self._client = lazy_boto3().client("sagemaker-runtime", region_name=aws_region)
+        self._service_client = lazy_boto3().client("sagemaker", region_name=aws_region)
         if os.environ.get('AWS_DEFAULT_REGION') is None:
             os.environ['AWS_DEFAULT_REGION'] = aws_region
+        self._sess = lazy_sagemaker().Session(sagemaker_client=self._service_client)
+        self.mode = Mode.SAGEMAKER
 
-        if self.mode == Mode.SAGEMAKER:
-            self._client = lazy_boto3().client("sagemaker-runtime", region_name=aws_region)
-            self._service_client = lazy_boto3().client("sagemaker", region_name=aws_region)
-            self._sess = lazy_sagemaker().Session(sagemaker_client=self._service_client)
-        elif self.mode == Mode.BEDROCK:
-            self._client = lazy_boto3().client("bedrock-runtime", region_name=aws_region)
-            self._service_client = lazy_boto3().client("bedrock", region_name=aws_region)
-            self._sess = None
-            self._endpoint_name = None
 
-    def _require_sagemaker(self) -> None:
-        if self.mode != Mode.SAGEMAKER:
-            raise CohereError("This method is only supported in SageMaker mode.")
 
     def _does_endpoint_exist(self, endpoint_name: str) -> bool:
         try:
@@ -60,7 +50,6 @@ def connect_to_endpoint(self, endpoint_name: str) -> None:
         Raises:
             CohereError: Connection to the endpoint failed.
         """
-        self._require_sagemaker()
         if not self._does_endpoint_exist(endpoint_name):
             raise CohereError(f"Endpoint {endpoint_name} does not exist.")
         self._endpoint_name = endpoint_name
@@ -148,7 +137,6 @@ def create_endpoint(
                 will be used to get the role. This should work when one uses the client inside SageMaker. If this errors
                 out, the default role "ServiceRoleSagemaker" will be used, which generally works outside of SageMaker.
         """
-        self._require_sagemaker()
         # First, check if endpoint already exists
         if self._does_endpoint_exist(endpoint_name):
             if recreate:
@@ -562,15 +550,11 @@ def embed(
         variant: Optional[str] = None,
         input_type: Optional[str] = None,
         model_id: Optional[str] = None,
-        output_dimension: Optional[int] = None,
-        embedding_types: Optional[List[str]] = None,
-    ) -> Union[Embeddings, Dict[str, List]]:
+    ) -> Embeddings:
         json_params = {
             'texts': texts,
             'truncate': truncate,
-            "input_type": input_type,
-            "output_dimension": output_dimension,
-            "embedding_types": embedding_types,
+            "input_type": input_type
         }
         for key, value in list(json_params.items()):
             if value is None:
@@ -607,10 +591,7 @@ def _sagemaker_embed(self, json_params: Dict[str, Any], variant: str):
             # ValidationError, e.g. when variant is bad
             raise CohereError(str(e))
 
-        embeddings = response['embeddings']
-        if isinstance(embeddings, dict):
-            return embeddings
-        return Embeddings(embeddings)
+        return Embeddings(response['embeddings'])
 
     def _bedrock_embed(self, json_params: Dict[str, Any], model_id: str):
         if not model_id:
@@ -631,10 +612,7 @@ def _bedrock_embed(self, json_params: Dict[str, Any], model_id: str):
             # ValidationError, e.g. when variant is bad
             raise CohereError(str(e))
 
-        embeddings = response['embeddings']
-        if isinstance(embeddings, dict):
-            return embeddings
-        return Embeddings(embeddings)
+        return Embeddings(response['embeddings'])
 
 
     def rerank(self,
@@ -827,7 +805,6 @@ def export_finetune(
             This should work when one uses the client inside SageMaker. If this errors out,
             the default role "ServiceRoleSagemaker" will be used, which generally works outside SageMaker.
         """
-        self._require_sagemaker()
         if name == "model":
             raise ValueError("name cannot be 'model'")
 
@@ -971,7 +948,6 @@ def summarize(
         additional_command: Optional[str] = "",
         variant: Optional[str] = None
     ) -> Summary:
-        self._require_sagemaker()
 
         if self._endpoint_name is None:
             raise CohereError("No endpoint connected. "
@@ -1013,7 +989,6 @@ def summarize(
 
 
     def delete_endpoint(self) -> None:
-        self._require_sagemaker()
         if self._endpoint_name is None:
             raise CohereError("No endpoint connected.")
         try: