embedding-flow 0.1.0__tar.gz

embedding_flow-0.1.0/LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2025
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+

embedding_flow-0.1.0/MANIFEST.in

@@ -0,0 +1,7 @@
+include README.md
+include LICENSE
+include requirements.txt
+recursive-include contracts *.py
+recursive-include transform *.py
+recursive-include load *.py
+

embedding_flow-0.1.0/PKG-INFO

@@ -0,0 +1,22 @@
+Metadata-Version: 2.4
+Name: embedding-flow
+Version: 0.1.0
+Summary: Pipeline to transform text chunks into embeddings and load to Qdrant
+Author: facuvega
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+License-File: LICENSE
+Requires-Dist: pandas>=2.0.0
+Requires-Dist: pyarrow>=12.0.0
+Requires-Dist: sentence-transformers>=2.2.0
+Requires-Dist: torch>=2.0.0
+Requires-Dist: qdrant-client>=1.7.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
+Dynamic: license-file

embedding_flow-0.1.0/README.md

@@ -0,0 +1,40 @@
+# embedding-flow
+
+Pipeline for transforming text chunks into 768-dimensional embeddings and loading to Qdrant.
+
+## Installation
+
+```bash
+pip install embedding-flow
+```
+
+## Usage
+
+```python
+from transform.transform import transform_embedding
+from load.load import load_embedding
+
+# Transform
+transformer = transform_embedding()
+output_path = transformer.transform_data("chunks.parquet")
+
+# Load to Qdrant
+loader = load_embedding()
+loader.load_data(output_path)
+```
+
+## Environment Variables
+
+```bash
+QDRANT_URL=http://localhost:6333
+QDRANT_COLLECTION=embeddings_collection
+VECTOR_SIZE=768
+```
+
+## Development
+
+```bash
+pip install -e ".[dev]"
+pytest tests/
+```
+

embedding_flow-0.1.0/contracts/contracts.py

@@ -0,0 +1,14 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+class transform_data(ABC):
+    @abstractmethod
+    def transform_data(self, url: str) -> Optional[str]:
+        """Transforma datos y retorna la ruta del archivo procesado, o None si falla"""
+        pass
+
+class load_data(ABC):
+    @abstractmethod
+    def load_data(self, url: str) -> bool:
+        """Carga datos y retorna True si fue exitoso, False si falló"""
+        pass

embedding_flow-0.1.0/embedding_flow.egg-info/PKG-INFO

@@ -0,0 +1,22 @@
+Metadata-Version: 2.4
+Name: embedding-flow
+Version: 0.1.0
+Summary: Pipeline to transform text chunks into embeddings and load to Qdrant
+Author: facuvega
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+License-File: LICENSE
+Requires-Dist: pandas>=2.0.0
+Requires-Dist: pyarrow>=12.0.0
+Requires-Dist: sentence-transformers>=2.2.0
+Requires-Dist: torch>=2.0.0
+Requires-Dist: qdrant-client>=1.7.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
+Dynamic: license-file

embedding_flow-0.1.0/embedding_flow.egg-info/SOURCES.txt

@@ -0,0 +1,19 @@
+LICENSE
+MANIFEST.in
+README.md
+pyproject.toml
+requirements.txt
+setup.py
+contracts/__init__.py
+contracts/contracts.py
+embedding_flow.egg-info/PKG-INFO
+embedding_flow.egg-info/SOURCES.txt
+embedding_flow.egg-info/dependency_links.txt
+embedding_flow.egg-info/requires.txt
+embedding_flow.egg-info/top_level.txt
+load/__init__.py
+load/load.py
+tests/test_load.py
+tests/test_transform.py
+transform/__init__.py
+transform/transform.py

embedding_flow-0.1.0/embedding_flow.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

embedding_flow-0.1.0/embedding_flow.egg-info/requires.txt

@@ -0,0 +1,9 @@
+pandas>=2.0.0
+pyarrow>=12.0.0
+sentence-transformers>=2.2.0
+torch>=2.0.0
+qdrant-client>=1.7.0
+
+[dev]
+pytest>=7.0.0
+pytest-cov>=4.0.0

embedding_flow-0.1.0/embedding_flow.egg-info/top_level.txt

@@ -0,0 +1,3 @@
+contracts
+load
+transform

embedding_flow-0.1.0/load/load.py

