embedkit 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

embedkit/base.py

@@ -2,7 +2,7 @@
 """Base classes for EmbedKit."""
 
 from abc import ABC, abstractmethod
-from typing import Union, List
+from typing import Union, List, Optional
 from pathlib import Path
 import numpy as np
 from dataclasses import dataclass
@@ -14,6 +14,7 @@ class EmbeddingResult:
     model_name: str
     model_provider: str
     input_type: str
+    source_images_b64: Optional[List[str]] = None
 
     @property
     def shape(self) -> tuple:

embedkit/providers/cohere.py

@@ -6,7 +6,7 @@ from pathlib import Path
 import numpy as np
 from enum import Enum
 
-from ..utils import pdf_to_images
+from ..utils import pdf_to_images, image_to_base64
 from ..base import EmbeddingProvider, EmbeddingError, EmbeddingResult
 
 
@@ -84,38 +84,11 @@ class CohereProvider(EmbeddingProvider):
             images = [images]
 
         try:
-            import base64
-
             b64_images = []
             for image in images:
-                if isinstance(image, (Path, str)):
-                    try:
-                        base64_only = base64.b64encode(Path(image).read_bytes()).decode(
-                            "utf-8"
-                        )
-                    except Exception as e:
-                        raise EmbeddingError(
-                            f"Failed to read image {image}: {e}"
-                        ) from e
-
-                    if isinstance(image, Path):
-                        image = str(image)
-
-                    if image.lower().endswith(".png"):
-                        content_type = "image/png"
-                    elif image.lower().endswith((".jpg", ".jpeg")):
-                        content_type = "image/jpeg"
-                    elif image.lower().endswith(".gif"):
-                        content_type = "image/gif"
-                    else:
-                        raise EmbeddingError(
-                            f"Unsupported image format for {image}; expected .png, .jpg, .jpeg, or .gif"
-                        )
-                    base64_image = f"data:{content_type};base64,{base64_only}"
-                else:
-                    raise EmbeddingError(f"Unsupported image type: {type(image)}")
-
-                b64_images.append(base64_image)
+                b64_image = image_to_base64(image)
+
+            b64_images.append(b64_image)
 
             response = client.embed(
                 model=self.model_name,
@@ -129,6 +102,7 @@ class CohereProvider(EmbeddingProvider):
                 model_name=self.model_name,
                 model_provider=self.provider_name,
                 input_type=input_type,
+                source_images_b64=b64_images,
             )
 
         except Exception as e:

embedkit/providers/colpali.py

@@ -8,7 +8,7 @@ import numpy as np
 import torch
 from PIL import Image
 
-from ..utils import pdf_to_images
+from ..utils import pdf_to_images, image_to_base64
 from ..base import EmbeddingProvider, EmbeddingError, EmbeddingResult
 
 logger = logging.getLogger(__name__)
@@ -92,6 +92,7 @@ class ColPaliProvider(EmbeddingProvider):
 
         try:
             pil_images = []
+            b64_images = []
             for img_path in images:
                 if not img_path.exists():
                     raise EmbeddingError(f"Image not found: {img_path}")
@@ -99,8 +100,14 @@ class ColPaliProvider(EmbeddingProvider):
                 with Image.open(img_path) as img:
                     pil_images.append(img.convert("RGB"))
 
+                for image in images:
+                    b64_image = image_to_base64(image)
+
+                b64_images.append(b64_image)
+
             processed = self._processor.process_images(pil_images).to(self.device)
 
+
             with torch.no_grad():
                 embeddings = self._model(**processed)
 
@@ -109,6 +116,7 @@ class ColPaliProvider(EmbeddingProvider):
                 model_name=self.model_name,
                 model_provider=self.provider_name,
                 input_type="image",
+                source_images_b64=b64_images,
             )
 
         except Exception as e:

embedkit/utils.py

@@ -1,6 +1,7 @@
 from pdf2image import convert_from_path
 from pathlib import Path
 from .config import get_temp_dir
+from typing import Union
 
 
 def pdf_to_images(pdf_path: Path) -> list[Path]:
@@ -19,3 +20,33 @@ def pdf_to_images(pdf_path: Path) -> list[Path]:
         image.save(output_path)
         image_paths.append(output_path)
     return image_paths
