Skip to content

core / google_gemini_client

core.google_gemini_client

Utility client for interacting with the Google Generative Language REST API.

This small helper wraps the HTTP endpoints used by Gemini so that we can interact with the service without pulling in the google-generativeai package. Using the REST interface keeps stderr free from the gRPC warnings the SDK emits during import/initialisation (e.g. the ALTS creds ignored message that was polluting the CLI output).

GeminiAPIError

Bases: RuntimeError

Raised when the Gemini service reports an error or blocks a prompt.

Source code in core\google_gemini_client.py
21
22
class GeminiAPIError(RuntimeError):
    """Raised when the Gemini service reports an error or blocks a prompt."""

GeminiClient

Lightweight REST client for Gemini models.

Source code in core\google_gemini_client.py
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
class GeminiClient:
    """Lightweight REST client for Gemini models."""

    def __init__(
        self,
        api_key: str,
        *,
        api_base: Optional[str] = None,
        api_version: Optional[str] = None,
        timeout: int = 120,
    ) -> None:
        if not api_key:
            raise ValueError("`api_key` must be a non-empty string.")

        env_base = os.getenv("GOOGLE_API_BASE")
        env_version = os.getenv("GOOGLE_API_VERSION")

        self._api_key = api_key
        self._api_base = (api_base or env_base or DEFAULT_API_BASE).rstrip("/")
        self._api_version = api_version or env_version or DEFAULT_API_VERSION
        self._timeout = timeout

    # ------------------------------------------------------------------
    # Public helpers
    # ------------------------------------------------------------------
    def generate_text(
        self,
        model: str,
        *,
        prompt: str,
        system_prompt: Optional[str] = None,
        temperature: Optional[float] = None,
        max_output_tokens: Optional[int] = None,
    ) -> Dict[str, Any]:
        """Generate text for a purely textual prompt."""
        contents = [
            {
                "role": "user",
                "parts": [{"text": prompt}],
            }
        ]

        generation_config: Dict[str, Any] = {}
        if temperature is not None:
            generation_config["temperature"] = temperature
        if max_output_tokens is not None:
            generation_config["maxOutputTokens"] = max_output_tokens

        payload: Dict[str, Any] = {"contents": contents}
        if system_prompt:
            payload["systemInstruction"] = {
                "parts": [{"text": system_prompt}],
            }
        if generation_config:
            payload["generationConfig"] = generation_config

        response = self._post_json(
            f"{_normalise_model_name(model)}:generateContent", payload
        )
        total_tokens = response.get("usageMetadata", {}).get("totalTokenCount", 0)
        content = self._extract_text(response)

        return {
            "tokens_used": total_tokens,
            "content": content
        }

    def generate_multimodal(
        self,
        model: str,
        *,
        text: str,
        image_bytes: bytes,
        system_prompt: Optional[str] = None,
        temperature: Optional[float] = None,
    ) -> str:
        """Generate text from a prompt that also contains an inline image."""
        inline_data = {
            "mimeType": "image/png",
            "data": base64.b64encode(image_bytes).decode("utf-8"),
        }

        parts: List[Dict[str, Any]] = [{"text": text}, {"inlineData": inline_data}]
        contents = [{"role": "user", "parts": parts}]

        payload: Dict[str, Any] = {"contents": contents}
        if system_prompt:
            payload["systemInstruction"] = {
                "parts": [{"text": system_prompt}],
            }
        if temperature is not None:
            payload["generationConfig"] = {"temperature": temperature}

        response = self._post_json(
            f"{_normalise_model_name(model)}:generateContent", payload
        )
        return self._extract_text(response)

    def embed_text(self, model: str, *, text: str) -> List[float]:
        """Fetch an embedding vector for the supplied text."""
        payload = {
            "content": {
                "parts": [{"text": text}],
            }
        }
        response = self._post_json(
            f"{_normalise_model_name(model)}:embedContent", payload
        )

        embedding = response.get("embedding")
        if isinstance(embedding, dict) and "values" in embedding:
            return list(map(float, embedding.get("values", [])))
        if isinstance(embedding, list):
            return [float(x) for x in embedding]

        raise GeminiAPIError("Gemini embedContent response did not contain embeddings.")

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------
    def _endpoint(self, path: str) -> str:
        return f"{self._api_base}/{self._api_version}/{path.lstrip('/')}"

    def _post_json(self, path: str, payload: Dict[str, Any]) -> Dict[str, Any]:
        response = requests.post(
            self._endpoint(path),
            params={"key": self._api_key},
            json=payload,
            timeout=self._timeout,
        )
        response.raise_for_status()
        return response.json()

    @staticmethod
    def _extract_text(response: Dict[str, Any]) -> str:
        feedback = response.get("promptFeedback")
        if isinstance(feedback, dict):
            reason = feedback.get("blockReason")
            if reason:
                raise GeminiAPIError(f"Prompt blocked by Gemini: {reason}")

        candidates: Iterable[Dict[str, Any]] = response.get("candidates", []) or []
        for candidate in candidates:
            if candidate.get("finishReason") == "SAFETY":
                # Skip candidates halted for safety reasons.
                continue
            content = candidate.get("content") or {}
            parts: Iterable[Dict[str, Any]] = content.get("parts", []) or []
            texts = [part.get("text", "") for part in parts if isinstance(part, dict)]
            text = "".join(texts).strip()
            if text:
                return text

        return ""

