Implement thinking content for Gemini (home-assistant#150347)

Co-authored-by: Joost Lekkerkerker <[email protected]>

Author: Shulyaka

homeassistant/components/google_generative_ai_conversation/entity.py

@@ -3,12 +3,13 @@
 from __future__ import annotations
 
 import asyncio
+import base64
 import codecs
 from collections.abc import AsyncGenerator, AsyncIterator, Callable
-from dataclasses import replace
+from dataclasses import dataclass, replace
 import mimetypes
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 from google.genai import Client
 from google.genai.errors import APIError, ClientError
@@ -27,6 +28,7 @@
     PartUnionDict,
     SafetySetting,
     Schema,
+    ThinkingConfig,
     Tool,
     ToolListUnion,
 )
@@ -201,6 +203,30 @@ def _create_google_tool_response_content(
     )
 
 
+@dataclass(slots=True)
+class PartDetails:
+    """Additional data for a content part."""
+
+    part_type: Literal["text", "thought", "function_call"]
+    """The part type for which this data is relevant for."""
+
+    index: int
+    """Start position or number of the tool."""
+
+    length: int = 0
+    """Length of the relevant data."""
+
+    thought_signature: str | None = None
+    """Base64 encoded thought signature, if available."""
+
+
+@dataclass(slots=True)
+class ContentDetails:
+    """Native data for AssistantContent."""
+
+    part_details: list[PartDetails]
+
+
 def _convert_content(
     content: (
         conversation.UserContent
@@ -209,32 +235,91 @@ def _convert_content(
     ),
 ) -> Content:
     """Convert HA content to Google content."""
-    if content.role != "assistant" or not content.tool_calls:
-        role = "model" if content.role == "assistant" else content.role
+    if content.role != "assistant":
         return Content(
-            role=role,
-            parts=[
-                Part.from_text(text=content.content if content.content else ""),
-            ],
+            role=content.role,
+            parts=[Part.from_text(text=content.content if content.content else "")],
         )
 
     # Handle the Assistant content with tool calls.
     assert type(content) is conversation.AssistantContent
     parts: list[Part] = []
+    part_details: list[PartDetails] = (
+        content.native.part_details
+        if isinstance(content.native, ContentDetails)
+        else []
+    )
+    details: PartDetails | None = None
 
     if content.content:
-        parts.append(Part.from_text(text=content.content))
+        index = 0
+        for details in part_details:
+            if details.part_type == "text":
+                if index < details.index:
+                    parts.append(
+                        Part.from_text(text=content.content[index : details.index])
+                    )
+                    index = details.index
+                parts.append(
+                    Part.from_text(
+                        text=content.content[index : index + details.length],
+                    )
+                )
+                if details.thought_signature:
+                    parts[-1].thought_signature = base64.b64decode(
+                        details.thought_signature
+                    )
+                index += details.length
+        if index < len(content.content):
+            parts.append(Part.from_text(text=content.content[index:]))
+
+    if content.thinking_content:
+        index = 0
+        for details in part_details:
+            if details.part_type == "thought":
+                if index < details.index:
+                    parts.append(
+                        Part.from_text(
+                            text=content.thinking_content[index : details.index]
+                        )
+                    )
+                    parts[-1].thought = True
+                    index = details.index
+                parts.append(
+                    Part.from_text(
+                        text=content.thinking_content[index : index + details.length],
+                    )
+                )
+                parts[-1].thought = True
+                if details.thought_signature:
+                    parts[-1].thought_signature = base64.b64decode(
+                        details.thought_signature
+                    )
+                index += details.length
+        if index < len(content.thinking_content):
+            parts.append(Part.from_text(text=content.thinking_content[index:]))
+            parts[-1].thought = True
 
     if content.tool_calls:
-        parts.extend(
-            [
+        for index, tool_call in enumerate(content.tool_calls):
+            parts.append(
                 Part.from_function_call(
                     name=tool_call.tool_name,
                     args=_escape_decode(tool_call.tool_args),
                 )
-                for tool_call in content.tool_calls
-            ]
-        )
+            )
+            if details := next(
+                (
+                    d
+                    for d in part_details
+                    if d.part_type == "function_call" and d.index == index
+                ),
+                None,
+            ):
+                if details.thought_signature:
+                    parts[-1].thought_signature = base64.b64decode(
+                        details.thought_signature
+                    )
 
     return Content(role="model", parts=parts)
 
@@ -243,14 +328,20 @@ async def _transform_stream(
     result: AsyncIterator[GenerateContentResponse],
 ) -> AsyncGenerator[conversation.AssistantContentDeltaDict]:
     new_message = True
+    part_details: list[PartDetails] = []
     try:
         async for response in result:
             LOGGER.debug("Received response chunk: %s", response)
-            chunk: conversation.AssistantContentDeltaDict = {}
 
             if new_message:
-                chunk["role"] = "assistant"
+                if part_details:
+                    yield {"native": ContentDetails(part_details=part_details)}
+                    part_details = []
+                yield {"role": "assistant"}
                 new_message = False
+                content_index = 0
+                thinking_content_index = 0
+                tool_call_index = 0
 
             # According to the API docs, this would mean no candidate is returned, so we can safely throw an error here.
             if response.prompt_feedback or not response.candidates:
@@ -284,23 +375,62 @@ async def _transform_stream(
                 else []
             )
 
-            content = "".join([part.text for part in response_parts if part.text])
-            tool_calls = []
             for part in response_parts:
-                if not part.function_call:
-                    continue
-                tool_call = part.function_call
-                tool_name = tool_call.name if tool_call.name else ""
-                tool_args = _escape_decode(tool_call.args)
-                tool_calls.append(
-                    llm.ToolInput(tool_name=tool_name, tool_args=tool_args)
-                )
+                chunk: conversation.AssistantContentDeltaDict = {}
+
+                if part.text:
+                    if part.thought:
+                        chunk["thinking_content"] = part.text
+                        if part.thought_signature:
+                            part_details.append(
+                                PartDetails(
+                                    part_type="thought",
+                                    index=thinking_content_index,
+                                    length=len(part.text),
+                                    thought_signature=base64.b64encode(
+                                        part.thought_signature
+                                    ).decode("utf-8"),
+                                )
+                            )
+                        thinking_content_index += len(part.text)
+                    else:
+                        chunk["content"] = part.text
+                        if part.thought_signature:
+                            part_details.append(
+                                PartDetails(
+                                    part_type="text",
+                                    index=content_index,
+                                    length=len(part.text),
+                                    thought_signature=base64.b64encode(
+                                        part.thought_signature
+                                    ).decode("utf-8"),
+                                )
+                            )
+                        content_index += len(part.text)
+
+                if part.function_call:
+                    tool_call = part.function_call
+                    tool_name = tool_call.name if tool_call.name else ""
+                    tool_args = _escape_decode(tool_call.args)
+                    chunk["tool_calls"] = [
+                        llm.ToolInput(tool_name=tool_name, tool_args=tool_args)
+                    ]
+                    if part.thought_signature:
+                        part_details.append(
+                            PartDetails(
+                                part_type="function_call",
+                                index=tool_call_index,
+                                thought_signature=base64.b64encode(
+                                    part.thought_signature
+                                ).decode("utf-8"),
+                            )
+                        )
+
+                yield chunk
 
-            if tool_calls:
-                chunk["tool_calls"] = tool_calls
+        if part_details:
+            yield {"native": ContentDetails(part_details=part_details)}
 
-            chunk["content"] = content
-            yield chunk
     except (
         APIError,
         ValueError,
@@ -522,6 +652,7 @@ def create_generate_content_config(self) -> GenerateContentConfig:
                     ),
                 ),
             ],
