llm-gemini 0.1a1__py3-none-any.whl → 0.1a3__py3-none-any.whl

{llm_gemini-0.1a1.dist-info → llm_gemini-0.1a3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: llm-gemini
-Version: 0.1a1
+Version: 0.1a3
 Summary: LLM plugin to access Google's Gemini family of models
 Author: Simon Willison
 License: Apache-2.0
@@ -61,6 +61,28 @@ llm chat -m gemini-pro
 
 If you have access to the Gemini 1.5 Pro preview you can use `-m gemini-1.5-pro-latest` to work with that model.
 
+### Embeddings
+
+The plugin also adds support for the `text-embedding-004` embedding model.
+
+Run that against a single string like this:
+```bash
+llm embed -m text-embedding-004 -c 'hello world'
+```
+This returns a JSON array of 768 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
+```
+You can then run similarity searches against that collection like this:
+```bash
+llm similar readmes -c 'upload csvs to stuff' -d embed.db
+```
+
+See the [LLM embeddings documentation](https://llm.datasette.io/en/stable/embeddings/cli.html) for further details.
+
 ## Development
 
 To set up this plugin locally, first checkout the code. Then create a new virtual environment:

llm_gemini-0.1a3.dist-info/RECORD

@@ -0,0 +1,7 @@
+llm_gemini.py,sha256=uwvFPki6KrMLPqHJsfLi4eKk_4kj7HJCWSnWIjRJzgM,3961
+llm_gemini-0.1a3.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+llm_gemini-0.1a3.dist-info/METADATA,sha256=SF6uTnGVgucK_UrAghpFrXC_rQ5UBOPPtV2Tvz6Lp9E,3045
+llm_gemini-0.1a3.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+llm_gemini-0.1a3.dist-info/entry_points.txt,sha256=n544bpgUPIBc5l_cnwsTxPc3gMGJHPtAyqBNp-CkMWk,26
+llm_gemini-0.1a3.dist-info/top_level.txt,sha256=WUQmG6_2QKbT_8W4HH93qyKl_0SUteL4Ra6_PhyNGKU,11
+llm_gemini-0.1a3.dist-info/RECORD,,

llm_gemini.py

@@ -56,14 +56,17 @@ class GeminiPro(llm.Model):
             {"key": key}
         )
         gathered = []
+        body = {
+            "contents": self.build_messages(prompt, conversation),
+            "safetySettings": SAFETY_SETTINGS,
+        }
+        if prompt.system:
+            body["systemInstruction"] = {"parts": [{"text": prompt.system}]}
         with httpx.stream(
             "POST",
             url,
             timeout=None,
-            json={
-                "contents": self.build_messages(prompt, conversation),
-                "safetySettings": SAFETY_SETTINGS,
-            },
+            json=body,
         ) as http_response:
             events = ijson.sendable_list()
             coro = ijson.items_coro(events, "item")
@@ -80,3 +83,45 @@ class GeminiPro(llm.Model):
                     gathered.append(event)
                     events.clear()
         response.response_json = gathered
+
+
+@llm.hookimpl
+def register_embedding_models(register):
+    register(
+        GeminiEmbeddingModel("text-embedding-004", "text-embedding-004"),
+    )
+
+
+class GeminiEmbeddingModel(llm.EmbeddingModel):
+    needs_key = "gemini"
+    key_env_var = "LLM_GEMINI_KEY"
+    batch_size = 20
+
+    def __init__(self, model_id, gemini_model_id):
+        self.model_id = model_id
+        self.gemini_model_id = gemini_model_id
+
+    def embed_batch(self, items):
+        headers = {
+            "Content-Type": "application/json",
+        }
+        data = {
+            "requests": [
+                {
+                    "model": "models/" + self.gemini_model_id,
+                    "content": {"parts": [{"text": item}]},
+                }
+                for item in items
+            ]
+        }
+
+        with httpx.Client() as client:
+            response = client.post(
+                f"https://generativelanguage.googleapis.com/v1beta/models/{self.gemini_model_id}:batchEmbedContents?key={self.get_key()}",
+                headers=headers,
+                json=data,
+                timeout=None,
+            )
+
+        response.raise_for_status()
+        return [item["values"] for item in response.json()["embeddings"]]

llm_gemini-0.1a1.dist-info/RECORD

@@ -1,7 +0,0 @@
-llm_gemini.py,sha256=hILc9KPySth6TA2uScJBvwtCF_j8uJ1rfswCR89B8KY,2634
-llm_gemini-0.1a1.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-llm_gemini-0.1a1.dist-info/METADATA,sha256=vJwoQbQzRkn_vpY-h_o_9m7Q-Rlo3c5bQWQFSLTrxBY,2262
-llm_gemini-0.1a1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-llm_gemini-0.1a1.dist-info/entry_points.txt,sha256=n544bpgUPIBc5l_cnwsTxPc3gMGJHPtAyqBNp-CkMWk,26
-llm_gemini-0.1a1.dist-info/top_level.txt,sha256=WUQmG6_2QKbT_8W4HH93qyKl_0SUteL4Ra6_PhyNGKU,11
-llm_gemini-0.1a1.dist-info/RECORD,,