generate_text(model, *, prompt, system_prompt=None, temperature=None, max_output_tokens=None)

Generate text for a purely textual prompt.

Source code in core\google_gemini_client.py
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
def generate_text(
    self,
    model: str,
    *,
    prompt: str,
    system_prompt: Optional[str] = None,
    temperature: Optional[float] = None,
    max_output_tokens: Optional[int] = None,
) -> Dict[str, Any]:
    """Generate text for a purely textual prompt."""
    contents = [
        {
            "role": "user",
            "parts": [{"text": prompt}],
        }
    ]

    generation_config: Dict[str, Any] = {}
    if temperature is not None:
        generation_config["temperature"] = temperature
    if max_output_tokens is not None:
        generation_config["maxOutputTokens"] = max_output_tokens

    payload: Dict[str, Any] = {"contents": contents}
    if system_prompt:
        payload["systemInstruction"] = {
            "parts": [{"text": system_prompt}],
        }
    if generation_config:
        payload["generationConfig"] = generation_config

    response = self._post_json(
        f"{_normalise_model_name(model)}:generateContent", payload
    )
    total_tokens = response.get("usageMetadata", {}).get("totalTokenCount", 0)
    content = self._extract_text(response)

    return {
        "tokens_used": total_tokens,
        "content": content
    }

generate_multimodal(model, *, text, image_bytes, system_prompt=None, temperature=None)

Generate text from a prompt that also contains an inline image.

Source code in core\google_gemini_client.py
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def generate_multimodal(
    self,
    model: str,
    *,
    text: str,
    image_bytes: bytes,
    system_prompt: Optional[str] = None,
    temperature: Optional[float] = None,
) -> str:
    """Generate text from a prompt that also contains an inline image."""
    inline_data = {
        "mimeType": "image/png",
        "data": base64.b64encode(image_bytes).decode("utf-8"),
    }

    parts: List[Dict[str, Any]] = [{"text": text}, {"inlineData": inline_data}]
    contents = [{"role": "user", "parts": parts}]

    payload: Dict[str, Any] = {"contents": contents}
    if system_prompt:
        payload["systemInstruction"] = {
            "parts": [{"text": system_prompt}],
        }
    if temperature is not None:
        payload["generationConfig"] = {"temperature": temperature}

    response = self._post_json(
        f"{_normalise_model_name(model)}:generateContent", payload
    )
    return self._extract_text(response)

embed_text(model, *, text)

Fetch an embedding vector for the supplied text.

Source code in core\google_gemini_client.py
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def embed_text(self, model: str, *, text: str) -> List[float]:
    """Fetch an embedding vector for the supplied text."""
    payload = {
        "content": {
            "parts": [{"text": text}],
        }
    }
    response = self._post_json(
        f"{_normalise_model_name(model)}:embedContent", payload
    )

    embedding = response.get("embedding")
    if isinstance(embedding, dict) and "values" in embedding:
        return list(map(float, embedding.get("values", [])))
    if isinstance(embedding, list):
        return [float(x) for x in embedding]

    raise GeminiAPIError("Gemini embedContent response did not contain embeddings.")