+            thinking_config=ThinkingConfig(include_thoughts=True),
         )
 
 

tests/components/google_generative_ai_conversation/snapshots/test_conversation.ambr

@@ -0,0 +1,66 @@
+# serializer version: 1
+# name: test_function_call
+  list([
+    Content(
+    parts=[
+      Part(
+        text='Please call the test function'
+      ),
+    ],
+    role='user'
+  ),
+    Content(
+    parts=[
+      Part(
+        text='Hi there!',
+        thought_signature=b'_thought_signature_2'
+      ),
+      Part(
+        text='The user asked me to call a function',
+        thought=True,
+        thought_signature=b'_thought_signature_1'
+      ),
+      Part(
+        function_call=FunctionCall(
+          args={
+            'param1': [
+              'test_value',
+              "param1's value",
+            ],
+            'param2': 2.7
+          },
+          name='test_tool'
+        ),
+        thought_signature=b'_thought_signature_3'
+      ),
+    ],
+    role='model'
+  ),
+    Content(
+    parts=[
+      Part(
+        function_response=FunctionResponse(
+          name='test_tool',
+          response={
+            'result': 'Test response'
+          }
+        )
+      ),
+    ],
+    role='user'
+  ),
+    Content(
+    parts=[
+      Part(
+        text="I've called the ",
+        thought_signature=b'_thought_signature_4'
+      ),
+      Part(
+        text='test function with the provided parameters.',
+        thought_signature=b'_thought_signature_5'
+      ),
+    ],
+    role='model'
+  ),
+  ])
+# ---