Code execution (#396)

* Modify _make_tool to accept code_execution proto object

* Make adding code execution compatible with current tests

* Update test_content.py to test code_execution

* Test cases running and passing for code execution updates

* Updated _join_contents to handle executable_code and code_execution_results

* Updated .text field to include executable code and code execution results

* Update versions.

Change-Id: I4e2c9081a9de9105bfc1d9d83769d98faee4d18e

* format

Change-Id: I38de8c9fe434785608ad5e04a2234bd584e2de03

---------

Co-authored-by: Mark Daoust <[email protected]>

Author: shilpakancharla

google/generativeai/types/content_types.py

@@ -623,24 +623,40 @@ def _encode_fd(fd: FunctionDeclaration | protos.FunctionDeclaration) -> protos.F
 class Tool:
     """A wrapper for `protos.Tool`, Contains a collection of related `FunctionDeclaration` objects."""
 
-    def __init__(self, function_declarations: Iterable[FunctionDeclarationType]):
+    def __init__(
+        self,
+        function_declarations: Iterable[FunctionDeclarationType] | None = None,
+        code_execution: protos.CodeExecution | None = None,
+    ):
         # The main path doesn't use this but is seems useful.
-        self._function_declarations = [_make_function_declaration(f) for f in function_declarations]
-        self._index = {}
-        for fd in self._function_declarations:
-            name = fd.name
-            if name in self._index:
-                raise ValueError("")
-            self._index[fd.name] = fd
+        if function_declarations:
+            self._function_declarations = [
+                _make_function_declaration(f) for f in function_declarations
+            ]
+            self._index = {}
+            for fd in self._function_declarations:
+                name = fd.name
+                if name in self._index:
+                    raise ValueError("")
+                self._index[fd.name] = fd
+        else:
+            # Consistent fields
+            self._function_declarations = []
+            self._index = {}
 
         self._proto = protos.Tool(
-            function_declarations=[_encode_fd(fd) for fd in self._function_declarations]
+            function_declarations=[_encode_fd(fd) for fd in self._function_declarations],
+            code_execution=code_execution,
         )
 
     @property
     def function_declarations(self) -> list[FunctionDeclaration | protos.FunctionDeclaration]:
         return self._function_declarations
 
+    @property
+    def code_execution(self) -> protos.CodeExecution:
+        return self._proto.code_execution
+
     def __getitem__(
         self, name: str | protos.FunctionCall
     ) -> FunctionDeclaration | protos.FunctionDeclaration:
@@ -673,13 +689,24 @@ def _make_tool(tool: ToolType) -> Tool:
     if isinstance(tool, Tool):
         return tool
     elif isinstance(tool, protos.Tool):
-        return Tool(function_declarations=tool.function_declarations)
+        if "code_execution" in tool:
+            code_execution = tool.code_execution
+        else:
+            code_execution = None
+        return Tool(function_declarations=tool.function_declarations, code_execution=code_execution)
     elif isinstance(tool, dict):
-        if "function_declarations" in tool:
+        if "function_declarations" in tool or "code_execution" in tool:
             return Tool(**tool)
         else:
             fd = tool
             return Tool(function_declarations=[protos.FunctionDeclaration(**fd)])
+    elif isinstance(tool, str):
+        if tool.lower() == "code_execution":
+            return Tool(code_execution=protos.CodeExecution())
+        else:
+            raise ValueError("The only string that can be passed as a tool is 'code_execution'.")
+    elif isinstance(tool, protos.CodeExecution):
+        return Tool(code_execution=tool)
     elif isinstance(tool, Iterable):
         return Tool(function_declarations=tool)
     else:
@@ -734,7 +761,12 @@ def to_proto(self):
 
 
 def _make_tools(tools: ToolsType) -> list[Tool]:
-    if isinstance(tools, Iterable) and not isinstance(tools, Mapping):
+    if isinstance(tools, str):
+        if tools.lower() == "code_execution":
+            return [_make_tool(tools)]
+        else:
+            raise ValueError("The only string that can be passed as a tool is 'code_execution'.")
+    elif isinstance(tools, Iterable) and not isinstance(tools, Mapping):
         tools = [_make_tool(t) for t in tools]
         if len(tools) > 1 and all(len(t.function_declarations) == 1 for t in tools):
             # flatten into a single tool.

google/generativeai/types/generation_types.py

@@ -261,26 +261,50 @@ def _join_contents(contents: Iterable[protos.Content]):
     for content in contents:
         parts.extend(content.parts)
 
-    merged_parts = [parts.pop(0)]
-    for part in parts:
-        if not merged_parts[-1].text:
-            merged_parts.append(part)
+    merged_parts = []
+    last = parts[0]
+    for part in parts[1:]:
+        if "text" in last and "text" in part:
+            last = protos.Part(text=last.text + part.text)
             continue
 
-        if not part.text:
-            merged_parts.append(part)
+        # Can we merge the new thing into last?
+        # If not, put last in list of parts, and new thing becomes last
+        if "executable_code" in last and "executable_code" in part:
+            last = protos.Part(
+                executable_code=_join_executable_code(last.executable_code, part.executable_code)
+            )
             continue
 
-        merged_part = protos.Part(merged_parts[-1])
-        merged_part.text += part.text
-        merged_parts[-1] = merged_part
+        if "code_execution_result" in last and "code_execution_result" in part:
+            last = protos.Part(
+                code_execution_result=_join_code_execution_result(
+                    last.code_execution_result, part.code_execution_result
+                )
+            )
+            continue
+
+        merged_parts.append(last)
+        last = part
+
+    merged_parts.append(last)
 
     return protos.Content(
         role=role,
         parts=merged_parts,
     )
 
 
+def _join_executable_code(code_1, code_2):
+    return protos.ExecutableCode(language=code_1.language, code=code_1.code + code_2.code)
+
+
+def _join_code_execution_result(result_1, result_2):
+    return protos.CodeExecutionResult(
+        outcome=result_2.outcome, output=result_1.output + result_2.output
+    )
+
+
 def _join_candidates(candidates: Iterable[protos.Candidate]):
     candidates = tuple(candidates)
 
@@ -413,13 +437,35 @@ def text(self):
                 "Invalid operation: The `response.text` quick accessor requires the response to contain a valid `Part`, "
                 "but none were returned. Please check the `candidate.safety_ratings` to determine if the response was blocked."
             )
