vertexai[patch]: support code_execution (#959)

Author: ccurme

libs/vertexai/langchain_google_vertexai/chat_models.py

@@ -114,9 +114,11 @@
 from google.cloud.aiplatform_v1beta1.types import (
     Blob,
     Candidate,
+    CodeExecutionResult,
     Part,
     HarmCategory,
     Content,
+    ExecutableCode,
     FileData,
     FunctionCall,
     FunctionResponse,
@@ -242,6 +244,29 @@ def _convert_to_prompt(part: Union[str, Dict]) -> Optional[Part]:
                 return Part(text=part["text"])
             else:
                 return None
+        if part["type"] == "executable_code":
+            if "executable_code" not in part or "language" not in part:
+                raise ValueError(
+                    "Executable code part must have 'code' and 'language' keys, got "
+                    f"{part}"
+                )
+            return Part(
+                executable_code=ExecutableCode(
+                    language=part["language"], code=part["executable_code"]
+                )
+            )
+        if part["type"] == "code_execution_result":
+            if "code_execution_result" not in part or "outcome" not in part:
+                raise ValueError(
+                    "Code execution result part must have 'code_execution_result' and "
+                    f"'outcome' keys, got {part}"
+                )
+            return Part(
+                code_execution_result=CodeExecutionResult(
+                    output=part["code_execution_result"], outcome=part["outcome"]
+                )
+            )
+
         if is_data_content_block(part):
             # LangChain standard format
             if part["type"] == "image" and part["source_type"] == "url":
@@ -542,7 +567,7 @@ def _parse_response_candidate(
 def _parse_response_candidate(
     response_candidate: "Candidate", streaming: bool = False
 ) -> AIMessage:
-    content: Union[None, str, List[str]] = None
+    content: Union[None, str, List[Union[str, dict[str, Any]]]] = None
     additional_kwargs = {}
     tool_calls = []
     invalid_tool_calls = []
@@ -610,6 +635,44 @@ def _parse_response_candidate(
                             error=str(e),
                         )
                     )
+        if hasattr(part, "executable_code") and part.executable_code is not None:
+            if part.executable_code.code and part.executable_code.language:
+                code_message = {
+                    "type": "executable_code",
+                    "executable_code": part.executable_code.code,
+                    "language": part.executable_code.language,
+                }
+                if not content:
+                    content = [code_message]
+                elif isinstance(content, str):
+                    content = [content, code_message]
+                elif isinstance(content, list):
+                    content.append(code_message)
+                else:
+                    raise Exception("Unexpected content type")
+
+        if (
+            hasattr(part, "code_execution_result")
+            and part.code_execution_result is not None
+        ):
+            if part.code_execution_result.output and part.code_execution_result.outcome:
+                execution_result = {
+                    "type": "code_execution_result",
+                    # Name output -> code_execution_result for consistency with
+                    # langchain-google-genai
+                    "code_execution_result": part.code_execution_result.output,
+                    "outcome": part.code_execution_result.outcome,
+                }
+
+                if not content:
+                    content = [execution_result]
+                elif isinstance(content, str):
+                    content = [content, execution_result]
+                elif isinstance(content, list):
+                    content.append(execution_result)
+                else:
+                    raise Exception("Unexpected content type")
+
     if content is None:
         content = ""
 
@@ -896,16 +959,30 @@ class GetPopulation(BaseModel):
 
         See ``ChatVertexAI.bind_tools()`` method for more.
 
-    Use Search with Gemini 2:
+    Built-in search:
         .. code-block:: python
 
             from google.cloud.aiplatform_v1beta1.types import Tool as VertexTool
+            from langchain_google_vertexai import ChatVertexAI
+
             llm = ChatVertexAI(model="gemini-2.0-flash-exp")
             resp = llm.invoke(
                 "When is the next total solar eclipse in US?",
                 tools=[VertexTool(google_search={})],
             )
 
+    Built-in code execution:
+        .. code-block:: python
+
+            from google.cloud.aiplatform_v1beta1.types import Tool as VertexTool
+            from langchain_google_vertexai import ChatVertexAI
+
+            llm = ChatVertexAI(model="gemini-2.0-flash-exp")
+            resp = llm.invoke(
+                "What is 3^3?",
+                tools=[VertexTool(code_execution={})],
+            )
+
     Structured output:
         .. code-block:: python
 

libs/vertexai/langchain_google_vertexai/functions_utils.py

@@ -32,6 +32,7 @@
 )
 from langchain_core.utils.json_schema import dereference_refs
 from pydantic import BaseModel
+from typing_extensions import NotRequired
 
 logger = logging.getLogger(__name__)
 
@@ -49,13 +50,15 @@
 ]
 _GoogleSearchLike = Union[gapic.Tool.GoogleSearch, Dict[str, Any]]
 _RetrievalLike = Union[gapic.Retrieval, Dict[str, Any]]