+
+
+def image_to_base64(image_path: Union[str, Path]):
+    import base64
+
+    try:
+        base64_only = base64.b64encode(Path(image_path).read_bytes()).decode(
+            "utf-8"
+        )
+    except Exception as e:
+        raise ValueError(
+            f"Failed to read image {image_path}: {e}"
+        ) from e
+
+    if isinstance(image_path, Path):
+        image_path_str = str(image_path)
+
+    if image_path_str.lower().endswith(".png"):
+        content_type = "image/png"
+    elif image_path_str.lower().endswith((".jpg", ".jpeg")):
+        content_type = "image/jpeg"
+    elif image_path_str.lower().endswith(".gif"):
+        content_type = "image/gif"
+    else:
+        raise ValueError(
+            f"Unsupported image format for {image_path}; expected .png, .jpg, .jpeg, or .gif"
+        )
+    base64_image = f"data:{content_type};base64,{base64_only}"
+
+    return base64_image

{embedkit-0.1.0.dist-info → embedkit-0.1.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: embedkit
-Version: 0.1.0
+Version: 0.1.1
 Summary: A simple toolkit for generating vector embeddings across multiple providers and models
 Author-email: JP Hwang <me@jphwang.com>
 License: MIT

embedkit-0.1.1.dist-info/RECORD

@@ -0,0 +1,12 @@
+embedkit/__init__.py,sha256=vm_dF7i_EGQsNEgBn7WPq-Vbo1xTnqV2devUvY18Z5E,3862
+embedkit/base.py,sha256=wB55ksH-LT5xkEj39FUE3mbJK14DSGyK2dUpwmMMEaE,1293
+embedkit/config.py,sha256=EVGODSKxQAr46bU8dyORFunsfRuj6dnvtSqa4MxUZCo,138
+embedkit/models.py,sha256=EBIYkyZeIhGaOPL-9bslHHdLaZ7qzOYLd0qxVZ7VX7w,226
+embedkit/utils.py,sha256=LlwUq2KIiAl6J8WpVAxiAz5V6Gj1m1ItFjBeCHdBmy8,1616
+embedkit/providers/__init__.py,sha256=HaS-HNQabvhn9xLNZCq3VUqPCb7rGG4pvgvpKP4AXcw,201
+embedkit/providers/cohere.py,sha256=rDyZI2UWrwcX-AgJoDotLAmMk2mBmwC-NWiyyp5msek,3589
+embedkit/providers/colpali.py,sha256=RDaGyYUCWv_QgkG7jVUgEgwJy1ReKaOY7iM_nxtvR60,4127
+embedkit-0.1.1.dist-info/METADATA,sha256=dCC3gxaZFrDz0S2Ab9A_Sd2cGlyRQNVQp8Z7VBfyjSo,1893
+embedkit-0.1.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+embedkit-0.1.1.dist-info/licenses/LICENSE,sha256=-g2Rad7b3rb2oVwOTwfMOIpscHT1zuaJoguamLRCBJs,1072
+embedkit-0.1.1.dist-info/RECORD,,

embedkit-0.1.0.dist-info/RECORD

@@ -1,12 +0,0 @@
-embedkit/__init__.py,sha256=vm_dF7i_EGQsNEgBn7WPq-Vbo1xTnqV2devUvY18Z5E,3862
-embedkit/base.py,sha256=ZwCeDnJXVsVVT5l7ybpP5wG2ZU9e19XgV3c9OJp9z2o,1233
-embedkit/config.py,sha256=EVGODSKxQAr46bU8dyORFunsfRuj6dnvtSqa4MxUZCo,138
-embedkit/models.py,sha256=EBIYkyZeIhGaOPL-9bslHHdLaZ7qzOYLd0qxVZ7VX7w,226
-embedkit/utils.py,sha256=TyFyDk6tMx-PaVotixSdJDx8U3JgrPi9nV2j-rW-clw,705
-embedkit/providers/__init__.py,sha256=HaS-HNQabvhn9xLNZCq3VUqPCb7rGG4pvgvpKP4AXcw,201
-embedkit/providers/cohere.py,sha256=u6zoAjXKkjaVfTZk1VgjwRqtQ7Bea1odlVBKWomB_1A,4737
-embedkit/providers/colpali.py,sha256=20YAEeTvkNoexax-KhU7lWjJBdWRHPzE4Zf-6XpP3v0,3896
-embedkit-0.1.0.dist-info/METADATA,sha256=18DAz2h--FOgMSO3VNgm9ZXENSXK9IsVkEYm-xb2a3c,1893
-embedkit-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-embedkit-0.1.0.dist-info/licenses/LICENSE,sha256=-g2Rad7b3rb2oVwOTwfMOIpscHT1zuaJoguamLRCBJs,1072
-embedkit-0.1.0.dist-info/RECORD,,