diff --git a/libs/core/langchain_core/messages/block_translators/ollama.py b/libs/core/langchain_core/messages/block_translators/ollama.py index a0f41ab7634..a1874d4c726 100644 --- a/libs/core/langchain_core/messages/block_translators/ollama.py +++ b/libs/core/langchain_core/messages/block_translators/ollama.py @@ -1,17 +1,152 @@ """Derivations of standard content blocks from Ollama content.""" +from typing import cast +from uuid import uuid4 + from langchain_core.messages import AIMessage, AIMessageChunk from langchain_core.messages import content as types def translate_content(message: AIMessage) -> list[types.ContentBlock]: - """Derive standard content blocks from a message with Ollama content.""" - raise NotImplementedError + """Derive standard content blocks from a message with Ollama v0 content.""" + content_blocks: list[types.ContentBlock] = [] + + # First, handle reasoning content from additional_kwargs. There will be only one + if "reasoning_content" in message.additional_kwargs: + reasoning_content = message.additional_kwargs["reasoning_content"] + # `reasoning_content` is only ever a str + if reasoning_content: + content_blocks.append( + cast( + "types.ReasoningContentBlock", + { + "type": "reasoning", + "reasoning": reasoning_content, + }, + ) + ) + + # Handle main content + if message.content and isinstance(message.content, str): + content_blocks.append( + cast("types.TextContentBlock", {"type": "text", "text": message.content}) + ) + elif isinstance(message.content, list): + for item in message.content: + if isinstance(item, dict): + if item.get("type") == "text": + block = {"type": "text", "text": item["text"]} + + if item.get("id"): + block["id"] = item["id"] + + content_blocks.append( + cast( + "types.TextContentBlock", + block, + ) + ) + else: + # Keep other content types as-is (multimodal, etc.) + # These should be handled during `_normalize_messages` + content_blocks.append(cast("types.ContentBlock", item)) + elif isinstance(item, str): + content_blocks.append( + cast("types.TextContentBlock", {"type": "text", "text": item}) + ) + + # Handle tool calls + content_blocks.extend( + [ + cast( + "types.ToolCall", + { + "type": "tool_call", + "id": tool_call.get("id", str(uuid4())), + "name": tool_call["name"], + "args": tool_call["args"], + }, + ) + for tool_call in message.tool_calls + ] + ) + + return content_blocks def translate_content_chunk(message: AIMessageChunk) -> list[types.ContentBlock]: """Derive standard content blocks from a message chunk with Ollama content.""" - raise NotImplementedError + content_blocks: list[types.ContentBlock] = [] + + if "reasoning_content" in message.additional_kwargs: + reasoning_content = message.additional_kwargs["reasoning_content"] + if reasoning_content: + content_blocks.append( + cast( + "types.ReasoningContentBlock", + { + "type": "reasoning", + "reasoning": reasoning_content, + }, + ) + ) + + # Handle main content + if isinstance(message.content, str) and message.content: + content_blocks.append( + cast("types.TextContentBlock", {"type": "text", "text": message.content}) + ) + elif isinstance(message.content, list): + for item in message.content: + if isinstance(item, dict): + if item.get("type") == "text": + block = {"type": "text", "text": item["text"]} + + if item.get("id"): + block["id"] = item["id"] + + content_blocks.append( + cast( + "types.TextContentBlock", + block, + ) + ) + else: + content_blocks.append(cast("types.ContentBlock", item)) + elif isinstance(item, str): + content_blocks.append( + cast("types.TextContentBlock", {"type": "text", "text": item}) + ) + + # Handle tool calls + content_blocks.extend( + [ + cast( + "types.ToolCall", + { + "type": "tool_call", + "id": tool_call.get("id", str(uuid4())), + "name": tool_call["name"], + "args": tool_call["args"], + }, + ) + for tool_call in message.tool_calls + ] + ) + + # Handle tool call chunks + for tool_call_chunk in message.tool_call_chunks: + tc: types.ToolCallChunk = { + "type": "tool_call_chunk", + "id": tool_call_chunk.get("id"), + "name": tool_call_chunk.get("name"), + "args": tool_call_chunk.get("args"), + } + if (idx := tool_call_chunk.get("index")) is not None: + tc["index"] = idx + content_blocks.append(tc) + + return content_blocks def _register_ollama_translator() -> None: diff --git a/libs/partners/ollama/langchain_ollama/chat_models.py b/libs/partners/ollama/langchain_ollama/chat_models.py index e2f2b290a77..e7be8ff119d 100644 --- a/libs/partners/ollama/langchain_ollama/chat_models.py +++ b/libs/partners/ollama/langchain_ollama/chat_models.py @@ -33,6 +33,7 @@ from langchain_core.messages import ( SystemMessage, ToolCall, ToolMessage, + ensure_id, is_data_content_block, ) from langchain_core.messages.ai import UsageMetadata @@ -59,6 +60,64 @@ from typing_extensions import Self, is_typeddict from ._utils import validate_model + +def _ensure_content_blocks_have_ids(content_blocks: list) -> list: + """Ensure all content blocks have IDs.""" + updated_blocks = [] + for block in content_blocks: + if isinstance(block, dict) and "id" not in block: + updated_block = dict(block) + updated_block["id"] = ensure_id(None) + updated_blocks.append(updated_block) + else: + updated_blocks.append(block) + return updated_blocks + + +def _update_ollama_message_to_v1(message: AIMessage) -> AIMessage: + """Update Ollama message to v1 format. + + Exclude `reasoning_content` from `additional_kwargs` since it is now part of + `content_blocks`. + + Ensure `response_metadata` has `output_version` set to `v1` for future use. + + """ + content_blocks_with_ids = _ensure_content_blocks_have_ids(message.content_blocks) + return message.model_copy( + update={ + "content": content_blocks_with_ids, + "additional_kwargs": { + k: v + for k, v in message.additional_kwargs.items() + if k != "reasoning_content" + }, + "response_metadata": { + **message.response_metadata, + "output_version": "v1", + }, + } + ) + + +def _update_ollama_message_chunk_to_v1(message: AIMessageChunk) -> AIMessageChunk: + content_blocks_with_ids = _ensure_content_blocks_have_ids(message.content_blocks) + return message.model_copy( + update={ + "content": content_blocks_with_ids, + "additional_kwargs": { + k: v + for k, v in message.additional_kwargs.items() + if k != "reasoning_content" + }, + "response_metadata": { + **message.response_metadata, + "output_version": "v1", + }, + } + ) + + log = logging.getLogger(__name__) @@ -817,13 +876,19 @@ class ChatOllama(BaseChatModel): messages, stop, run_manager, verbose=self.verbose, **kwargs ) generation_info = final_chunk.generation_info + message = AIMessage( + content=final_chunk.text, + usage_metadata=cast(AIMessageChunk, final_chunk.message).usage_metadata, + tool_calls=cast(AIMessageChunk, final_chunk.message).tool_calls, + additional_kwargs=final_chunk.message.additional_kwargs, + response_metadata={"model_provider": "ollama"}, + ) + + if self.output_version == "v1": + message = _update_ollama_message_to_v1(message) + chat_generation = ChatGeneration( - message=AIMessage( - content=final_chunk.text, - usage_metadata=cast(AIMessageChunk, final_chunk.message).usage_metadata, - tool_calls=cast(AIMessageChunk, final_chunk.message).tool_calls, - additional_kwargs=final_chunk.message.additional_kwargs, - ), + message=message, generation_info=generation_info, ) return ChatResult(generations=[chat_generation]) @@ -883,6 +948,7 @@ class ChatOllama(BaseChatModel): stream_resp ), tool_calls=_get_tool_calls_from_response(stream_resp), + response_metadata={"model_provider": "ollama"}, ), generation_info=generation_info, ) @@ -897,6 +963,14 @@ class ChatOllama(BaseChatModel): **kwargs: Any, ) -> Iterator[ChatGenerationChunk]: for chunk in self._iterate_over_stream(messages, stop, **kwargs): + if self.output_version == "v1": + chunk_message = cast(AIMessageChunk, chunk.message) + converted_message = _update_ollama_message_chunk_to_v1(chunk_message) + chunk = ChatGenerationChunk( + message=converted_message, + generation_info=chunk.generation_info, + ) + if run_manager: run_manager.on_llm_new_token( chunk.text, @@ -959,6 +1033,7 @@ class ChatOllama(BaseChatModel): stream_resp ), tool_calls=_get_tool_calls_from_response(stream_resp), + response_metadata={"model_provider": "ollama"}, ), generation_info=generation_info, ) @@ -973,6 +1048,14 @@ class ChatOllama(BaseChatModel): **kwargs: Any, ) -> AsyncIterator[ChatGenerationChunk]: async for chunk in self._aiterate_over_stream(messages, stop, **kwargs): + if self.output_version == "v1": + chunk_message = cast(AIMessageChunk, chunk.message) + converted_message = _update_ollama_message_chunk_to_v1(chunk_message) + chunk = ChatGenerationChunk( + message=converted_message, + generation_info=chunk.generation_info, + ) + if run_manager: await run_manager.on_llm_new_token( chunk.text, @@ -991,13 +1074,19 @@ class ChatOllama(BaseChatModel): messages, stop, run_manager, verbose=self.verbose, **kwargs ) generation_info = final_chunk.generation_info + message = AIMessage( + content=final_chunk.text, + usage_metadata=cast(AIMessageChunk, final_chunk.message).usage_metadata, + tool_calls=cast(AIMessageChunk, final_chunk.message).tool_calls, + additional_kwargs=final_chunk.message.additional_kwargs, + response_metadata={"model_provider": "ollama"}, + ) + + if self.output_version == "v1": + message = _update_ollama_message_to_v1(message) + chat_generation = ChatGeneration( - message=AIMessage( - content=final_chunk.text, - usage_metadata=cast(AIMessageChunk, final_chunk.message).usage_metadata, - tool_calls=cast(AIMessageChunk, final_chunk.message).tool_calls, - additional_kwargs=final_chunk.message.additional_kwargs, - ), + message=message, generation_info=generation_info, ) return ChatResult(generations=[chat_generation]) diff --git a/libs/partners/ollama/tests/unit_tests/test_v1_chat_models.py b/libs/partners/ollama/tests/unit_tests/test_v1_chat_models.py new file mode 100644 index 00000000000..48ea3c74490 --- /dev/null +++ b/libs/partners/ollama/tests/unit_tests/test_v1_chat_models.py @@ -0,0 +1,428 @@ +"""Unit tests for ChatOllama v1 format support.""" + +from typing import Any +from unittest.mock import MagicMock, patch + +from langchain_core.messages import ( + AIMessage, + AIMessageChunk, + HumanMessage, +) +from langchain_core.messages.block_translators.ollama import ( + translate_content, + translate_content_chunk, +) +from langchain_core.utils.utils import LC_AUTO_PREFIX + +from langchain_ollama.chat_models import ChatOllama + +MODEL_NAME = "llama3.1" + + +class TestV1BlockTranslator: + """Test block translator functions.""" + + def test_translate_content_text_only(self) -> None: + """Test translation of text message to v1 format.""" + + message = AIMessage( + content="Hello, world!", + ) + + blocks = translate_content(message) + + assert len(blocks) == 1 + assert blocks[0]["type"] == "text" + assert blocks[0]["text"] == "Hello, world!" + assert "id" not in blocks[0] # ID should not be added during translation + + def test_translate_content_with_reasoning(self) -> None: + """Test translation of message with reasoning to v1 format.""" + + message = AIMessage( + content="The answer is 42.", + additional_kwargs={"reasoning_content": "Let me think about this..."}, + ) + + blocks = translate_content(message) + + assert len(blocks) == 2 + + # Reasoning should come before main content + reasoning_block = blocks[0] + assert reasoning_block["type"] == "reasoning" + assert reasoning_block.get("reasoning") == "Let me think about this..." + assert "id" not in reasoning_block + + text_block = blocks[1] + assert text_block["type"] == "text" + assert text_block["text"] == "The answer is 42." + + def test_translate_content_with_tool_calls(self) -> None: + """Test translation of message with tool calls to v1 format.""" + + tool_call = { + "name": "multiply", + "args": {"a": 3, "b": 4}, + "id": "call_123", + } + + message = AIMessage( + content="I'll multiply these numbers.", + tool_calls=[tool_call], + ) + + blocks = translate_content(message) + + assert len(blocks) == 2 + + text_block = blocks[0] + assert text_block["type"] == "text" + assert text_block["text"] == "I'll multiply these numbers." + + tool_call_block = blocks[1] + assert tool_call_block["type"] == "tool_call" + assert tool_call_block["name"] == "multiply" + assert tool_call_block["args"] == {"a": 3, "b": 4} + assert tool_call_block["id"] == "call_123" + + def test_translate_content_chunk(self) -> None: + """Test translation of chunk to v1 format.""" + + chunk = AIMessageChunk( + content="Hello", + additional_kwargs={"reasoning_content": "Thinking..."}, + ) + + blocks = translate_content_chunk(chunk) + + assert len(blocks) == 2 + + reasoning_block = blocks[0] + assert reasoning_block["type"] == "reasoning" + assert reasoning_block.get("reasoning") == "Thinking..." + + text_block = blocks[1] + assert text_block["type"] == "text" + assert text_block["text"] == "Hello" + + +class TestV1ChatOllama: + """Test ChatOllama with v1 format functionality.""" + + def test_v1_parameter(self) -> None: + llm = ChatOllama(model=MODEL_NAME) + assert llm.output_version == "v0" + + llm = ChatOllama(model=MODEL_NAME, output_version="v1") + assert llm.output_version == "v1" + + def test_v0_output_format(self) -> None: + """Test that previous output format is retained.""" + mock_response = [ + { + "model": "test-model", + "message": { + "role": "assistant", + "content": "Hello, world!", + "thinking": "I should greet the user.", + }, + "done": True, + "done_reason": "stop", + } + ] + + with patch("langchain_ollama.chat_models.Client") as mock_client_class: + mock_client = MagicMock() + mock_client_class.return_value = mock_client + mock_client.chat.return_value = mock_response + + llm = ChatOllama(model=MODEL_NAME, reasoning=True) + result = llm.invoke([HumanMessage("Hello")]) + + # Should be v0 format - reasoning in additional_kwargs + assert isinstance(result.content, str) + assert result.content == "Hello, world!" + reasoning_content = result.additional_kwargs.get("reasoning_content") + assert reasoning_content == "I should greet the user." + + # Test using `.content_blocks` + content_blocks = result.content_blocks + assert isinstance(content_blocks, list) + assert len(content_blocks) == 2 + + # First block should be reasoning + assert content_blocks[0]["type"] == "reasoning" + assert content_blocks[0].get("reasoning") == "I should greet the user." + # assert "reasoning_content" not in result.additional_kwargs TODO: use this? + + assert content_blocks[1]["type"] == "text" + assert content_blocks[1]["text"] == "Hello, world!" + + # ID should not be added unless flag is v1 + assert "id" not in content_blocks[1] + + def test_v1_output_format(self) -> None: + mock_response = [ + { + "model": "test-model", + "message": { + "role": "assistant", + "content": "Hello, world!", + "thinking": "I should greet the user.", + }, + "done": True, + "done_reason": "stop", + } + ] + + with patch("langchain_ollama.chat_models.Client") as mock_client_class: + mock_client = MagicMock() + mock_client_class.return_value = mock_client + mock_client.chat.return_value = mock_response + + llm = ChatOllama(model=MODEL_NAME, output_version="v1", reasoning=True) + result = llm.invoke([HumanMessage("Hello")]) + + # Should be stored in v1 format (content blocks) since flag is set + assert isinstance(result.content, list) + assert len(result.content) == 2 + + reasoning_block = result.content[0] + assert isinstance(reasoning_block, dict) + assert reasoning_block["type"] == "reasoning" + assert reasoning_block["reasoning"] == "I should greet the user." + assert "id" in reasoning_block + assert reasoning_block["id"] is not None + assert type(reasoning_block["id"]) is str + assert reasoning_block["id"].startswith(LC_AUTO_PREFIX) + + text_block = result.content[1] + assert isinstance(text_block, dict) + assert text_block["type"] == "text" + assert text_block["text"] == "Hello, world!" + assert "id" in text_block + assert text_block["id"] is not None + assert type(text_block["id"]) is str + assert text_block["id"].startswith(LC_AUTO_PREFIX) + + assert "reasoning_content" not in result.additional_kwargs + + def test_v1_streaming_output_format(self) -> None: + """Test that v1 format works with streaming.""" + mock_responses = [ + { + "model": "test-model", + "message": {"role": "assistant", "content": "Hello"}, + "done": False, + }, + { + "model": "test-model", + "message": {"role": "assistant", "content": ", world!"}, + "done": True, + "done_reason": "stop", + }, + ] + + with patch("langchain_ollama.chat_models.Client") as mock_client_class: + mock_client = MagicMock() + mock_client_class.return_value = mock_client + mock_client.chat.return_value = mock_responses + + llm = ChatOllama(model=MODEL_NAME, output_version="v1") + chunks = list(llm.stream([HumanMessage("Hello")])) + + for chunk in chunks: + assert isinstance(chunk.content, list) + if chunk.content: # Skip empty chunks + content_block = chunk.content[0] + assert isinstance(content_block, dict) + assert content_block["type"] == "text" + + # TODO: should chunks have auto-generated IDs? what to do here? + + def test_v1_with_tool_calls(self) -> None: + """Test v1 format with tool calls.""" + mock_response = [ + { + "model": "test-model", + "message": { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "function": { + "name": "multiply", + "arguments": {"a": 3, "b": 4}, + } + } + ], + }, + "done": True, + "done_reason": "stop", + } + ] + + with patch("langchain_ollama.chat_models.Client") as mock_client_class: + mock_client = MagicMock() + mock_client_class.return_value = mock_client + mock_client.chat.return_value = mock_response + + llm = ChatOllama(model=MODEL_NAME, output_version="v1") + result = llm.invoke([HumanMessage("Multiply 3 and 4")]) + + assert isinstance(result.content, list) + tool_call_block = None + for block in result.content: + if isinstance(block, dict) and block.get("type") == "tool_call": + tool_call_block = block + break + assert tool_call_block is not None + assert tool_call_block["name"] == "multiply" + assert tool_call_block["args"] == {"a": 3, "b": 4} + + async def test_v1_async_generation(self) -> None: + """Test v1 format with async generation.""" + mock_response = [ + { + "model": "test-model", + "message": { + "role": "assistant", + "content": "Hello, async world!", + "thinking": "Async thinking...", + }, + "done": True, + "done_reason": "stop", + } + ] + + with patch("langchain_ollama.chat_models.AsyncClient") as mock_client_class: + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + # The chat method should return a coroutine that yields the async iterator + async def async_chat_coroutine(*args: Any, **kwargs: Any) -> Any: + async def async_generator() -> Any: + for response in mock_response: + yield response + + return async_generator() + + mock_client.chat = async_chat_coroutine + + llm = ChatOllama(model=MODEL_NAME, output_version="v1", reasoning=True) + result = await llm.ainvoke([HumanMessage("Hello")]) + + # Should be v1 format + assert isinstance(result.content, list) + assert len(result.content) == 2 + + reasoning_block = result.content[0] + assert isinstance(reasoning_block, dict) + assert reasoning_block["type"] == "reasoning" + assert reasoning_block["reasoning"] == "Async thinking..." + + text_block = result.content[1] + assert isinstance(text_block, dict) + assert text_block["type"] == "text" + assert text_block["text"] == "Hello, async world!" + + async def test_v1_async_streaming(self) -> None: + """Test v1 format with async streaming.""" + mock_responses = [ + { + "model": "test-model", + "message": {"role": "assistant", "content": "Async"}, + "done": False, + }, + { + "model": "test-model", + "message": {"role": "assistant", "content": " streaming!"}, + "done": True, + "done_reason": "stop", + }, + ] + + with patch("langchain_ollama.chat_models.AsyncClient") as mock_client_class: + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + async def async_chat_coroutine(*args: Any, **kwargs: Any) -> Any: + async def async_generator() -> Any: + for response in mock_responses: + yield response + + return async_generator() + + mock_client.chat = async_chat_coroutine + + llm = ChatOllama(model=MODEL_NAME, output_version="v1") + chunks = [chunk async for chunk in llm.astream([HumanMessage("Hello")])] + + # Each chunk should have v1 format + for chunk in chunks: + # chunk is an AIMessageChunk directly + assert isinstance(chunk.content, list) + if chunk.content: # Skip empty chunks + content_block = chunk.content[0] + assert isinstance(content_block, dict) + assert content_block["type"] == "text" + + +class TestV1EdgeCases: + """Test edge cases for v1 format support.""" + + def test_empty_content_handling(self) -> None: + """Test handling of empty content.""" + + message = AIMessage(content="", additional_kwargs={}) + blocks = translate_content(message) + + assert isinstance(blocks, list) + assert len(blocks) == 0 + + def test_list_content_passthrough(self) -> None: + """Test that existing list content is handled properly.""" + + message = AIMessage(content=[{"type": "text", "text": "Already a list"}]) + blocks = translate_content(message) + + assert isinstance(blocks, list) + assert len(blocks) == 1 + assert blocks[0]["type"] == "text" + assert blocks[0]["text"] == "Already a list" + + # IDs should be preserved if they exist + message = AIMessage( + content=[{"type": "text", "text": "Already a list", "id": "existing_id"}] + ) + blocks = translate_content(message) + + assert isinstance(blocks, list) + assert len(blocks) == 1 + assert blocks[0].get("id") == "existing_id" + + def test_multimodal_content_preservation(self) -> None: + """Test that multimodal content is preserved in v1 conversion.""" + + image_block = {"type": "image", "url": "https://example.com/image.png"} + + message = AIMessage(content=[image_block]) + blocks = translate_content(message) + + assert isinstance(blocks, list) + assert len(blocks) == 1 + assert blocks[0] == image_block + + def test_unknown_content_block_preservation(self) -> None: + """Test that unknown content block types are preserved.""" + # TODO: check this, shouldn't this become NonStandardContentBlock? + + unknown_block = {"type": "custom_block", "custom_field": "custom_value"} + + message = AIMessage(content=[unknown_block]) + blocks = translate_content(message) + + assert isinstance(blocks, list) + assert len(blocks) == 1 + assert blocks[0] == unknown_block