llm-gemini 0.13.1__py3-none-any.whl → 0.14.1__py3-none-any.whl

{llm_gemini-0.13.1.dist-info → llm_gemini-0.14.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: llm-gemini
-Version: 0.13.1
+Version: 0.14.1
 Summary: LLM plugin to access Google's Gemini family of models
 Author: Simon Willison
 License: Apache-2.0
@@ -17,6 +17,7 @@ Requires-Dist: ijson
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-recording; extra == "test"
+Requires-Dist: pytest-asyncio; extra == "test"
 Requires-Dist: nest-asyncio; extra == "test"
 
 # llm-gemini
@@ -145,7 +146,7 @@ llm chat -m gemini-1.5-pro-latest
 
 ## Embeddings
 
-The plugin also adds support for the `text-embedding-004` embedding model.
+The plugin also adds support for the `gemini-embedding-exp-03-07` and `text-embedding-004` embedding models.
 
 Run that against a single string like this:
 ```bash
@@ -153,10 +154,20 @@ llm embed -m text-embedding-004 -c 'hello world'
 ```
 This returns a JSON array of 768 numbers.
 
+The `gemini-embedding-exp-03-07` model is larger, returning 3072 numbers. You can also use variants of it that are truncated down to smaller sizes:
+
+- `gemini-embedding-exp-03-07` - 3072 numbers
+- `gemini-embedding-exp-03-07-2048` - 2048 numbers
+- `gemini-embedding-exp-03-07-1024` - 1024 numbers
+- `gemini-embedding-exp-03-07-512` - 512 numbers
+- `gemini-embedding-exp-03-07-256` - 256 numbers
+- `gemini-embedding-exp-03-07-128` - 128 numbers
+
 This command will embed every `README.md` file in child directories of the current directory and store the results in a SQLite database called `embed.db` in a collection called `readmes`:
 
 ```bash
-llm embed-multi readmes --files . '*/README.md' -d embed.db -m text-embedding-004
+llm embed-multi readmes -d embed.db -m gemini-embedding-exp-03-07-128 \
+  --files . '*/README.md'
 ```
 You can then run similarity searches against that collection like this:
 ```bash

llm_gemini-0.14.1.dist-info/RECORD

@@ -0,0 +1,7 @@
+llm_gemini.py,sha256=eddn8U4kN2g9DOrotcKbyzNxGonR7XLiK7i_m7dgZ84,15296
+llm_gemini-0.14.1.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+llm_gemini-0.14.1.dist-info/METADATA,sha256=jH04kG_VQYRVqjyVF2ZWHAWBqyxf-giSyX1NRO8ntss,7558
+llm_gemini-0.14.1.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
+llm_gemini-0.14.1.dist-info/entry_points.txt,sha256=n544bpgUPIBc5l_cnwsTxPc3gMGJHPtAyqBNp-CkMWk,26
+llm_gemini-0.14.1.dist-info/top_level.txt,sha256=WUQmG6_2QKbT_8W4HH93qyKl_0SUteL4Ra6_PhyNGKU,11
+llm_gemini-0.14.1.dist-info/RECORD,,

llm_gemini.py

@@ -88,18 +88,24 @@ def resolve_type(attachment):
     return mime_type
 
 
-def cleanup_schema(schema):
+def cleanup_schema(schema, in_properties=False):
     "Gemini supports only a subset of JSON schema"
     keys_to_remove = ("$schema", "additionalProperties", "title")
-    # Recursively remove them
+
     if isinstance(schema, dict):
-        for key in keys_to_remove:
-            schema.pop(key, None)
-        for value in schema.values():
-            cleanup_schema(value)
+        # Only remove keys if we're not inside a 'properties' block.
+        if not in_properties:
+            for key in keys_to_remove:
+                schema.pop(key, None)
+        for key, value in list(schema.items()):
+            # If the key is 'properties', set the flag for its value.
+            if key == "properties" and isinstance(value, dict):
+                cleanup_schema(value, in_properties=True)
+            else:
+                cleanup_schema(value, in_properties=in_properties)
     elif isinstance(schema, list):
-        for value in schema:
-            cleanup_schema(value)
+        for item in schema:
+            cleanup_schema(item, in_properties=in_properties)
     return schema
 
 
@@ -378,9 +384,19 @@ class AsyncGeminiPro(_SharedGemini, llm.AsyncKeyModel):
 
 @llm.hookimpl
 def register_embedding_models(register):
+    register(GeminiEmbeddingModel("text-embedding-004", "text-embedding-004"))
+    # gemini-embedding-exp-03-07 in different truncation sizes
     register(
-        GeminiEmbeddingModel("text-embedding-004", "text-embedding-004"),
+        GeminiEmbeddingModel(
+            "gemini-embedding-exp-03-07", "gemini-embedding-exp-03-07"
+        ),
     )
+    for i in (128, 256, 512, 1024, 2048):
+        register(
+            GeminiEmbeddingModel(
+                f"gemini-embedding-exp-03-07-{i}", f"gemini-embedding-exp-03-07", i
+            ),
+        )
 
 
 class GeminiEmbeddingModel(llm.EmbeddingModel):
@@ -388,9 +404,10 @@ class GeminiEmbeddingModel(llm.EmbeddingModel):
     key_env_var = "LLM_GEMINI_KEY"
     batch_size = 20
 
-    def __init__(self, model_id, gemini_model_id):
+    def __init__(self, model_id, gemini_model_id, truncate=None):
         self.model_id = model_id
         self.gemini_model_id = gemini_model_id
+        self.truncate = truncate
 
     def embed_batch(self, items):
         headers = {
@@ -416,4 +433,7 @@ class GeminiEmbeddingModel(llm.EmbeddingModel):
             )
 
         response.raise_for_status()
-        return [item["values"] for item in response.json()["embeddings"]]
+        values = [item["values"] for item in response.json()["embeddings"]]
+        if self.truncate:
+            values = [value[: self.truncate] for value in values]
+        return values

llm_gemini-0.13.1.dist-info/RECORD

@@ -1,7 +0,0 @@
-llm_gemini.py,sha256=JWaRtT8vJzdxrRbCAgoNrKqes2df_T4gqqeDmiw0oRI,14370
-llm_gemini-0.13.1.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-llm_gemini-0.13.1.dist-info/METADATA,sha256=S87OFYnm9K0EVvJIRrJoZT-J0gxhPA4BrnIGkVD8q8c,7016
-llm_gemini-0.13.1.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
-llm_gemini-0.13.1.dist-info/entry_points.txt,sha256=n544bpgUPIBc5l_cnwsTxPc3gMGJHPtAyqBNp-CkMWk,26
-llm_gemini-0.13.1.dist-info/top_level.txt,sha256=WUQmG6_2QKbT_8W4HH93qyKl_0SUteL4Ra6_PhyNGKU,11
-llm_gemini-0.13.1.dist-info/RECORD,,