@@ -0,0 +1,101 @@
+from contracts.contracts import load_data
+from qdrant_client import QdrantClient
+from qdrant_client.models import Distance, VectorParams, PointStruct
+import pandas as pd
+import logging
+import os
+from typing import List
+import uuid
+
+# Configurar logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+class load_embedding(load_data):
+    def __init__(self):
+        """Inicializa el cliente de Qdrant"""
+        self.qdrant_url = os.getenv("QDRANT_URL", "http://localhost:6333")
+        self.qdrant_api_key = os.getenv("QDRANT_API_KEY", None)
+        self.collection_name = os.getenv("QDRANT_COLLECTION", "embeddings_collection")
+        self.vector_size = int(os.getenv("VECTOR_SIZE", "768"))
+        
+        # Inicializar cliente
+        self.client = QdrantClient(
+            url=self.qdrant_url,
+            api_key=self.qdrant_api_key
+        )
+        
+        # Crear colección si no existe
+        self._ensure_collection_exists()
+    
+    def _ensure_collection_exists(self):
+        """Crea la colección en Qdrant si no existe"""
+        try:
+            collections = self.client.get_collections().collections
+            collection_exists = any(col.name == self.collection_name for col in collections)
+            
+            if not collection_exists:
+                self.client.create_collection(
+                    collection_name=self.collection_name,
+                    vectors_config=VectorParams(
+                        size=self.vector_size,
+                        distance=Distance.COSINE
+                    )
+                )
+                logger.info(f"✅ Colección '{self.collection_name}' creada en Qdrant")
+            else:
+                logger.info(f"ℹ️ Colección '{self.collection_name}' ya existe")
+        except Exception as e:
+            logger.error(f"❌ Error al verificar/crear colección: {e}", exc_info=True)
+            raise
+    
+    def load_data(self, parquet_path: str) -> bool:
+        """
+        Carga los embeddings desde un parquet a Qdrant
+        
+        Args:
+            parquet_path: Ruta al archivo parquet con embeddings
+            
+        Returns:
+            True si la carga fue exitosa, False si falló
+        """
+        try:
+            # Leer parquet
+            df = pd.read_parquet(parquet_path)
+            
+            if "embedding" not in df.columns:
+                raise ValueError(f"El parquet {parquet_path} no contiene columna 'embedding'")
+            
+            # Preparar puntos para Qdrant
+            points: List[PointStruct] = []
+            
+            for idx, row in df.iterrows():
+                # Generar ID único
+                point_id = str(uuid.uuid4())
+                
+                # Preparar payload (todos los campos excepto embedding)
+                payload = {col: row[col] for col in df.columns if col != "embedding"}
+                
+                # Crear punto
+                point = PointStruct(
+                    id=point_id,
+                    vector=row["embedding"],
+                    payload=payload
+                )
+                points.append(point)
+            
+            # Insertar en Qdrant (en batch)
+            self.client.upsert(
+                collection_name=self.collection_name,
+                points=points
+            )
+            
+            logger.info(f"✅ {len(points)} embeddings cargados a Qdrant desde {parquet_path}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"❌ Error al cargar embeddings a Qdrant desde {parquet_path}: {e}", exc_info=True)
+            return False

embedding_flow-0.1.0/pyproject.toml

@@ -0,0 +1,36 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "embedding-flow"
+version = "0.1.0"
+description = "Pipeline to transform text chunks into embeddings and load to Qdrant"
+authors = [{name = "facuvega"}]
+requires-python = ">=3.10"
+dependencies = [
+    "pandas>=2.0.0",
+    "pyarrow>=12.0.0",
+    "sentence-transformers>=2.2.0",
+    "torch>=2.0.0",
+    "qdrant-client>=1.7.0",
+]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0.0",
+    "pytest-cov>=4.0.0",
+]
+
+[tool.setuptools.packages.find]
+include = ["contracts*", "transform*", "load*"]
+exclude = ["tests*", "venv*"]
+

embedding_flow-0.1.0/requirements.txt

@@ -0,0 +1,17 @@
+# Core dependencies
+pandas>=2.0.0
+pyarrow>=12.0.0
+
+# ML & Embeddings
+sentence-transformers>=2.2.0
+torch>=2.0.0
+
+# Vector Database
+qdrant-client>=1.7.0
+
+# Airflow (si se necesita localmente, sino está en el servidor)
+# apache-airflow>=2.7.0
+
+# Utilities
+python-dotenv>=1.0.0
+

embedding_flow-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

embedding_flow-0.1.0/setup.py

@@ -0,0 +1,6 @@
+from setuptools import setup, find_packages
+
+setup(
+    packages=find_packages(exclude=["tests", "venv"]),
+)
+

embedding_flow-0.1.0/tests/test_load.py

