|
| 1 | +"""Round-trip Gemini 3 thought signatures through llama-stack's Vertex AI path. |
| 2 | +
|
| 3 | +Gemini 3.x models (for example ``gemini-3-flash`` and ``gemini-3.5-flash``) |
| 4 | +attach a ``thought_signature`` to the first ``functionCall`` part of a |
| 5 | +tool-calling turn. The signature MUST be replayed verbatim on the following |
| 6 | +turn or Gemini rejects the request with HTTP 400. |
| 7 | +
|
| 8 | +llama-stack converts Gemini responses into the OpenAI chat-completion shape |
| 9 | +before they re-enter its own history, and that shape has no field for a |
| 10 | +thought signature, so the signature is dropped and every multi-turn tool call |
| 11 | +against a Gemini 3 model fails. This module monkeypatches llama-stack's |
| 12 | +``vertexai`` converter so the signature survives the round trip. |
| 13 | +
|
| 14 | +Strategy: both patched functions are thin wrappers around the upstream |
| 15 | +originals. We copy none of llama-stack's conversion logic; we only smuggle the |
| 16 | +signature in and out through the opaque tool-call ``id`` (which llama-stack |
| 17 | +round-trips untouched and only ever compares for equality). |
| 18 | +
|
| 19 | +- On the way out (Gemini -> OpenAI): ``_extract_candidate_parts`` produces a |
| 20 | + random tool-call id per ``functionCall`` part. We call the original, then |
| 21 | + re-walk the candidate's parts in the same deterministic order, pair each |
| 22 | + ``functionCall`` part with the tool call the original emitted, and rewrite |
| 23 | + that tool call's id to embed the base64-encoded signature. |
| 24 | +
|
| 25 | +- On the way back (OpenAI -> Gemini): ``_convert_assistant_message`` builds |
| 26 | + the Gemini ``parts``. We call the original, then re-pair each |
| 27 | + ``function_call`` part with its source tool call (same order) and attach the |
| 28 | + decoded signature. |
| 29 | +
|
| 30 | +This file shadows behaviour tied to a specific llama-stack release. Remove it |
| 31 | +once the upstream Vertex AI provider carries thought signatures natively. |
| 32 | +""" |
| 33 | + |
| 34 | +import base64 |
| 35 | +from typing import Any |
| 36 | + |
| 37 | +from log import get_logger |
| 38 | + |
| 39 | +logger = get_logger(__name__) |
| 40 | + |
| 41 | +# Sentinel separating the real tool-call id from a smuggled Gemini |
| 42 | +# thought_signature. Chosen to be vanishingly unlikely in a normal id. |
| 43 | +_THOUGHT_SIG_SEP = "::gts::" |
| 44 | + |
| 45 | +# Set once the patch has been applied so repeated startup calls are no-ops. |
| 46 | +_PATCH_APPLIED = False |
| 47 | + |
| 48 | + |
| 49 | +def _encode_thought_signature_into_id(call_id: str, signature: Any) -> str: |
| 50 | + """Append a base64-encoded Gemini thought_signature to a tool-call id. |
| 51 | +
|
| 52 | + The signature is bytes; the id must stay a plain string that round-trips |
| 53 | + through llama-stack history. Returns ``call_id`` unchanged when there is no |
| 54 | + signature to carry or it cannot be encoded. |
| 55 | + """ |
| 56 | + if not signature: |
| 57 | + return call_id |
| 58 | + try: |
| 59 | + raw = signature.encode("utf-8") if isinstance(signature, str) else bytes(signature) |
| 60 | + encoded = base64.b64encode(raw).decode("ascii") |
| 61 | + except (TypeError, ValueError): |
| 62 | + return call_id |
| 63 | + return f"{call_id}{_THOUGHT_SIG_SEP}{encoded}" |
| 64 | + |
| 65 | + |
| 66 | +def _decode_thought_signature_from_id(call_id: str) -> bytes | None: |
| 67 | + """Recover the thought_signature bytes smuggled into a tool-call id.""" |
| 68 | + if not call_id or _THOUGHT_SIG_SEP not in call_id: |
| 69 | + return None |
| 70 | + _, _, encoded = call_id.partition(_THOUGHT_SIG_SEP) |
| 71 | + try: |
| 72 | + return base64.b64decode(encoded) |
| 73 | + except (ValueError, TypeError): |
| 74 | + return None |
| 75 | + |
| 76 | + |
| 77 | +def _iter_function_call_parts(candidate: Any) -> list[Any]: |
| 78 | + """Return the ``functionCall`` parts of a Gemini candidate, in order. |
| 79 | +
|
| 80 | + Mirrors the iteration order llama-stack's ``_extract_candidate_parts`` uses |
| 81 | + so the parts line up one-to-one with the tool calls it produces. |
| 82 | + """ |
| 83 | + content_obj = getattr(candidate, "content", None) |
| 84 | + parts = getattr(content_obj, "parts", None) or [] |
| 85 | + fc_parts: list[Any] = [] |
| 86 | + for part in parts: |
| 87 | + # Thinking parts and text parts are skipped before the function-call |
| 88 | + # branch upstream; replicate that ordering precisely. |
| 89 | + if getattr(part, "thought", None): |
| 90 | + continue |
| 91 | + if getattr(part, "text", None) is not None: |
| 92 | + continue |
| 93 | + if getattr(part, "function_call", None) is not None: |
| 94 | + fc_parts.append(part) |
| 95 | + return fc_parts |
| 96 | + |
| 97 | + |
| 98 | +def apply_patch() -> bool: |
| 99 | + """Monkeypatch the Vertex AI converter to carry Gemini thought signatures. |
| 100 | +
|
| 101 | + Idempotent. Returns ``True`` if the patch is in effect after the call, |
| 102 | + ``False`` if the converter module could not be imported (for example when |
| 103 | + the Vertex AI provider is not installed), in which case nothing is changed. |
| 104 | + """ |
| 105 | + global _PATCH_APPLIED |
| 106 | + if _PATCH_APPLIED: |
| 107 | + return True |
| 108 | + |
| 109 | + try: |
| 110 | + from llama_stack.providers.remote.inference.vertexai import ( |
| 111 | + converters, |
| 112 | + ) |
| 113 | + except ImportError: |
| 114 | + logger.info( |
| 115 | + "Vertex AI converter not importable; skipping Gemini thought-signature patch" |
| 116 | + ) |
| 117 | + return False |
| 118 | + |
| 119 | + original_extract = converters._extract_candidate_parts |
| 120 | + original_convert_assistant = converters._convert_assistant_message |
| 121 | + |
| 122 | + def patched_extract_candidate_parts(candidate: Any) -> Any: |
| 123 | + text_parts, thinking_parts, tool_calls = original_extract(candidate) |
| 124 | + if not tool_calls: |
| 125 | + return text_parts, thinking_parts, tool_calls |
| 126 | + fc_parts = _iter_function_call_parts(candidate) |
| 127 | + # The original emits exactly one tool call per function-call part, in |
| 128 | + # the same order. Pair them and embed any signature into the id. |
| 129 | + for tool_call, part in zip(tool_calls, fc_parts): |
| 130 | + signature = getattr(part, "thought_signature", None) |
| 131 | + if not signature: |
| 132 | + continue |
| 133 | + tool_call.id = _encode_thought_signature_into_id(tool_call.id, signature) |
| 134 | + return text_parts, thinking_parts, tool_calls |
| 135 | + |
| 136 | + def patched_convert_assistant_message(msg: dict[str, Any]) -> dict[str, Any] | None: |
| 137 | + result = original_convert_assistant(msg) |
| 138 | + if result is None: |
| 139 | + return None |
| 140 | + tool_calls = msg.get("tool_calls") or [] |
| 141 | + if not tool_calls: |
| 142 | + return result |
| 143 | + # Re-pair each Gemini function_call part with its source tool call, in |
| 144 | + # order, and attach the decoded signature. The original appends one |
| 145 | + # function_call part per tool call after any leading text part, so we |
| 146 | + # walk the function_call parts and the tool calls together. |
| 147 | + fc_parts = [p for p in result.get("parts", []) if "function_call" in p] |
| 148 | + for part, tool_call in zip(fc_parts, tool_calls): |
| 149 | + call_id = converters._to_dict(tool_call).get("id", "") |
| 150 | + signature = _decode_thought_signature_from_id(call_id) |
| 151 | + if signature is not None: |
| 152 | + part["thought_signature"] = signature |
| 153 | + return result |
| 154 | + |
| 155 | + converters._extract_candidate_parts = patched_extract_candidate_parts |
| 156 | + converters._convert_assistant_message = patched_convert_assistant_message |
| 157 | + _PATCH_APPLIED = True |
| 158 | + logger.info("Applied Gemini 3 thought-signature patch to Vertex AI converter") |
| 159 | + return True |
0 commit comments