-        if len(parts) != 1 or "text" not in parts[0]:
-            raise ValueError(
-                "Invalid operation: The `response.text` quick accessor requires a simple (single-`Part`) text response. "
-                "This response is not simple text. Please use the `result.parts` accessor or the full "
-                "`result.candidates[index].content.parts` lookup instead."
-            )
-        return parts[0].text
+
+        texts = []
+        for part in parts:
+            if "text" in part:
+                texts.append(part.text)
+                continue
+            if "executable_code" in part:
+                language = part.executable_code.language.name.lower()
+                if language == "language_unspecified":
+                    language = ""
+                else:
+                    language = f" {language}"
+                texts.extend([f"```{language}", part.executable_code.code, "```"])
+                continue
+            if "code_execution_result" in part:
+                outcome_result = part.code_execution_result.outcome.name.lower().replace(
+                    "outcome_", ""
+                )
+                if outcome_result == "ok" or outcome_result == "unspecified":
+                    outcome_result = ""
+                else:
+                    outcome_result = f" {outcome_result}"
+                texts.extend([f"```{outcome_result}", part.code_execution_result.output, "```"])
+                continue
+
+            part_type = protos.Part.pb(part).whichOneof("data")
+            raise ValueError(f"Could not convert `part.{part_type}` to text.")
+
+        return "\n".join(texts)
 
     @property
     def prompt_feedback(self):

google/generativeai/version.py

@@ -14,4 +14,4 @@
 # limitations under the License.
 from __future__ import annotations
 
-__version__ = "0.7.0"
+__version__ = "0.7.1"

setup.py