@@ -0,0 +1,63 @@
+import pytest
+import pandas as pd
+import tempfile
+import numpy as np
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+from load.load import load_embedding
+
+
+@patch('load.load.QdrantClient')
+def test_load_with_embeddings(mock_qdrant_client):
+    """Test que load carga correctamente embeddings de 768 dims"""
+    # Mock del cliente
+    mock_client = MagicMock()
+    mock_qdrant_client.return_value = mock_client
+    mock_client.get_collections.return_value.collections = []
+    
+    # Crear parquet con embeddings de 768 dims
+    with tempfile.NamedTemporaryFile(suffix='.parquet', delete=False) as tmp:
+        df = pd.DataFrame({
+            'text': ['test 1', 'test 2'],
+            'embedding': [np.random.rand(768).tolist(), np.random.rand(768).tolist()]
+        })
+        df.to_parquet(tmp.name, index=False)
+        tmp_path = tmp.name
+    
+    try:
+        loader = load_embedding()
+        success = loader.load_data(tmp_path)
+        
+        assert success is True
+        mock_client.upsert.assert_called_once()
+        
+        # Verificar que se llamó con 2 puntos
+        call_args = mock_client.upsert.call_args
+        points = call_args.kwargs['points']
+        assert len(points) == 2
+        
+    finally:
+        Path(tmp_path).unlink(missing_ok=True)
+
+
+@patch('load.load.QdrantClient')
+def test_load_without_embeddings(mock_qdrant_client):
+    """Test que load falla sin columna 'embedding'"""
+    mock_client = MagicMock()
+    mock_qdrant_client.return_value = mock_client
+    mock_client.get_collections.return_value.collections = []
+    
+    with tempfile.NamedTemporaryFile(suffix='.parquet', delete=False) as tmp:
+        df = pd.DataFrame({
+            'text': ['test 1', 'test 2']
+        })
+        df.to_parquet(tmp.name, index=False)
+        tmp_path = tmp.name
+    
+    try:
+        loader = load_embedding()
+        success = loader.load_data(tmp_path)
+        assert success is False
+    finally:
+        Path(tmp_path).unlink(missing_ok=True)
+

embedding_flow-0.1.0/tests/test_transform.py

@@ -0,0 +1,54 @@
+import pytest
+import pandas as pd
+import tempfile
+from pathlib import Path
+from transform.transform import transform_embedding
+
+
+def test_transform_creates_embeddings():
+    """Test que transform genera embeddings de 768 dimensiones"""
+    # Crear parquet temporal con texto
+    with tempfile.NamedTemporaryFile(suffix='.parquet', delete=False) as tmp:
+        df = pd.DataFrame({
+            'text': ['test text 1', 'test text 2', 'test text 3']
+        })
+        df.to_parquet(tmp.name, index=False)
+        tmp_path = tmp.name
+    
+    try:
+        # Transformar
+        transformer = transform_embedding()
+        output_path = transformer.transform_data(tmp_path)
+        
+        # Verificar
+        assert output_path is not None
+        assert Path(output_path).exists()
+        
+        # Verificar dimensiones
+        result_df = pd.read_parquet(output_path)
+        assert 'embedding' in result_df.columns
+        assert len(result_df['embedding'][0]) == 768
+        
+    finally:
+        # Cleanup
+        Path(tmp_path).unlink(missing_ok=True)
+        if output_path:
+            Path(output_path).unlink(missing_ok=True)
+
+
+def test_transform_missing_text_column():
+    """Test que transform falla sin columna 'text'"""
+    with tempfile.NamedTemporaryFile(suffix='.parquet', delete=False) as tmp:
+        df = pd.DataFrame({
+            'wrong_column': ['data 1', 'data 2']
+        })
+        df.to_parquet(tmp.name, index=False)
+        tmp_path = tmp.name
+    
+    try:
+        transformer = transform_embedding()
+        output_path = transformer.transform_data(tmp_path)
+        assert output_path is None
+    finally:
+        Path(tmp_path).unlink(missing_ok=True)
+

embedding_flow-0.1.0/transform/transform.py

@@ -0,0 +1,42 @@
+from contracts.contracts import transform_data
+from pathlib import Path
+import pandas as pd
+from sentence_transformers import SentenceTransformer
+import logging
+from typing import Optional
+
+# Configurar logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+class transform_embedding(transform_data):
+    def transform_data(self, parquet_path: str) -> Optional[str]:
+        try:
+            # Cargar datos del parquet
+            df = pd.read_parquet(parquet_path)
+            if "text" not in df.columns:
+                raise ValueError("El parquet no contiene una columna 'text' para generar embeddings.")
+
+            # Inicializar modelo con embeddings de 768 dimensiones
+            model = SentenceTransformer("all-mpnet-base-v2")
+
+            # Generar embeddings
+            embeddings = model.encode(df["text"].tolist(), show_progress_bar=True)
+            df["embedding"] = embeddings.tolist()  # guardar como lista
+
+            # Guardar parquet procesado
+            output_dir = Path("datos_embeddings")
+            output_dir.mkdir(parents=True, exist_ok=True)
+
+            output_file = output_dir / Path(parquet_path).name
+            df.to_parquet(output_file, index=False)
+
+            logger.info(f"✅ Embeddings (768 dim) generados y guardados en: {output_file}")
+            return str(output_file)
+
+        except Exception as e:
+            logger.error(f"❌ Error al transformar en embeddings {parquet_path}: {e}", exc_info=True)
+            return None