Move models param to the endpoint (#516)

Author: billytrend-cohere

src/cohere/aws_client.py

@@ -25,9 +25,6 @@ def __init__(
             aws_session_token: typing.Optional[str] = None,
             aws_region: typing.Optional[str] = None,
             timeout: typing.Optional[float] = None,
-            chat_model: typing.Optional[str] = None,
-            embed_model: typing.Optional[str] = None,
-            generate_model: typing.Optional[str] = None,
             service: typing.Union[typing.Literal["bedrock"], typing.Literal["sagemaker"]],
     ):
         Client.__init__(
@@ -44,9 +41,6 @@ def __init__(
                     aws_secret_key=aws_secret_key,
                     aws_session_token=aws_session_token,
                     aws_region=aws_region,
-                    chat_model=chat_model,
-                    embed_model=embed_model,
-                    generate_model=generate_model,
                 ),
                 timeout=timeout,
             ),
@@ -62,9 +56,6 @@ def get_event_hooks(
         aws_secret_key: typing.Optional[str] = None,
         aws_session_token: typing.Optional[str] = None,
         aws_region: typing.Optional[str] = None,
-        chat_model: typing.Optional[str] = None,
-        embed_model: typing.Optional[str] = None,
-        generate_model: typing.Optional[str] = None,
 ) -> typing.Dict[str, typing.List[EventHook]]:
     return {
         "request": [
@@ -74,17 +65,10 @@ def get_event_hooks(
                 aws_secret_key=aws_secret_key,
                 aws_session_token=aws_session_token,
                 aws_region=aws_region,
-                chat_model=chat_model,
-                embed_model=embed_model,
-                generate_model=generate_model,
             ),
         ],
         "response": [
-            map_response_from_bedrock(
-                chat_model=chat_model,
-                embed_model=embed_model,
-                generate_model=generate_model,
-            )
+            map_response_from_bedrock()
         ],
     }
 
@@ -138,17 +122,12 @@ def stream_generator(response: httpx.Response, endpoint: str) -> typing.Iterator
                     yield (json.dumps(parsed.dict()) + "\n").encode("utf-8") # type: ignore
 
 
-def map_response_from_bedrock(
-        chat_model: typing.Optional[str] = None,
-        embed_model: typing.Optional[str] = None,
-        generate_model: typing.Optional[str] = None,
-):
+def map_response_from_bedrock():
     def _hook(
             response: httpx.Response,
     ) -> None:
         stream = response.headers["content-type"] == "application/vnd.amazon.eventstream"
-        endpoint = get_endpoint_from_url(
-            response.url.path, chat_model, embed_model, generate_model)
+        endpoint = response.request.extensions["endpoint"]
         output: typing.Iterator[bytes]
 
         if stream:
@@ -179,9 +158,6 @@ def map_request_to_bedrock(
         aws_secret_key: typing.Optional[str] = None,
         aws_session_token: typing.Optional[str] = None,
         aws_region: typing.Optional[str] = None,
-        chat_model: typing.Optional[str] = None,
-        embed_model: typing.Optional[str] = None,
-        generate_model: typing.Optional[str] = None,
 ) -> EventHook:
     session = boto3.Session(
         region_name=aws_region,
@@ -192,23 +168,18 @@ def map_request_to_bedrock(
     credentials = session.get_credentials()
     signer = SigV4Auth(credentials, service, session.region_name)
 
-    model_lookup = {
-        "embed": embed_model,
-        "chat": chat_model,
-        "generate": generate_model,
-    }
-
     def _event_hook(request: httpx.Request) -> None:
         headers = request.headers.copy()
         del headers["connection"]
 
         endpoint = request.url.path.split("/")[-1]
         body = json.loads(request.read())
+        model = body["model"]
 
         url = get_url(
             platform=service,
             aws_region=aws_region,
-            model=model_lookup[endpoint],  # type: ignore
+            model=model,  # type: ignore
             stream="stream" in body and body["stream"],
         )
         request.url = URL(url)
@@ -217,6 +188,9 @@ def _event_hook(request: httpx.Request) -> None:
         if "stream" in body:
             del body["stream"]
 
+        if "model" in body:
+            del body["model"]
+
         new_body = json.dumps(body).encode("utf-8")
         request.stream = ByteStream(new_body)
         request._content = new_body
@@ -231,6 +205,7 @@ def _event_hook(request: httpx.Request) -> None:
         signer.add_auth(aws_request)
 
         request.headers = httpx.Headers(aws_request.prepare().headers)
+        request.extensions["endpoint"] = endpoint
 
     return _event_hook
 

src/cohere/bedrock_client.py

@@ -17,9 +17,6 @@ def __init__(
             aws_session_token: typing.Optional[str] = None,
             aws_region: typing.Optional[str] = None,
             timeout: typing.Optional[float] = None,
-            chat_model: typing.Optional[str] = None,
-            embed_model: typing.Optional[str] = None,
-            generate_model: typing.Optional[str] = None,
     ):
         AwsClient.__init__(
             self,
@@ -29,7 +26,4 @@ def __init__(
             aws_session_token=aws_session_token,
             aws_region=aws_region,
             timeout=timeout,
-            chat_model=chat_model,
-            embed_model=embed_model,
-            generate_model=generate_model,
         )

src/cohere/sagemaker_client.py

@@ -14,9 +14,6 @@ def __init__(
             aws_session_token: typing.Optional[str] = None,
             aws_region: typing.Optional[str] = None,
             timeout: typing.Optional[float] = None,
-            chat_model: typing.Optional[str] = None,
-            embed_model: typing.Optional[str] = None,
-            generate_model: typing.Optional[str] = None,
     ):
         AwsClient.__init__(
             self,
@@ -26,7 +23,4 @@ def __init__(
             aws_session_token=aws_session_token,
             aws_region=aws_region,
             timeout=timeout,
-            chat_model=chat_model,
-            embed_model=embed_model,
-            generate_model=generate_model,
         )

tests/test_aws_client.py

@@ -1,58 +1,73 @@
 import os
 import unittest
 
+import typing
 import cohere
 from parameterized import parameterized_class  # type: ignore
 
 package_dir = os.path.dirname(os.path.abspath(__file__))
 embed_job = os.path.join(package_dir, 'embed_job.jsonl')
 
 
+models = {
+    "bedrock": {
+        "chat_model": "cohere.command-r-plus-v1:0",
+        "embed_model": "cohere.embed-multilingual-v3",
+        "generate_model": "cohere.command-text-v14",
+    },
+    "sagemaker": {
+        "chat_model": "cohere.command-r-plus-v1:0",
+        "embed_model": "cohere.embed-multilingual-v3",
+        "generate_model": "cohere-command-light",
+    },
+}
+
+
 @parameterized_class([
     {
         "client": cohere.BedrockClient(
             timeout=10000,
             aws_region="us-east-1",
-            chat_model="cohere.command-r-plus-v1:0",
-            embed_model="cohere.embed-multilingual-v3",
-            generate_model="cohere.command-text-v14",
             aws_access_key="...",
             aws_secret_key="...",
             aws_session_token="...",
-        )
+        ),
+        "models": models["bedrock"],
     },
     {
         "client": cohere.SagemakerClient(
             timeout=10000,
             aws_region="us-east-1",
-            chat_model="cohere.command-r-plus-v1:0",
-            embed_model="cohere.embed-multilingual-v3",
-            generate_model="cohere-command-light",
             aws_access_key="...",
             aws_secret_key="...",
             aws_session_token="...",
-        )
+        ),
+        "models": models["sagemaker"],
     }
 ])
 @unittest.skip("skip tests until they work in CI")
 class TestClient(unittest.TestCase):
-    client: cohere.AwsClient;
+    client: cohere.AwsClient
+    models: typing.Dict[str, str]
 
     def test_embed(self) -> None:
         response = self.client.embed(
+            model=self.models["embed_model"],
             texts=["I love Cohere!"],
             input_type="search_document",
         )
         print(response)
 
     def test_generate(self) -> None:
         response = self.client.generate(
+            model=self.models["generate_model"],
             prompt='Please explain to me how LLMs work',
         )
         print(response)
 
     def test_generate_stream(self) -> None:
         response = self.client.generate_stream(
+            model=self.models["generate_model"],
             prompt='Please explain to me how LLMs work',
         )
         for event in response:
@@ -62,6 +77,7 @@ def test_generate_stream(self) -> None:
 
     def test_chat(self) -> None:
         response = self.client.chat(
+            model=self.models["chat_model"],
             message='Please explain to me how LLMs work',
         )
         print(response)
@@ -73,6 +89,7 @@ def test_chat(self) -> None:
     def test_chat_stream(self) -> None:
         response_types = set()
         response = self.client.chat_stream(
+            model=self.models["chat_model"],
             message='Please explain to me how LLMs work',
         )
         for event in response: