fix(agents-api): Patch unexpected response from litellm (finish_reaso…

…n) (#894)

<!-- ELLIPSIS_HIDDEN -->


> [!IMPORTANT]
> Add function to handle 'eos' finish reasons in litellm responses and
reformat print statements in `codec.py`.
> 
>   - **Behavior**:
> - Add `patch_litellm_response()` in `litellm.py` to convert 'eos'
finish reasons to 'stop'.
> - Apply `patch_litellm_response()` to `model_response` in
`acompletion()`.
>   - **Misc**:
> - Reformat print statements in `serialize()` and `deserialize()` in
`codec.py` for readability.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=julep-ai%2Fjulep&utm_source=github&utm_medium=referral)<sup>
for 653b8f3. It will automatically
update as commits are pushed.</sup>


<!-- ELLIPSIS_HIDDEN -->

agents-api/agents_api/clients/litellm.py

@@ -27,6 +27,25 @@
 litellm.drop_params = True
 
 
+def patch_litellm_response(
+    model_response: ModelResponse | CustomStreamWrapper,
+) -> ModelResponse | CustomStreamWrapper:
+    """
+    Patches the response we get from litellm to handle unexpected response formats.
+    """
+
+    if isinstance(model_response, ModelResponse):
+        for choice in model_response.choices:
+            if choice.finish_reason == "eos":
+                choice.finish_reason = "stop"
+
+    elif isinstance(model_response, CustomStreamWrapper):
+        if model_response.received_finish_reason == "eos":
+            model_response.received_finish_reason = "stop"
+
+    return model_response
+
+
 @wraps(_acompletion)
 @beartype
 async def acompletion(
@@ -45,14 +64,18 @@ async def acompletion(
             for message in messages
         ]
 
-    return await _acompletion(
+    model_response = await _acompletion(
         model=model,
         messages=messages,
         **settings,
         base_url=None if custom_api_key else litellm_url,
         api_key=custom_api_key or litellm_master_key,
     )
 
+    model_response = patch_litellm_response(model_response)
+
+    return model_response
+
 
 @wraps(_aembedding)
 @beartype

agents-api/agents_api/worker/codec.py

@@ -31,7 +31,9 @@ def serialize(x: Any) -> bytes:
 
     duration = time.time() - start_time
     if duration > 1:
-        print(f"||| [SERIALIZE] Time taken: {duration}s // Object size: {sys.getsizeof(x) / 1000}kb")
+        print(
+            f"||| [SERIALIZE] Time taken: {duration}s // Object size: {sys.getsizeof(x) / 1000}kb"
+        )
 
     return compressed
 
@@ -43,7 +45,9 @@ def deserialize(b: bytes) -> Any:
 
     duration = time.time() - start_time
     if duration > 1:
-        print(f"||| [DESERIALIZE] Time taken: {duration}s // Object size: {sys.getsizeof(b) / 1000}kb")
+        print(
+            f"||| [DESERIALIZE] Time taken: {duration}s // Object size: {sys.getsizeof(b) / 1000}kb"
+        )
 
     return object