datapizza-ai-embedders-google 0.0.2__py3-none-any.whl

datapizza/embedders/google/__init__.py

@@ -0,0 +1,3 @@
+from .google import GoogleEmbedder
+
+__all__ = ["GoogleEmbedder"]

datapizza/embedders/google/google.py

@@ -0,0 +1,62 @@
+from datapizza.core.embedder import BaseEmbedder
+
+from google import genai
+from google.genai import types
+
+
+class GoogleEmbedder(BaseEmbedder):
+    def __init__(self, *, api_key: str, model_name: str | None = None, task_type: str = "RETRIEVAL_DOCUMENT"):
+        self.api_key = api_key
+        self.model_name = model_name
+        self.task_type = task_type
+
+        self.client = None
+        self.a_client = None
+
+    def _set_client(self):
+        if not self.client:
+            client = genai.Client(api_key=self.api_key)
+            self.client = client
+
+    def _set_a_client(self):
+        if not self.a_client:
+            client = genai.Client(api_key=self.api_key)
+            self.a_client = client
+
+    def embed(
+        self, text: str | list[str], model_name: str | None = None
+    ) -> list[float] | list[list[float]]:
+        model = model_name or self.model_name
+        if not model:
+            raise ValueError("Model name is required.")
+
+        texts = [text] if isinstance(text, str) else text
+
+        client = self._get_client()
+
+        result = client.models.embed_content(
+            model=model,
+            contents=texts,
+            config=types.EmbedContentConfig(task_type=self.task_type)
+        )
+
+        res = [embedding.values for embedding in result.embeddings]
+
+        return res[0] if isinstance(text, str) else res
+
+    async def a_embed(
+        self, text: str | list[str], model_name: str | None = None
+    ) -> list[float]:
+        model = model_name or self.model_name
+        if not model:
+            raise ValueError("Model name is required.")
+
+        texts = [text] if isinstance(text, str) else text
+
+        client = self._get_a_client()
+
+        result = await client.models.embed_content_async(model=model, contents=texts, config=types.EmbedContentConfig(task_type=self.task_type))
+
+        res = [embedding.values for embedding in result.embeddings]
+
+        return res[0] if isinstance(text, str) else res

datapizza_ai_embedders_google-0.0.2.dist-info/METADATA

@@ -0,0 +1,32 @@
+Metadata-Version: 2.4
+Name: datapizza-ai-embedders-google
+Version: 0.0.2
+Summary: Google embedder for the datapizza-ai framework
+Author-email: Datapizza <datapizza@datapizza.tech>
+License: MIT
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: <4,>=3.10.0
+Requires-Dist: datapizza-ai-core>=0.0.1
+Requires-Dist: google-genai<2.0.0,>=1.3.0
+Description-Content-Type: text/markdown
+
+# Google Embedder
+
+Google Generative AI embedder implementation for datapizza-ai
+
+## Installation
+
+```bash
+pip install datapizza-ai-embedders-google
+```
+
+## Usage
+
+```python
+from datapizza.embedders.google import GoogleEmbedder
+
+embedder = GoogleEmbedder(api_key="your-google-api-key")
+embeddings = embedder.embed("Hello world", model_name="models/text-embedding-004")
+```

datapizza_ai_embedders_google-0.0.2.dist-info/RECORD

@@ -0,0 +1,5 @@
+datapizza/embedders/google/__init__.py,sha256=lFSxh2zQ5enR9g2rRKcVW3gRuUTdTVlavZkd878zx8M,65
+datapizza/embedders/google/google.py,sha256=PddSrcn02d4-nD0a9CP4T9zAY60q1YdJ7PF68czGsls,1969
+datapizza_ai_embedders_google-0.0.2.dist-info/METADATA,sha256=GLnJ1wBtyl6HW1JETrBn2wa0MphcvY6jXs8fANlUCKk,864
+datapizza_ai_embedders_google-0.0.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+datapizza_ai_embedders_google-0.0.2.dist-info/RECORD,,

datapizza_ai_embedders_google-0.0.2.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any