First commmit

packages/embedder/.python-version

@@ -0,0 +1 @@
+3.10

packages/embedder/pyproject.toml

@@ -0,0 +1,16 @@
+[project]
+name = "embedder"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+authors = [
+    { name = "Anibal Angulo", email = "a8065384@banorte.com" }
+]
+requires-python = ">=3.12"
+dependencies = [
+    "google-cloud-aiplatform>=1.106.0",
+]
+
+[build-system]
+requires = ["uv_build>=0.8.3,<0.9.0"]
+build-backend = "uv_build"

packages/embedder/src/embedder/base.py

@@ -0,0 +1,79 @@
+from abc import ABC, abstractmethod
+from typing import List
+
+import numpy as np
+
+
+class BaseEmbedder(ABC):
+    """Base class for all embedding models."""
+
+    @abstractmethod
+    def generate_embedding(self, text: str) -> List[float]:
+        """
+        Generate embeddings for text.
+
+        Args:
+            text: Single text string or list of texts
+
+        Returns:
+            Single embedding vector or list of embedding vectors
+        """
+        pass
+
+    @abstractmethod
+    def generate_embeddings_batch(self, texts: List[str]) -> List[List[float]]:
+        """
+        Generate embeddings for a batch of texts.
+
+        Args:
+            texts: List of text strings
+
+        Returns:
+            List of embedding vectors
+        """
+        pass
+
+    def preprocess_text(
+        self,
+        text: str,
+        *,
+        into_lowercase: bool = False,
+        normalize_whitespace: bool = True,
+        remove_punctuation: bool = False,
+    ) -> str:
+        """Preprocess text before embedding."""
+        # Basic preprocessing
+        text = text.strip()
+
+        if into_lowercase:
+            text = text.lower()
+
+        if normalize_whitespace:
+            text = " ".join(text.split())
+
+        if remove_punctuation:
+            import string
+
+            text = text.translate(str.maketrans("", "", string.punctuation))
+
+        return text
+
+    def normalize_embedding(self, embedding: List[float]) -> List[float]:
+        """Normalize embedding vector to unit length."""
+        norm = np.linalg.norm(embedding)
+        if norm > 0:
+            return (np.array(embedding) / norm).tolist()
+        return embedding
+
+    @abstractmethod
+    async def async_generate_embedding(self, text: str) -> List[float]:
+        """
+        Generate embeddings for text.
+
+        Args:
+            text: Single text string or list of texts
+
+        Returns:
+            Single embedding vector or list of embedding vectors
+        """
+        pass

packages/embedder/src/embedder/vertex_ai.py

@@ -0,0 +1,77 @@
+import logging
+import time
+from typing import List
+
+from google import genai
+from google.genai import types
+from tenacity import retry, stop_after_attempt, wait_exponential
+
+from .base import BaseEmbedder
+
+logger = logging.getLogger(__name__)
+
+
+class VertexAIEmbedder(BaseEmbedder):
+    """Embedder using Vertex AI text embedding models."""
+
+    def __init__(
+            self, model_name: str, project: str, location: str, task: str = "RETRIEVAL_DOCUMENT"
+    ) -> None:
+        self.model_name = model_name
+        self.client = genai.Client(
+            vertexai=True,
+            project=project,
+            location=location,
+        )
+        self.task = task
+
+    # @retry(
+    #     stop=stop_after_attempt(5), wait=wait_exponential(multiplier=1, min=4, max=30)
+    # )
+    def generate_embedding(self, text: str) -> List[float]:
+        preprocessed_text = self.preprocess_text(text)
+        result = self.client.models.embed_content(
+            model=self.model_name, contents=preprocessed_text, config=types.EmbedContentConfig(task_type=self.task)
+        )
+        return result.embeddings[0].values
+
+    # @retry(
+    #     stop=stop_after_attempt(5), wait=wait_exponential(multiplier=1, min=4, max=30)
+    # )
+    def generate_embeddings_batch(
+        self, texts: List[str], batch_size: int = 10
+    ) -> List[List[float]]:
+        """Generate embeddings for a batch of texts."""
+        if not texts:
+            return []
+
+        # Preprocess texts
+        preprocessed_texts = [self.preprocess_text(text) for text in texts]
+
+        # Process in batches if necessary
+        all_embeddings = []
+
+        for i in range(0, len(preprocessed_texts), batch_size):
+            batch = preprocessed_texts[i : i + batch_size]
+
+            # Generate embeddings for batch
+            result = self.client.models.embed_content(
+                model=self.model_name, contents=batch, config=types.EmbedContentConfig(task_type=self.task)
+            )
+
+            # Extract values
+            batch_embeddings = [emb.values for emb in result.embeddings]
+            all_embeddings.extend(batch_embeddings)
+
+            # Rate limiting
+            if i + batch_size < len(preprocessed_texts):
+                time.sleep(0.1)  # Small delay between batches
+
+        return all_embeddings
+
+    async def async_generate_embedding(self, text: str) -> List[float]:
+        preprocessed_text = self.preprocess_text(text)
+        result = await self.client.aio.models.embed_content(
+            model=self.model_name, contents=preprocessed_text, config=types.EmbedContentConfig(task_type=self.task)
+        )
+        return result.embeddings[0].values