Skip to content

Commit da79133

Browse files
fix: use streaming for Anthropic API to handle long responses
- Switch from messages.create() to messages.stream() for both sync and async - Required for max_tokens > 21333 or long-running requests - Add warnings for empty responses and truncation (stop_reason == max_tokens)
1 parent 455333a commit da79133

1 file changed

Lines changed: 35 additions & 4 deletions

File tree

‎rlm/clients/anthropic.py‎

Lines changed: 35 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import logging
12
from collections import defaultdict
23
from typing import Any
34

@@ -6,6 +7,8 @@
67
from rlm.clients.base_lm import BaseLM
78
from rlm.core.types import ModelUsageSummary, UsageSummary
89

10+
logger = logging.getLogger(__name__)
11+
912

1013
class AnthropicClient(BaseLM):
1114
"""
@@ -42,9 +45,23 @@ def completion(self, prompt: str | list[dict[str, Any]], model: str | None = Non
4245
if system:
4346
kwargs["system"] = system
4447

45-
response = self.client.messages.create(**kwargs)
48+
# Use streaming to handle long responses (required for max_tokens > 21333)
49+
parts: list[str] = []
50+
with self.client.messages.stream(**kwargs) as stream:
51+
for text in stream.text_stream:
52+
parts.append(text)
53+
response = stream.get_final_message()
54+
55+
full_response = "".join(parts)
56+
57+
if not full_response:
58+
logger.warning("Received empty response from Anthropic API")
59+
60+
if response.stop_reason == "max_tokens":
61+
logger.warning(f"Response truncated at {self.max_tokens} tokens")
62+
4663
self._track_cost(response, model)
47-
return response.content[0].text
64+
return full_response
4865

4966
async def acompletion(
5067
self, prompt: str | list[dict[str, Any]], model: str | None = None
@@ -59,9 +76,23 @@ async def acompletion(
5976
if system:
6077
kwargs["system"] = system
6178

62-
response = await self.async_client.messages.create(**kwargs)
79+
# Use streaming to handle long responses (required for max_tokens > 21333)
80+
parts: list[str] = []
81+
async with self.async_client.messages.stream(**kwargs) as stream:
82+
async for text in stream.text_stream:
83+
parts.append(text)
84+
response = await stream.get_final_message()
85+
86+
full_response = "".join(parts)
87+
88+
if not full_response:
89+
logger.warning("Received empty response from Anthropic API")
90+
91+
if response.stop_reason == "max_tokens":
92+
logger.warning(f"Response truncated at {self.max_tokens} tokens")
93+
6394
self._track_cost(response, model)
64-
return response.content[0].text
95+
return full_response
6596

6697
def _prepare_messages(
6798
self, prompt: str | list[dict[str, Any]]

0 commit comments

Comments
 (0)