@@ -42,7 +42,7 @@ def get_version():
     release_status = "Development Status :: 5 - Production/Stable"
 
 dependencies = [
-    "google-ai-generativelanguage==0.6.5",
+    "google-ai-generativelanguage==0.6.6",
     "google-api-core",
     "google-api-python-client",
     "google-auth>=2.15.0",  # 2.15 adds API key auth support

tests/test_content.py

@@ -15,7 +15,7 @@
 import dataclasses
 import pathlib
 import typing_extensions
-from typing import Any, Union
+from typing import Any, Union, Iterable
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -367,7 +367,7 @@ def test_to_tools(self, tools):
             raise ValueError("This shouldn't happen")
         tools = function_library.to_proto()
 
-        tools = type(tools[0]).to_dict(tools[0])
+        tools = type(tools[0]).to_dict(tools[0], including_default_value_fields=False)
         tools["function_declarations"][0].pop("parameters", None)
 
         expected = dict(
@@ -378,6 +378,24 @@ def test_to_tools(self, tools):
 
         self.assertEqual(tools, expected)
 
+    @parameterized.named_parameters(
+        ["string", "code_execution"],
+        ["proto_object", protos.CodeExecution()],
+        ["proto_passed_in", protos.Tool(code_execution=protos.CodeExecution())],
+        ["empty_dictionary", {"code_execution": {}}],
+        ["string_list", ["code_execution"]],
+        ["proto_object_list", [protos.CodeExecution()]],
+        ["proto_passed_in_list", [protos.Tool(code_execution=protos.CodeExecution())]],
+        ["empty_dictionary_list", [{"code_execution": {}}]],
+    )
+    def test_code_execution(self, tools):
+        if isinstance(tools, Iterable):
+            t = content_types._make_tools(tools)
+            self.assertIsInstance(t[0].code_execution, protos.CodeExecution)
+        else:
+            t = content_types._make_tool(tools)  # Pass code execution into tools
+            self.assertIsInstance(t.code_execution, protos.CodeExecution)
+
     def test_two_fun_is_one_tool(self):
         def a():
             pass

tests/test_generation.py

@@ -124,6 +124,61 @@ def test_join_contents(self):
 
         self.assertEqual(expected, type(result).to_dict(result))
 
+    def test_join_parts(self):
+        contents = [
+            protos.Content(role="assistant", parts=[protos.Part(text="A")]),
+            protos.Content(role="assistant", parts=[protos.Part(text="B")]),
+            protos.Content(role="assistant", parts=[protos.Part(executable_code={"code": "C"})]),
+            protos.Content(role="assistant", parts=[protos.Part(executable_code={"code": "D"})]),
+            protos.Content(
+                role="assistant", parts=[protos.Part(code_execution_result={"output": "E"})]
+            ),
+            protos.Content(
+                role="assistant", parts=[protos.Part(code_execution_result={"output": "F"})]
+            ),
+            protos.Content(role="assistant", parts=[protos.Part(text="G")]),
+            protos.Content(role="assistant", parts=[protos.Part(text="H")]),
+        ]
+        g = generation_types._join_contents(contents=contents)
+        expected = protos.Content(
+            role="assistant",
+            parts=[
+                protos.Part(text="AB"),
+                protos.Part(executable_code={"code": "CD"}),
+                protos.Part(code_execution_result={"output": "EF"}),
+                protos.Part(text="GH"),
+            ],
+        )
+        self.assertEqual(expected, g)
+
+    def test_code_execution_text(self):
+        content = protos.Content(
+            role="assistant",
+            parts=[
+                protos.Part(text="AB"),
+                protos.Part(executable_code={"language": "PYTHON", "code": "CD"}),
+                protos.Part(code_execution_result={"outcome": "OUTCOME_OK", "output": "EF"}),
+                protos.Part(text="GH"),
+            ],
+        )
+        response = generation_types.GenerateContentResponse(
+            done=True,
+            iterator=None,
+            result=protos.GenerateContentResponse({"candidates": [{"content": content}]}),
+        )
+        expected = textwrap.dedent(
+            """\
+            AB
+            ``` python
+            CD
+            ```
+            ```
+            EF
+            ```
+            GH"""
+        )
+        self.assertEqual(expected, response.text)
+
     def test_many_join_contents(self):
         import string