test_utils.py
1 # SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 2 # 3 # SPDX-License-Identifier: Apache-2.0 4 5 from unittest.mock import call, patch 6 7 from openai.types.chat import chat_completion_chunk 8 9 from haystack.components.generators.utils import _convert_streaming_chunks_to_chat_message, print_streaming_chunk 10 from haystack.dataclasses import ( 11 ComponentInfo, 12 ReasoningContent, 13 StreamingChunk, 14 ToolCall, 15 ToolCallDelta, 16 ToolCallResult, 17 ) 18 19 20 def test_convert_streaming_chunks_to_chat_message_tool_calls_in_any_chunk(): 21 chunks = [ 22 StreamingChunk( 23 content="", 24 meta={ 25 "model": "gpt-4o-mini-2024-07-18", 26 "index": 0, 27 "tool_calls": None, 28 "finish_reason": None, 29 "received_at": "2025-02-19T16:02:55.910076", 30 }, 31 component_info=ComponentInfo(name="test", type="test"), 32 ), 33 StreamingChunk( 34 content="", 35 meta={ 36 "model": "gpt-4o-mini-2024-07-18", 37 "index": 0, 38 "tool_calls": [ 39 chat_completion_chunk.ChoiceDeltaToolCall( 40 index=0, 41 id="call_ZOj5l67zhZOx6jqjg7ATQwb6", 42 function=chat_completion_chunk.ChoiceDeltaToolCallFunction( 43 arguments="", name="rag_pipeline_tool" 44 ), 45 type="function", 46 ) 47 ], 48 "finish_reason": None, 49 "received_at": "2025-02-19T16:02:55.913919", 50 }, 51 component_info=ComponentInfo(name="test", type="test"), 52 index=0, 53 start=True, 54 tool_calls=[ 55 ToolCallDelta(id="call_ZOj5l67zhZOx6jqjg7ATQwb6", tool_name="rag_pipeline_tool", arguments="", index=0) 56 ], 57 ), 58 StreamingChunk( 59 content="", 60 meta={ 61 "model": "gpt-4o-mini-2024-07-18", 62 "index": 0, 63 "tool_calls": [ 64 chat_completion_chunk.ChoiceDeltaToolCall( 65 index=0, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments='{"qu') 66 ) 67 ], 68 "finish_reason": None, 69 "received_at": "2025-02-19T16:02:55.914439", 70 }, 71 component_info=ComponentInfo(name="test", type="test"), 72 index=0, 73 tool_calls=[ToolCallDelta(arguments='{"qu', index=0)], 74 ), 75 StreamingChunk( 76 content="", 77 meta={ 78 "model": "gpt-4o-mini-2024-07-18", 79 "index": 0, 80 "tool_calls": [ 81 chat_completion_chunk.ChoiceDeltaToolCall( 82 index=0, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments='ery":') 83 ) 84 ], 85 "finish_reason": None, 86 "received_at": "2025-02-19T16:02:55.924146", 87 }, 88 component_info=ComponentInfo(name="test", type="test"), 89 index=0, 90 tool_calls=[ToolCallDelta(arguments='ery":', index=0)], 91 ), 92 StreamingChunk( 93 content="", 94 meta={ 95 "model": "gpt-4o-mini-2024-07-18", 96 "index": 0, 97 "tool_calls": [ 98 chat_completion_chunk.ChoiceDeltaToolCall( 99 index=0, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments=' "Wher') 100 ) 101 ], 102 "finish_reason": None, 103 "received_at": "2025-02-19T16:02:55.924420", 104 }, 105 component_info=ComponentInfo(name="test", type="test"), 106 index=0, 107 tool_calls=[ToolCallDelta(arguments=' "Wher', index=0)], 108 ), 109 StreamingChunk( 110 content="", 111 meta={ 112 "model": "gpt-4o-mini-2024-07-18", 113 "index": 0, 114 "tool_calls": [ 115 chat_completion_chunk.ChoiceDeltaToolCall( 116 index=0, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments="e do") 117 ) 118 ], 119 "finish_reason": None, 120 "received_at": "2025-02-19T16:02:55.944398", 121 }, 122 component_info=ComponentInfo(name="test", type="test"), 123 index=0, 124 tool_calls=[ToolCallDelta(arguments="e do", index=0)], 125 ), 126 StreamingChunk( 127 content="", 128 meta={ 129 "model": "gpt-4o-mini-2024-07-18", 130 "index": 0, 131 "tool_calls": [ 132 chat_completion_chunk.ChoiceDeltaToolCall( 133 index=0, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments="es Ma") 134 ) 135 ], 136 "finish_reason": None, 137 "received_at": "2025-02-19T16:02:55.944958", 138 }, 139 component_info=ComponentInfo(name="test", type="test"), 140 index=0, 141 tool_calls=[ToolCallDelta(arguments="es Ma", index=0)], 142 ), 143 StreamingChunk( 144 content="", 145 meta={ 146 "model": "gpt-4o-mini-2024-07-18", 147 "index": 0, 148 "tool_calls": [ 149 chat_completion_chunk.ChoiceDeltaToolCall( 150 index=0, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments="rk liv") 151 ) 152 ], 153 "finish_reason": None, 154 "received_at": "2025-02-19T16:02:55.945507", 155 }, 156 component_info=ComponentInfo(name="test", type="test"), 157 index=0, 158 tool_calls=[ToolCallDelta(arguments="rk liv", index=0)], 159 ), 160 StreamingChunk( 161 content="", 162 meta={ 163 "model": "gpt-4o-mini-2024-07-18", 164 "index": 0, 165 "tool_calls": [ 166 chat_completion_chunk.ChoiceDeltaToolCall( 167 index=0, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments='e?"}') 168 ) 169 ], 170 "finish_reason": None, 171 "received_at": "2025-02-19T16:02:55.946018", 172 }, 173 component_info=ComponentInfo(name="test", type="test"), 174 index=0, 175 tool_calls=[ToolCallDelta(arguments='e?"}', index=0)], 176 ), 177 StreamingChunk( 178 content="", 179 meta={ 180 "model": "gpt-4o-mini-2024-07-18", 181 "index": 0, 182 "tool_calls": [ 183 chat_completion_chunk.ChoiceDeltaToolCall( 184 index=1, 185 id="call_STxsYY69wVOvxWqopAt3uWTB", 186 function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments="", name="get_weather"), 187 type="function", 188 ) 189 ], 190 "finish_reason": None, 191 "received_at": "2025-02-19T16:02:55.946578", 192 }, 193 component_info=ComponentInfo(name="test", type="test"), 194 index=1, 195 start=True, 196 tool_calls=[ 197 ToolCallDelta(id="call_STxsYY69wVOvxWqopAt3uWTB", tool_name="get_weather", arguments="", index=1) 198 ], 199 ), 200 StreamingChunk( 201 content="", 202 meta={ 203 "model": "gpt-4o-mini-2024-07-18", 204 "index": 0, 205 "tool_calls": [ 206 chat_completion_chunk.ChoiceDeltaToolCall( 207 index=1, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments='{"ci') 208 ) 209 ], 210 "finish_reason": None, 211 "received_at": "2025-02-19T16:02:55.946981", 212 }, 213 component_info=ComponentInfo(name="test", type="test"), 214 index=1, 215 tool_calls=[ToolCallDelta(arguments='{"ci', index=1)], 216 ), 217 StreamingChunk( 218 content="", 219 meta={ 220 "model": "gpt-4o-mini-2024-07-18", 221 "index": 0, 222 "tool_calls": [ 223 chat_completion_chunk.ChoiceDeltaToolCall( 224 index=1, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments='ty": ') 225 ) 226 ], 227 "finish_reason": None, 228 "received_at": "2025-02-19T16:02:55.947411", 229 }, 230 component_info=ComponentInfo(name="test", type="test"), 231 index=1, 232 tool_calls=[ToolCallDelta(arguments='ty": ', index=1)], 233 ), 234 StreamingChunk( 235 content="", 236 meta={ 237 "model": "gpt-4o-mini-2024-07-18", 238 "index": 0, 239 "tool_calls": [ 240 chat_completion_chunk.ChoiceDeltaToolCall( 241 index=1, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments='"Berli') 242 ) 243 ], 244 "finish_reason": None, 245 "received_at": "2025-02-19T16:02:55.947643", 246 }, 247 component_info=ComponentInfo(name="test", type="test"), 248 index=1, 249 tool_calls=[ToolCallDelta(arguments='"Berli', index=1)], 250 ), 251 StreamingChunk( 252 content="", 253 meta={ 254 "model": "gpt-4o-mini-2024-07-18", 255 "index": 0, 256 "tool_calls": [ 257 chat_completion_chunk.ChoiceDeltaToolCall( 258 index=1, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments='n"}') 259 ) 260 ], 261 "finish_reason": None, 262 "received_at": "2025-02-19T16:02:55.947939", 263 }, 264 component_info=ComponentInfo(name="test", type="test"), 265 index=1, 266 tool_calls=[ToolCallDelta(arguments='n"}', index=1)], 267 ), 268 StreamingChunk( 269 content="", 270 meta={ 271 "model": "gpt-4o-mini-2024-07-18", 272 "index": 0, 273 "tool_calls": None, 274 "finish_reason": "tool_calls", 275 "received_at": "2025-02-19T16:02:55.948772", 276 }, 277 component_info=ComponentInfo(name="test", type="test"), 278 finish_reason="tool_calls", 279 ), 280 StreamingChunk( 281 content="", 282 meta={ 283 "model": "gpt-4o-mini-2024-07-18", 284 "index": 0, 285 "tool_calls": None, 286 "finish_reason": None, 287 "received_at": "2025-02-19T16:02:55.948772", 288 "usage": { 289 "completion_tokens": 42, 290 "prompt_tokens": 282, 291 "total_tokens": 324, 292 "completion_tokens_details": { 293 "accepted_prediction_tokens": 0, 294 "audio_tokens": 0, 295 "reasoning_tokens": 0, 296 "rejected_prediction_tokens": 0, 297 }, 298 "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, 299 }, 300 }, 301 component_info=ComponentInfo(name="test", type="test"), 302 ), 303 ] 304 305 # Convert chunks to a chat message 306 result = _convert_streaming_chunks_to_chat_message(chunks=chunks) 307 308 assert not result.texts 309 assert not result.text 310 311 # Verify both tool calls were found and processed 312 assert len(result.tool_calls) == 2 313 assert result.tool_calls[0].id == "call_ZOj5l67zhZOx6jqjg7ATQwb6" 314 assert result.tool_calls[0].tool_name == "rag_pipeline_tool" 315 assert result.tool_calls[0].arguments == {"query": "Where does Mark live?"} 316 assert result.tool_calls[1].id == "call_STxsYY69wVOvxWqopAt3uWTB" 317 assert result.tool_calls[1].tool_name == "get_weather" 318 assert result.tool_calls[1].arguments == {"city": "Berlin"} 319 320 # Verify meta information 321 assert result.meta["model"] == "gpt-4o-mini-2024-07-18" 322 assert result.meta["finish_reason"] == "tool_calls" 323 assert result.meta["index"] == 0 324 assert result.meta["completion_start_time"] == "2025-02-19T16:02:55.910076" 325 assert result.meta["usage"] == { 326 "completion_tokens": 42, 327 "prompt_tokens": 282, 328 "total_tokens": 324, 329 "completion_tokens_details": { 330 "accepted_prediction_tokens": 0, 331 "audio_tokens": 0, 332 "reasoning_tokens": 0, 333 "rejected_prediction_tokens": 0, 334 }, 335 "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, 336 } 337 338 339 def test_convert_streaming_chunk_to_chat_message_two_tool_calls_in_same_chunk(): 340 chunks = [ 341 StreamingChunk( 342 content="", 343 meta={ 344 "model": "mistral-small-latest", 345 "index": 0, 346 "tool_calls": None, 347 "finish_reason": None, 348 "usage": None, 349 }, 350 component_info=ComponentInfo( 351 type="haystack_integrations.components.generators.mistral.chat.chat_generator.MistralChatGenerator", 352 name=None, 353 ), 354 ), 355 StreamingChunk( 356 content="", 357 meta={ 358 "model": "mistral-small-latest", 359 "index": 0, 360 "finish_reason": "tool_calls", 361 "usage": { 362 "completion_tokens": 35, 363 "prompt_tokens": 77, 364 "total_tokens": 112, 365 "completion_tokens_details": None, 366 "prompt_tokens_details": None, 367 }, 368 }, 369 component_info=ComponentInfo( 370 type="haystack_integrations.components.generators.mistral.chat.chat_generator.MistralChatGenerator", 371 name=None, 372 ), 373 index=0, 374 tool_calls=[ 375 ToolCallDelta(index=0, tool_name="weather", arguments='{"city": "Paris"}', id="FL1FFlqUG"), 376 ToolCallDelta(index=1, tool_name="weather", arguments='{"city": "Berlin"}', id="xSuhp66iB"), 377 ], 378 start=True, 379 finish_reason="tool_calls", 380 ), 381 ] 382 383 # Convert chunks to a chat message 384 result = _convert_streaming_chunks_to_chat_message(chunks=chunks) 385 386 assert not result.texts 387 assert not result.text 388 389 # Verify both tool calls were found and processed 390 assert len(result.tool_calls) == 2 391 assert result.tool_calls[0].id == "FL1FFlqUG" 392 assert result.tool_calls[0].tool_name == "weather" 393 assert result.tool_calls[0].arguments == {"city": "Paris"} 394 assert result.tool_calls[1].id == "xSuhp66iB" 395 assert result.tool_calls[1].tool_name == "weather" 396 assert result.tool_calls[1].arguments == {"city": "Berlin"} 397 398 399 def test_convert_streaming_chunk_to_chat_message_empty_tool_call_delta(): 400 chunks = [ 401 StreamingChunk( 402 content="", 403 meta={ 404 "model": "gpt-4o-mini-2024-07-18", 405 "index": 0, 406 "tool_calls": None, 407 "finish_reason": None, 408 "received_at": "2025-02-19T16:02:55.910076", 409 }, 410 component_info=ComponentInfo(name="test", type="test"), 411 ), 412 StreamingChunk( 413 content="", 414 meta={ 415 "model": "gpt-4o-mini-2024-07-18", 416 "index": 0, 417 "tool_calls": [ 418 chat_completion_chunk.ChoiceDeltaToolCall( 419 index=0, 420 id="call_ZOj5l67zhZOx6jqjg7ATQwb6", 421 function=chat_completion_chunk.ChoiceDeltaToolCallFunction( 422 arguments='{"query":', name="rag_pipeline_tool" 423 ), 424 type="function", 425 ) 426 ], 427 "finish_reason": None, 428 "received_at": "2025-02-19T16:02:55.913919", 429 }, 430 component_info=ComponentInfo(name="test", type="test"), 431 index=0, 432 start=True, 433 tool_calls=[ 434 ToolCallDelta( 435 id="call_ZOj5l67zhZOx6jqjg7ATQwb6", tool_name="rag_pipeline_tool", arguments='{"query":', index=0 436 ) 437 ], 438 ), 439 StreamingChunk( 440 content="", 441 meta={ 442 "model": "gpt-4o-mini-2024-07-18", 443 "index": 0, 444 "tool_calls": [ 445 chat_completion_chunk.ChoiceDeltaToolCall( 446 index=0, 447 function=chat_completion_chunk.ChoiceDeltaToolCallFunction( 448 arguments=' "Where does Mark live?"}' 449 ), 450 ) 451 ], 452 "finish_reason": None, 453 "received_at": "2025-02-19T16:02:55.924420", 454 }, 455 component_info=ComponentInfo(name="test", type="test"), 456 index=0, 457 tool_calls=[ToolCallDelta(arguments=' "Where does Mark live?"}', index=0)], 458 ), 459 StreamingChunk( 460 content="", 461 meta={ 462 "model": "gpt-4o-mini-2024-07-18", 463 "index": 0, 464 "tool_calls": [ 465 chat_completion_chunk.ChoiceDeltaToolCall( 466 index=0, function=chat_completion_chunk.ChoiceDeltaToolCallFunction() 467 ) 468 ], 469 "finish_reason": "tool_calls", 470 "received_at": "2025-02-19T16:02:55.948772", 471 }, 472 tool_calls=[ToolCallDelta(index=0)], 473 component_info=ComponentInfo(name="test", type="test"), 474 finish_reason="tool_calls", 475 index=0, 476 ), 477 StreamingChunk( 478 content="", 479 meta={ 480 "model": "gpt-4o-mini-2024-07-18", 481 "index": 0, 482 "tool_calls": None, 483 "finish_reason": None, 484 "received_at": "2025-02-19T16:02:55.948772", 485 "usage": { 486 "completion_tokens": 42, 487 "prompt_tokens": 282, 488 "total_tokens": 324, 489 "completion_tokens_details": { 490 "accepted_prediction_tokens": 0, 491 "audio_tokens": 0, 492 "reasoning_tokens": 0, 493 "rejected_prediction_tokens": 0, 494 }, 495 "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, 496 }, 497 }, 498 component_info=ComponentInfo(name="test", type="test"), 499 ), 500 ] 501 502 # Convert chunks to a chat message 503 result = _convert_streaming_chunks_to_chat_message(chunks=chunks) 504 505 assert not result.texts 506 assert not result.text 507 508 # Verify both tool calls were found and processed 509 assert len(result.tool_calls) == 1 510 assert result.tool_calls[0].id == "call_ZOj5l67zhZOx6jqjg7ATQwb6" 511 assert result.tool_calls[0].tool_name == "rag_pipeline_tool" 512 assert result.tool_calls[0].arguments == {"query": "Where does Mark live?"} 513 assert result.meta["finish_reason"] == "tool_calls" 514 515 516 def test_convert_streaming_chunk_to_chat_message_with_empty_tool_call_arguments(): 517 chunks = [ 518 # Message start with input tokens 519 StreamingChunk( 520 content="", 521 meta={ 522 "type": "message_start", 523 "message": { 524 "id": "msg_123", 525 "type": "message", 526 "role": "assistant", 527 "content": [], 528 "model": "claude-sonnet-4-20250514", 529 "stop_reason": None, 530 "stop_sequence": None, 531 "usage": {"input_tokens": 25, "output_tokens": 0}, 532 }, 533 }, 534 index=0, 535 tool_calls=[], 536 tool_call_result=None, 537 start=True, 538 finish_reason=None, 539 ), 540 # Initial text content 541 StreamingChunk( 542 content="", 543 meta={"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, 544 index=1, 545 tool_calls=[], 546 tool_call_result=None, 547 start=True, 548 finish_reason=None, 549 ), 550 StreamingChunk( 551 content="Let me check", 552 meta={"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Let me check"}}, 553 index=2, 554 tool_calls=[], 555 tool_call_result=None, 556 start=False, 557 finish_reason=None, 558 ), 559 StreamingChunk( 560 content=" the weather", 561 meta={"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": " the weather"}}, 562 index=3, 563 tool_calls=[], 564 tool_call_result=None, 565 start=False, 566 finish_reason=None, 567 ), 568 # Tool use content 569 StreamingChunk( 570 content="", 571 meta={ 572 "type": "content_block_start", 573 "index": 1, 574 "content_block": {"type": "tool_use", "id": "toolu_123", "name": "weather", "input": {}}, 575 }, 576 index=5, 577 tool_calls=[ToolCallDelta(index=1, id="toolu_123", tool_name="weather", arguments=None)], 578 tool_call_result=None, 579 start=True, 580 finish_reason=None, 581 ), 582 StreamingChunk( 583 content="", 584 meta={"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": ""}}, 585 index=7, 586 tool_calls=[ToolCallDelta(index=1, id=None, tool_name=None, arguments="")], 587 tool_call_result=None, 588 start=False, 589 finish_reason=None, 590 ), 591 # Final message delta 592 StreamingChunk( 593 content="", 594 meta={ 595 "type": "message_delta", 596 "delta": {"stop_reason": "tool_use", "stop_sequence": None}, 597 "usage": {"completion_tokens": 40}, 598 }, 599 index=8, 600 tool_calls=[], 601 tool_call_result=None, 602 start=False, 603 finish_reason="tool_calls", 604 ), 605 ] 606 607 message = _convert_streaming_chunks_to_chat_message(chunks=chunks) 608 609 assert message.texts == ["Let me check the weather"] 610 assert len(message.tool_calls) == 1 611 assert message.tool_calls[0].arguments == {} 612 assert message.tool_calls[0].id == "toolu_123" 613 assert message.tool_calls[0].tool_name == "weather" 614 615 616 def test_print_streaming_chunk_content_only(): 617 chunk = StreamingChunk( 618 content="Hello, world!", 619 meta={"model": "test-model"}, 620 component_info=ComponentInfo(name="test", type="test"), 621 start=True, 622 ) 623 with patch("builtins.print") as mock_print: 624 print_streaming_chunk(chunk) 625 expected_calls = [call("[ASSISTANT]\n", flush=True, end=""), call("Hello, world!", flush=True, end="")] 626 mock_print.assert_has_calls(expected_calls) 627 628 629 def test_print_streaming_chunk_tool_call(): 630 chunk = StreamingChunk( 631 content="", 632 meta={"model": "test-model"}, 633 component_info=ComponentInfo(name="test", type="test"), 634 start=True, 635 index=0, 636 tool_calls=[ToolCallDelta(id="call_123", tool_name="test_tool", arguments='{"param": "value"}', index=0)], 637 ) 638 with patch("builtins.print") as mock_print: 639 print_streaming_chunk(chunk) 640 expected_calls = [ 641 call("[TOOL CALL]\nTool: test_tool \nArguments: ", flush=True, end=""), 642 call('{"param": "value"}', flush=True, end=""), 643 ] 644 mock_print.assert_has_calls(expected_calls) 645 646 647 def test_print_streaming_chunk_tool_call_result(): 648 chunk = StreamingChunk( 649 content="", 650 meta={"model": "test-model"}, 651 component_info=ComponentInfo(name="test", type="test"), 652 index=0, 653 tool_call_result=ToolCallResult( 654 result="Tool execution completed successfully", 655 origin=ToolCall(id="call_123", tool_name="test_tool", arguments={}), 656 error=False, 657 ), 658 ) 659 with patch("builtins.print") as mock_print: 660 print_streaming_chunk(chunk) 661 expected_calls = [call("[TOOL RESULT]\nTool execution completed successfully", flush=True, end="")] 662 mock_print.assert_has_calls(expected_calls) 663 664 665 def test_print_streaming_chunk_with_finish_reason(): 666 chunk = StreamingChunk( 667 content="Final content.", 668 meta={"model": "test-model"}, 669 component_info=ComponentInfo(name="test", type="test"), 670 start=True, 671 finish_reason="stop", 672 ) 673 with patch("builtins.print") as mock_print: 674 print_streaming_chunk(chunk) 675 expected_calls = [ 676 call("[ASSISTANT]\n", flush=True, end=""), 677 call("Final content.", flush=True, end=""), 678 call("\n\n", flush=True, end=""), 679 ] 680 mock_print.assert_has_calls(expected_calls) 681 682 683 def test_print_streaming_chunk_empty_chunk(): 684 chunk = StreamingChunk( 685 content="", meta={"model": "test-model"}, component_info=ComponentInfo(name="test", type="test") 686 ) 687 with patch("builtins.print") as mock_print: 688 print_streaming_chunk(chunk) 689 mock_print.assert_not_called() 690 691 692 def test_convert_streaming_chunks_to_chat_message_usage_not_in_last_chunk(): 693 """ 694 Test that usage info is correctly extracted even when it's not in the last chunk. 695 This can happen with some API providers like Qwen3 where usage info may be returned 696 in a different chunk than the final one. 697 """ 698 chunks = [ 699 StreamingChunk( 700 content="", 701 meta={"model": "qwen-plus", "index": 0, "finish_reason": None, "received_at": "2025-01-01T00:00:00.000000"}, 702 component_info=ComponentInfo(name="test", type="test"), 703 ), 704 StreamingChunk( 705 content="Hello", 706 meta={"model": "qwen-plus", "index": 0, "finish_reason": None, "received_at": "2025-01-01T00:00:00.100000"}, 707 component_info=ComponentInfo(name="test", type="test"), 708 index=0, 709 start=True, 710 ), 711 StreamingChunk( 712 content=" world", 713 meta={"model": "qwen-plus", "index": 0, "finish_reason": None, "received_at": "2025-01-01T00:00:00.200000"}, 714 component_info=ComponentInfo(name="test", type="test"), 715 index=0, 716 ), 717 # Chunk with usage info (not the last chunk) 718 StreamingChunk( 719 content="", 720 meta={ 721 "model": "qwen-plus", 722 "received_at": "2025-01-01T00:00:00.300000", 723 "usage": {"completion_tokens": 10, "prompt_tokens": 20, "total_tokens": 30}, 724 }, 725 component_info=ComponentInfo(name="test", type="test"), 726 index=None, 727 ), 728 # Final chunk with finish_reason but no usage (simulating Qwen3 behavior) 729 StreamingChunk( 730 content="", 731 meta={ 732 "model": "qwen-plus", 733 "index": 0, 734 "finish_reason": "stop", 735 "received_at": "2025-01-01T00:00:00.400000", 736 "usage": None, # No usage info in final chunk 737 }, 738 component_info=ComponentInfo(name="test", type="test"), 739 finish_reason="stop", 740 ), 741 ] 742 743 result = _convert_streaming_chunks_to_chat_message(chunks=chunks) 744 745 assert result.text == "Hello world" 746 assert result.meta["model"] == "qwen-plus" 747 assert result.meta["finish_reason"] == "stop" 748 # Usage should be extracted from the chunk that has it, not just the last chunk 749 assert result.meta["usage"] == {"completion_tokens": 10, "prompt_tokens": 20, "total_tokens": 30} 750 751 752 def test_convert_streaming_chunks_to_chat_message_with_reasoning(): 753 """Test that reasoning content is correctly accumulated from streaming chunks.""" 754 chunks = [ 755 StreamingChunk( 756 content="", 757 meta={"model": "test-model", "received_at": "2025-01-01T00:00:00"}, 758 component_info=ComponentInfo(name="test", type="test"), 759 reasoning=ReasoningContent(reasoning_text="Let me think about this..."), 760 index=0, 761 ), 762 StreamingChunk( 763 content="", 764 meta={"model": "test-model", "received_at": "2025-01-01T00:00:01"}, 765 component_info=ComponentInfo(name="test", type="test"), 766 reasoning=ReasoningContent(reasoning_text=" The capital of France is Paris."), 767 index=0, 768 ), 769 StreamingChunk( 770 content="Paris", 771 meta={"model": "test-model", "received_at": "2025-01-01T00:00:02"}, 772 component_info=ComponentInfo(name="test", type="test"), 773 ), 774 StreamingChunk( 775 content="", 776 meta={"model": "test-model", "received_at": "2025-01-01T00:00:03"}, 777 component_info=ComponentInfo(name="test", type="test"), 778 finish_reason="stop", 779 ), 780 ] 781 782 result = _convert_streaming_chunks_to_chat_message(chunks=chunks) 783 784 assert result.text == "Paris" 785 assert result.reasoning is not None 786 assert isinstance(result.reasoning, ReasoningContent) 787 assert result.reasoning.reasoning_text == "Let me think about this... The capital of France is Paris." 788 assert result.meta["finish_reason"] == "stop" 789 790 791 def test_convert_streaming_chunks_to_chat_message_without_reasoning(): 792 """Test that messages without reasoning work correctly (backward compatibility).""" 793 chunks = [ 794 StreamingChunk( 795 content="Hello", 796 meta={"model": "test-model", "received_at": "2025-01-01T00:00:00"}, 797 component_info=ComponentInfo(name="test", type="test"), 798 ), 799 StreamingChunk( 800 content=" world", 801 meta={"model": "test-model", "received_at": "2025-01-01T00:00:01"}, 802 component_info=ComponentInfo(name="test", type="test"), 803 finish_reason="stop", 804 ), 805 ] 806 807 result = _convert_streaming_chunks_to_chat_message(chunks=chunks) 808 809 assert result.text == "Hello world" 810 assert result.reasoning is None 811 812 813 def test_print_streaming_chunk_with_reasoning(): 814 """Test that print_streaming_chunk handles reasoning content correctly.""" 815 chunk = StreamingChunk( 816 content="", 817 meta={"model": "test-model"}, 818 component_info=ComponentInfo(name="test", type="test"), 819 start=True, 820 reasoning=ReasoningContent(reasoning_text="I am thinking about this question."), 821 index=0, 822 ) 823 with patch("builtins.print") as mock_print: 824 print_streaming_chunk(chunk) 825 expected_calls = [ 826 call("[REASONING]\n", flush=True, end=""), 827 call("I am thinking about this question.", flush=True, end=""), 828 ] 829 mock_print.assert_has_calls(expected_calls) 830 831 832 def test_print_streaming_chunk_with_reasoning_continuation(): 833 """Test that print_streaming_chunk handles reasoning continuation correctly.""" 834 chunk = StreamingChunk( 835 content="", 836 meta={"model": "test-model"}, 837 component_info=ComponentInfo(name="test", type="test"), 838 start=False, # Not the first chunk 839 reasoning=ReasoningContent(reasoning_text="continued reasoning..."), 840 index=0, 841 ) 842 with patch("builtins.print") as mock_print: 843 print_streaming_chunk(chunk) 844 # Should only print the reasoning text without the header since it's a continuation 845 expected_calls = [call("continued reasoning...", flush=True, end="")] 846 mock_print.assert_has_calls(expected_calls)