llm-gemini 0.3__py3-none-any.whl → 0.4__py3-none-any.whl

{llm_gemini-0.3.dist-info → llm_gemini-0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: llm-gemini
-Version: 0.3
+Version: 0.4
 Summary: LLM plugin to access Google's Gemini family of models
 Author: Simon Willison
 License: Apache-2.0
@@ -16,6 +16,7 @@ Requires-Dist: httpx
 Requires-Dist: ijson
 Provides-Extra: test
 Requires-Dist: pytest ; extra == 'test'
+Requires-Dist: pytest-recording ; extra == 'test'
 
 # llm-gemini
 
@@ -35,13 +36,13 @@ llm install llm-gemini
 ## Usage
 
 Configure the model by setting a key called "gemini" to your [API key](https://aistudio.google.com/app/apikey):
-
 ```bash
 llm keys set gemini
 ```
 ```
 <paste key here>
 ```
+You can also set the API key by assigning it to the environment variable `LLM_GEMINI_KEY`.
 
 Now run the model using `-m gemini-1.5-pro-latest`, for example:
 
@@ -55,6 +56,12 @@ llm -m gemini-1.5-pro-latest "A joke about a pelican and a walrus"
 >
 > The pelican taps its beak thoughtfully. "I believe," it says, "it's a billfish."
 
+Other models are:
+
+- `gemini-1.5-flash-latest`
+- `gemini-1.5-flash-8b-latest` - the least expensive
+- `gemini-exp-1114` - recent experimental
+
 ### Images, audio and video
 
 Gemini models are multi-modal. You can provide images, audio or video files as input like this:
@@ -78,8 +85,22 @@ And video:
 ```bash
 llm -m gemini-1.5-pro-latest 'describe what happens' -a video.mp4
 ```
+The Gemini prompting guide includes [extensive advice](https://ai.google.dev/gemini-api/docs/file-prompting-strategies) on multi-modal prompting.
+
+### JSON output
+
+Use `-o json_object 1` to force the output to be JSON:
+
+```bash
+llm -m gemini-1.5-flash-latest -o json_object 1 \
+  '3 largest cities in California, list of {"name": "..."}'
+```
+Outputs:
+```json
+{"cities": [{"name": "Los Angeles"}, {"name": "San Diego"}, {"name": "San Jose"}]}
+```
 
-## Code execution
+### Code execution
 
 Gemini models can [write and execute code](https://ai.google.dev/gemini-api/docs/code-execution) - they can decide to write Python code, execute it in a secure sandbox and use the result as part of their response.
 
@@ -98,13 +119,7 @@ To chat interactively with the model, run `llm chat`:
 llm chat -m gemini-1.5-pro-latest
 ```
 
-Other models are:
-
-- `gemini-1.5-flash-latest`
-- gemini-1.5-flash-8b-latest` - the least expensive
-
-
-### Embeddings
+## Embeddings
 
 The plugin also adds support for the `text-embedding-004` embedding model.
 

llm_gemini-0.4.dist-info/RECORD

@@ -0,0 +1,7 @@
+llm_gemini.py,sha256=_7yQ14ffRpyK3ChAOc6M2ufylg2kUKTQK4C3VKOiUgM,10672
+llm_gemini-0.4.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+llm_gemini-0.4.dist-info/METADATA,sha256=rxQq1JCu1dvMwuVSXy8Tn1txFBdJX1NjCjdFJ6VEa1g,4862
+llm_gemini-0.4.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
+llm_gemini-0.4.dist-info/entry_points.txt,sha256=n544bpgUPIBc5l_cnwsTxPc3gMGJHPtAyqBNp-CkMWk,26
+llm_gemini-0.4.dist-info/top_level.txt,sha256=WUQmG6_2QKbT_8W4HH93qyKl_0SUteL4Ra6_PhyNGKU,11
+llm_gemini-0.4.dist-info/RECORD,,

{llm_gemini-0.3.dist-info → llm_gemini-0.4.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.2.0)
+Generator: setuptools (75.5.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

llm_gemini.py

@@ -4,9 +4,6 @@ import llm
 from pydantic import Field
 from typing import Optional
 
-import urllib.parse
-
-# We disable all of these to avoid random unexpected errors
 SAFETY_SETTINGS = [
     {
         "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
@@ -29,15 +26,20 @@ SAFETY_SETTINGS = [
 
 @llm.hookimpl
 def register_models(register):
-    register(GeminiPro("gemini-pro"))
-    register(GeminiPro("gemini-1.5-pro-latest"))
-    register(GeminiPro("gemini-1.5-flash-latest"))
-    register(GeminiPro("gemini-1.5-pro-001"))
-    register(GeminiPro("gemini-1.5-flash-001"))
-    register(GeminiPro("gemini-1.5-pro-002"))
-    register(GeminiPro("gemini-1.5-flash-002"))
-    register(GeminiPro("gemini-1.5-flash-8b-latest"))
-    register(GeminiPro("gemini-1.5-flash-8b-001"))
+    # Register both sync and async versions of each model
+    for model_id in [
+        "gemini-pro",
+        "gemini-1.5-pro-latest",
+        "gemini-1.5-flash-latest",
+        "gemini-1.5-pro-001",
+        "gemini-1.5-flash-001",
+        "gemini-1.5-pro-002",
+        "gemini-1.5-flash-002",
+        "gemini-1.5-flash-8b-latest",
+        "gemini-1.5-flash-8b-001",
+        "gemini-exp-1114",
+    ]:
+        register(GeminiPro(model_id), AsyncGeminiPro(model_id))
 
 
 def resolve_type(attachment):
@@ -48,7 +50,7 @@ def resolve_type(attachment):
     return mime_type
 
 
-class GeminiPro(llm.Model):
+class _SharedGemini:
     needs_key = "gemini"
     key_env_var = "LLM_GEMINI_KEY"
     can_stream = True
@@ -80,6 +82,7 @@ class GeminiPro(llm.Model):
         "video/webm",
         "video/wmv",
         "video/3gpp",
+        "video/quicktime",
     )
 
     class Options(llm.Options):
@@ -88,7 +91,11 @@ class GeminiPro(llm.Model):
             default=None,
         )
         temperature: Optional[float] = Field(
-            description="Controls the randomness of the output. Use higher values for more creative responses, and lower values for more deterministic responses.",
+            description=(
+                "Controls the randomness of the output. Use higher values for "
+                "more creative responses, and lower values for more "
+                "deterministic responses."
+            ),
             default=None,
             ge=0.0,
             le=2.0,
@@ -98,16 +105,30 @@ class GeminiPro(llm.Model):
             default=None,
         )
         top_p: Optional[float] = Field(
-            description="Changes how the model selects tokens for output. Tokens are selected from the most to least probable until the sum of their probabilities equals the topP value.",
+            description=(
+                "Changes how the model selects tokens for output. Tokens are "
+                "selected from the most to least probable until the sum of "
+                "their probabilities equals the topP value."
+            ),
             default=None,
             ge=0.0,
             le=1.0,
         )
         top_k: Optional[int] = Field(
-            description="Changes how the model selects tokens for output. A topK of 1 means the selected token is the most probable among all the tokens in the model's vocabulary, while a topK of 3 means that the next token is selected from among the 3 most probable using the temperature.",
+            description=(
+                "Changes how the model selects tokens for output. A topK of 1 "
+                "means the selected token is the most probable among all the "
+                "tokens in the model's vocabulary, while a topK of 3 means "
+                "that the next token is selected from among the 3 most "
+                "probable using the temperature."
+            ),
             default=None,
             ge=1,
         )
+        json_object: Optional[bool] = Field(
+            description="Output a valid JSON object {...}",
+            default=None,
+        )
 
     def __init__(self, model_id):
         self.model_id = model_id
@@ -127,11 +148,14 @@ class GeminiPro(llm.Model):
                             }
                         }
                     )
-                parts.append({"text": response.prompt.prompt})
+                if response.prompt.prompt:
+                    parts.append({"text": response.prompt.prompt})
                 messages.append({"role": "user", "parts": parts})
                 messages.append({"role": "model", "parts": [{"text": response.text()}]})
 
-        parts = [{"text": prompt.prompt}]
+        parts = []
+        if prompt.prompt:
+            parts.append({"text": prompt.prompt})
         for attachment in prompt.attachments:
             mime_type = resolve_type(attachment)
             parts.append(
@@ -146,14 +170,7 @@ class GeminiPro(llm.Model):
         messages.append({"role": "user", "parts": parts})
         return messages
 
-    def execute(self, prompt, stream, response, conversation):
-        key = self.get_key()
-        url = "https://generativelanguage.googleapis.com/v1beta/models/{}:streamGenerateContent?".format(
-            self.model_id
-        ) + urllib.parse.urlencode(
-            {"key": key}
-        )
-        gathered = []
+    def build_request_body(self, prompt, conversation):
         body = {
             "contents": self.build_messages(prompt, conversation),
             "safetySettings": SAFETY_SETTINGS,
@@ -169,7 +186,9 @@ class GeminiPro(llm.Model):
             "top_p": "topP",
             "top_k": "topK",
         }
-        # If any of those are set in prompt.options...
+        if prompt.options and prompt.options.json_object:
+            body["generationConfig"] = {"response_mime_type": "application/json"}
+
         if any(
             getattr(prompt.options, key, None) is not None for key in config_map.keys()
         ):
@@ -180,10 +199,30 @@ class GeminiPro(llm.Model):
                     generation_config[other_key] = config_value
             body["generationConfig"] = generation_config
 
+        return body
+
+    def process_part(self, part):
+        if "text" in part:
+            return part["text"]
+        elif "executableCode" in part:
+            return f'```{part["executableCode"]["language"].lower()}\n{part["executableCode"]["code"].strip()}\n```\n'
+        elif "codeExecutionResult" in part:
+            return f'```\n{part["codeExecutionResult"]["output"].strip()}\n```\n'
+        return ""
+
+
+class GeminiPro(_SharedGemini, llm.Model):
+    def execute(self, prompt, stream, response, conversation):
+        key = self.get_key()
+        url = f"https://generativelanguage.googleapis.com/v1beta/models/{self.model_id}:streamGenerateContent"
+        gathered = []
+        body = self.build_request_body(prompt, conversation)
+
         with httpx.stream(
             "POST",
             url,
             timeout=None,
+            headers={"x-goog-api-key": key},
             json=body,
         ) as http_response:
             events = ijson.sendable_list()
@@ -196,14 +235,7 @@ class GeminiPro(llm.Model):
                         raise llm.ModelError(event["error"]["message"])
                     try:
                         part = event["candidates"][0]["content"]["parts"][0]
-                        if "text" in part:
-                            yield part["text"]
-                        elif "executableCode" in part:
-                            # For code_execution
-                            yield f'```{part["executableCode"]["language"].lower()}\n{part["executableCode"]["code"].strip()}\n```\n'
-                        elif "codeExecutionResult" in part:
-                            # For code_execution
-                            yield f'```\n{part["codeExecutionResult"]["output"].strip()}\n```\n'
+                        yield self.process_part(part)
                     except KeyError:
                         yield ""
                     gathered.append(event)
@@ -211,6 +243,39 @@ class GeminiPro(llm.Model):
         response.response_json = gathered
 
 
+class AsyncGeminiPro(_SharedGemini, llm.AsyncModel):
+    async def execute(self, prompt, stream, response, conversation):
+        key = self.get_key()
+        url = f"https://generativelanguage.googleapis.com/v1beta/models/{self.model_id}:streamGenerateContent"
+        gathered = []
+        body = self.build_request_body(prompt, conversation)
+
+        async with httpx.AsyncClient() as client:
+            async with client.stream(
+                "POST",
+                url,
+                timeout=None,
+                headers={"x-goog-api-key": key},
+                json=body,
+            ) as http_response:
+                events = ijson.sendable_list()
+                coro = ijson.items_coro(events, "item")
+                async for chunk in http_response.aiter_bytes():
+                    coro.send(chunk)
+                    if events:
+                        event = events[0]
+                        if isinstance(event, dict) and "error" in event:
+                            raise llm.ModelError(event["error"]["message"])
+                        try:
+                            part = event["candidates"][0]["content"]["parts"][0]
+                            yield self.process_part(part)
+                        except KeyError:
+                            yield ""
+                        gathered.append(event)
+                        events.clear()
+        response.response_json = gathered
+
+
 @llm.hookimpl
 def register_embedding_models(register):
     register(
@@ -230,6 +295,7 @@ class GeminiEmbeddingModel(llm.EmbeddingModel):
     def embed_batch(self, items):
         headers = {
             "Content-Type": "application/json",
+            "x-goog-api-key": self.get_key(),
         }
         data = {
             "requests": [
@@ -243,7 +309,7 @@ class GeminiEmbeddingModel(llm.EmbeddingModel):
 
         with httpx.Client() as client:
             response = client.post(
-                f"https://generativelanguage.googleapis.com/v1beta/models/{self.gemini_model_id}:batchEmbedContents?key={self.get_key()}",
+                f"https://generativelanguage.googleapis.com/v1beta/models/{self.gemini_model_id}:batchEmbedContents",
                 headers=headers,
                 json=data,
                 timeout=None,

llm_gemini-0.3.dist-info/RECORD

@@ -1,7 +0,0 @@
-llm_gemini.py,sha256=DQO3ROfJSajqUYmgeuW-4_FJ1yvMoFVKb44ly20oqGw,8628
-llm_gemini-0.3.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-llm_gemini-0.3.dist-info/METADATA,sha256=ROGQUiOTfQHn1FXN3x6cgFFnNsJ75TtHTOyb_EJvFBA,4234
-llm_gemini-0.3.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
-llm_gemini-0.3.dist-info/entry_points.txt,sha256=n544bpgUPIBc5l_cnwsTxPc3gMGJHPtAyqBNp-CkMWk,26
-llm_gemini-0.3.dist-info/top_level.txt,sha256=WUQmG6_2QKbT_8W4HH93qyKl_0SUteL4Ra6_PhyNGKU,11
-llm_gemini-0.3.dist-info/RECORD,,