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

datapizza/embedders/cohere/__init__.py

@@ -0,0 +1,3 @@
+from .cohere import CohereEmbedder
+
+__all__ = ["CohereEmbedder"]

datapizza/embedders/cohere/cohere.py

@@ -0,0 +1,74 @@
+from datapizza.core.embedder import BaseEmbedder
+
+
+class CohereEmbedder(BaseEmbedder):
+    def __init__(
+        self,
+        *,
+        api_key: str,
+        model_name: str | None = None,
+        base_url: str | None = None,
+        input_type: str = "search_document",
+    ):
+        self.api_key = api_key
+        self.base_url = base_url
+        self.model_name = model_name
+
+        self.input_type = input_type
+
+        self.client = None
+        self.a_client = None
+
+    def _set_client(self):
+        import cohere
+
+        if not self.client:
+            self.client = cohere.ClientV2(base_url=self.base_url, api_key=self.api_key)
+
+    def _set_a_client(self):
+        import cohere
+
+        if not self.a_client:
+            self.a_client = cohere.AsyncClientV2(
+                base_url=self.base_url,
+                api_key=self.api_key,
+            )
+
+    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()
+
+        response = client.embed(
+            texts=texts,
+            model=model,
+            input_type=self.input_type,
+            embedding_types=["float"],
+        )
+        embeddings = response.embeddings.float
+        return embeddings[0] if isinstance(text, str) else embeddings
+
+    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()
+        response = await client.embed(
+            texts=texts,
+            model=model,
+            input_type=self.input_type,
+            embedding_types=["float"],
+        )
+        embeddings = response.embeddings.float
+        return embeddings[0] if isinstance(text, str) else embeddings

datapizza_ai_embedders_cohere-0.0.2.dist-info/METADATA

@@ -0,0 +1,12 @@
+Metadata-Version: 2.4
+Name: datapizza-ai-embedders-cohere
+Version: 0.0.2
+Summary: Cohere 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: cohere<6.0.0,>=5.14.0
+Requires-Dist: datapizza-ai-core>=0.0.0

datapizza_ai_embedders_cohere-0.0.2.dist-info/RECORD

@@ -0,0 +1,5 @@
+datapizza/embedders/cohere/__init__.py,sha256=OO44WbItRJdFA3setpBHi0eb9ngEfEElh7oGWYZzQ9I,65
+datapizza/embedders/cohere/cohere.py,sha256=94z2nX4Le7bHrSKiCRteLujA8meIEMkuq037XCsv_Vs,2133
+datapizza_ai_embedders_cohere-0.0.2.dist-info/METADATA,sha256=hH2DYq82hoVti5OSHT3CL_Bm54JSQvH4qKNlaRGjhNY,445
+datapizza_ai_embedders_cohere-0.0.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+datapizza_ai_embedders_cohere-0.0.2.dist-info/RECORD,,

datapizza_ai_embedders_cohere-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