3
3
import json
4
4
import os
5
5
from base64 import b64encode
6
- from typing import Optional
6
+ from typing import Optional , cast
7
7
8
8
import httpx
9
9
import pytest
@@ -42,7 +42,10 @@ def test_stream() -> None:
42
42
chunks_with_model_name = 0
43
43
for token in llm .stream ("I'm Pickle Rick" ):
44
44
assert isinstance (token .content , str )
45
- full = token if full is None else full + token
45
+ if full is None :
46
+ full = cast (BaseMessageChunk , token )
47
+ else :
48
+ full = full + token
46
49
assert isinstance (token , AIMessageChunk )
47
50
if token .usage_metadata is not None :
48
51
if token .usage_metadata .get ("input_tokens" ):
@@ -81,7 +84,10 @@ async def test_astream() -> None:
81
84
chunks_with_output_token_counts = 0
82
85
async for token in llm .astream ("I'm Pickle Rick" ):
83
86
assert isinstance (token .content , str )
84
- full = token if full is None else full + token
87
+ if full is None :
88
+ full = cast (BaseMessageChunk , token )
89
+ else :
90
+ full = full + token
85
91
assert isinstance (token , AIMessageChunk )
86
92
if token .usage_metadata is not None :
87
93
if token .usage_metadata .get ("input_tokens" ):
@@ -697,7 +703,10 @@ def test_citations() -> None:
697
703
# Test streaming
698
704
full : Optional [BaseMessageChunk ] = None
699
705
for chunk in llm .stream (messages ):
700
- full = chunk if full is None else full + chunk
706
+ if full is None :
707
+ full = cast (BaseMessageChunk , chunk )
708
+ else :
709
+ full = full + chunk
701
710
assert isinstance (full , AIMessageChunk )
702
711
assert isinstance (full .content , list )
703
712
assert any ("citations" in block for block in full .content )
@@ -722,7 +731,10 @@ def test_thinking() -> None:
722
731
# Test streaming
723
732
full : Optional [BaseMessageChunk ] = None
724
733
for chunk in llm .stream ("Hello" ):
725
- full = chunk if full is None else full + chunk
734
+ if full is None :
735
+ full = cast (BaseMessageChunk , chunk )
736
+ else :
737
+ full = full + chunk
726
738
assert isinstance (full , AIMessageChunk )
727
739
assert isinstance (full .content , list )
728
740
assert any ("thinking" in block for block in full .content )
@@ -756,7 +768,10 @@ def test_redacted_thinking() -> None:
756
768
# Test streaming
757
769
full : Optional [BaseMessageChunk ] = None
758
770
for chunk in llm .stream (query ):
759
- full = chunk if full is None else full + chunk
771
+ if full is None :
772
+ full = cast (BaseMessageChunk , chunk )
773
+ else :
774
+ full = full + chunk
760
775
assert isinstance (full , AIMessageChunk )
761
776
assert isinstance (full .content , list )
762
777
stream_has_reasoning = False
0 commit comments