+_CodeExecutionLike = Union[gapic.Tool.CodeExecution, Dict[str, Any]]
 
 
 class _ToolDictLike(TypedDict):
     function_declarations: Optional[List[_FunctionDeclarationLike]]
     google_search_retrieval: Optional[_GoogleSearchRetrievalLike]
     google_search: Optional[_GoogleSearchLike]
     retrieval: Optional[_RetrievalLike]
+    code_execution: NotRequired[_CodeExecutionLike]
 
 
 _ToolType = Union[gapic.Tool, vertexai.Tool, _ToolDictLike, _FunctionDeclarationLike]
@@ -306,6 +309,8 @@ def _format_to_gapic_tool(tools: _ToolsType) -> gapic.Tool:
                 gapic_tool.function_declarations.extend(rt.function_declarations)
             if "google_search" in rt:
                 gapic_tool.google_search = rt.google_search
+            if "code_execution" in rt:
+                gapic_tool.code_execution = rt.code_execution
         elif isinstance(tool, dict):
             # not _ToolDictLike
             if not any(
@@ -315,6 +320,7 @@ def _format_to_gapic_tool(tools: _ToolsType) -> gapic.Tool:
                     "google_search_retrieval",
                     "google_search",
                     "retrieval",
+                    "code_execution",
                 ]
             ):
                 fd = _format_to_gapic_function_declaration(tool)
@@ -345,6 +351,10 @@ def _format_to_gapic_tool(tools: _ToolsType) -> gapic.Tool:
                 )
             if "retrieval" in tool:
                 gapic_tool.retrieval = gapic.Retrieval(tool["retrieval"])
+            if "code_execution" in tool:
+                gapic_tool.code_execution = gapic.Tool.CodeExecution(
+                    tool["code_execution"]
+                )
         else:
             fd = _format_to_gapic_function_declaration(tool)
             gapic_tool.function_declarations.append(fd)

libs/vertexai/tests/integration_tests/test_chat_models.py

@@ -1390,3 +1390,57 @@ class People(BaseModel):
     response = llm_with_tools.invoke("Chester, no hair color provided.")
     assert isinstance(response, AIMessage)
     assert response.tool_calls[0]["name"] == "People"
+
+
+def test_search_builtin() -> None:
+    llm = ChatVertexAI(model=_DEFAULT_MODEL_NAME).bind_tools([{"google_search": {}}])
+    input_message = {
+        "role": "user",
+        "content": "What is today's news?",
+    }
+    response = llm.invoke([input_message])
+    assert "grounding_metadata" in response.response_metadata
+
+    # Test streaming
+    full: Optional[BaseMessageChunk] = None
+    for chunk in llm.stream([input_message]):
+        assert isinstance(chunk, AIMessageChunk)
+        full = chunk if full is None else full + chunk
+    assert isinstance(full, AIMessageChunk)
+    assert "grounding_metadata" in full.response_metadata
+
+    # Test we can process chat history
+    next_message = {
+        "role": "user",
+        "content": "Tell me more about that last story.",
+    }
+    _ = llm.invoke([input_message, full, next_message])
+
+
+def test_code_execution_builtin() -> None:
+    llm = ChatVertexAI(model=_DEFAULT_MODEL_NAME).bind_tools([{"code_execution": {}}])
+    input_message = {
+        "role": "user",
+        "content": "What is 3^3?",
+    }
+    response = llm.invoke([input_message])
+    content_blocks = [block for block in response.content if isinstance(block, dict)]
+    expected_block_types = {"executable_code", "code_execution_result"}
+    assert set(block.get("type") for block in content_blocks) == expected_block_types
+
+    # Test streaming
+    full: Optional[BaseMessageChunk] = None
+    for chunk in llm.stream([input_message]):
+        assert isinstance(chunk, AIMessageChunk)
+        full = chunk if full is None else full + chunk
+    assert isinstance(full, AIMessageChunk)
+    content_blocks = [block for block in full.content if isinstance(block, dict)]
+    expected_block_types = {"executable_code", "code_execution_result"}
+    assert set(block.get("type") for block in content_blocks) == expected_block_types
+
+    # Test we can process chat history
+    next_message = {
+        "role": "user",
+        "content": "Can you add some comments to the code?",
+    }
+    _ = llm.invoke([input_message, full, next_message])