Skip to content

Commit aa0a2fc

Browse files
fix the streaming test (#8367)
1 parent b5390e9 commit aa0a2fc

File tree

1 file changed

+6
-6
lines changed

1 file changed

+6
-6
lines changed

tests/streaming/test_streaming.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -231,7 +231,6 @@ def __call__(self, x: str, **kwargs):
231231

232232
@pytest.mark.anyio
233233
async def test_streaming_handles_space_correctly():
234-
dspy.settings.configure(lm=dspy.LM("openai/gpt-4o-mini", cache=False), adapter=dspy.ChatAdapter())
235234
my_program = dspy.Predict("question->answer")
236235
program = dspy.streamify(
237236
my_program, stream_listeners=[dspy.streaming.StreamListener(signature_field_name="answer")]
@@ -250,11 +249,12 @@ async def gpt_4o_mini_stream(*args, **kwargs):
250249
)
251250

252251
with mock.patch("litellm.acompletion", side_effect=gpt_4o_mini_stream):
253-
output = program(question="What is the capital of France?")
254-
all_chunks = []
255-
async for value in output:
256-
if isinstance(value, dspy.streaming.StreamResponse):
257-
all_chunks.append(value)
252+
with dspy.context(lm=dspy.LM("openai/gpt-4o-mini", cache=False), adapter=dspy.ChatAdapter()):
253+
output = program(question="What is the capital of France?")
254+
all_chunks = []
255+
async for value in output:
256+
if isinstance(value, dspy.streaming.StreamResponse):
257+
all_chunks.append(value)
258258

259259
assert all_chunks[0].chunk == "How are you doing?"
260260

0 commit comments

Comments
 (0)