content blocks in (a)invoke

This commit is contained in:
Chester Curme 2025-03-11 16:43:31 -04:00
parent 03035b8d73
commit 7445085ac2
2 changed files with 57 additions and 28 deletions

View File

@ -991,16 +991,25 @@ class BaseChatOpenAI(BaseChatModel):
token_usage = response.usage.model_dump() if response.usage else {} token_usage = response.usage.model_dump() if response.usage else {}
generation_info = {} generation_info = {}
content_blocks = []
for output in response.output: for output in response.output:
if output.type == "message": if output.type == "message":
joined = "".join( for content in output.content:
content.text if content.type == "output_text":
for content in output.content block = {
if content.type == "output_text" "type": "text",
) "text": content.text,
"annotations": [
annotation.model_dump()
for annotation in content.annotations
],
}
content_blocks.append(block)
usage_metadata = _create_usage_metadata_responses(token_usage) usage_metadata = _create_usage_metadata_responses(token_usage)
message = AIMessage( message = AIMessage(
content=joined, id=output.id, usage_metadata=usage_metadata content=content_blocks, # type: ignore[arg-type]
id=output.id,
usage_metadata=usage_metadata,
) )
if output.status: if output.status:
generation_info["status"] = output.status generation_info["status"] = output.status

View File

@ -1237,7 +1237,17 @@ def test_web_search() -> None:
tools=[{"type": "web_search_preview"}], tools=[{"type": "web_search_preview"}],
) )
assert isinstance(response, AIMessage) assert isinstance(response, AIMessage)
assert response.content assert isinstance(response.content, list)
for block in response.content:
assert isinstance(block, dict)
if block["type"] == "text":
assert isinstance(block["text"], str)
for annotation in block["annotations"]:
for key in ["end_index", "start_index", "title", "type", "url"]:
assert key in annotation
text_content = response.text()
assert isinstance(text_content, str)
assert text_content
assert response.usage_metadata assert response.usage_metadata
assert response.usage_metadata["input_tokens"] > 0 assert response.usage_metadata["input_tokens"] > 0
assert response.usage_metadata["output_tokens"] > 0 assert response.usage_metadata["output_tokens"] > 0
@ -1246,16 +1256,16 @@ def test_web_search() -> None:
assert response.response_metadata["status"] assert response.response_metadata["status"]
# Test streaming # Test streaming
full: Optional[BaseMessageChunk] = None # full: Optional[BaseMessageChunk] = None
for chunk in llm.stream( # for chunk in llm.stream(
"What was a positive news story from today?", # "What was a positive news story from today?",
tools=[{"type": "web_search_preview"}], # tools=[{"type": "web_search_preview"}],
): # ):
assert isinstance(chunk, AIMessageChunk) # assert isinstance(chunk, AIMessageChunk)
full = chunk if full is None else full + chunk # full = chunk if full is None else full + chunk
assert isinstance(full, AIMessageChunk) # assert isinstance(full, AIMessageChunk)
assert full.content # assert full.content
assert full.usage_metadata # assert full.usage_metadata
async def test_web_search_async() -> None: async def test_web_search_async() -> None:
@ -1265,7 +1275,17 @@ async def test_web_search_async() -> None:
tools=[{"type": "web_search_preview"}], tools=[{"type": "web_search_preview"}],
) )
assert isinstance(response, AIMessage) assert isinstance(response, AIMessage)
assert response.content assert isinstance(response.content, list)
for block in response.content:
assert isinstance(block, dict)
if block["type"] == "text":
assert isinstance(block["text"], str)
for annotation in block["annotations"]:
for key in ["end_index", "start_index", "title", "type", "url"]:
assert key in annotation
text_content = response.text()
assert isinstance(text_content, str)
assert text_content
assert response.usage_metadata assert response.usage_metadata
assert response.usage_metadata["input_tokens"] > 0 assert response.usage_metadata["input_tokens"] > 0
assert response.usage_metadata["output_tokens"] > 0 assert response.usage_metadata["output_tokens"] > 0
@ -1273,13 +1293,13 @@ async def test_web_search_async() -> None:
assert response.response_metadata["model_name"] assert response.response_metadata["model_name"]
assert response.response_metadata["status"] assert response.response_metadata["status"]
full: Optional[BaseMessageChunk] = None # full: Optional[BaseMessageChunk] = None
async for chunk in llm.astream( # async for chunk in llm.astream(
"What was a positive news story from today?", # "What was a positive news story from today?",
tools=[{"type": "web_search_preview"}], # tools=[{"type": "web_search_preview"}],
): # ):
assert isinstance(chunk, AIMessageChunk) # assert isinstance(chunk, AIMessageChunk)
full = chunk if full is None else full + chunk # full = chunk if full is None else full + chunk
assert isinstance(full, AIMessageChunk) # assert isinstance(full, AIMessageChunk)
assert full.content # assert full.content
assert full.usage_metadata # assert full.usage_metadata