Skip to content

Commit 96d13ae

Browse files
authored
Fix truss predict for content-type as text. (#572)
1 parent bb97a05 commit 96d13ae

File tree

1 file changed

+7
-2
lines changed

1 file changed

+7
-2
lines changed

truss/remote/baseten/service.py

+7-2
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ def __init__(
1717
service_url: str,
1818
truss_handle: Optional[TrussHandle] = None,
1919
):
20-
2120
super().__init__(is_draft=is_draft, service_url=service_url)
2221
self._model_id = model_id
2322
self._model_version_id = model_version_id
@@ -43,7 +42,13 @@ def decode_content():
4342
for chunk in response.iter_content(
4443
chunk_size=8192, decode_unicode=True
4544
):
46-
yield chunk.decode(response.encoding or DEFAULT_STREAM_ENCODING)
45+
# Depending on the content-type of the response,
46+
# iter_content will either emit a byte stream, or a stream
47+
# of strings. Only decode in the bytes case.
48+
if isinstance(chunk, bytes):
49+
yield chunk.decode(response.encoding or DEFAULT_STREAM_ENCODING)
50+
else:
51+
yield chunk
4752

4853
return decode_content()
4954

0 commit comments

Comments
 (0)