Feature: improve relationship builders for better async and reduced memory utilization

docs/howtos/applications/cost.ipynb

@@ -26,7 +26,19 @@
    "cell_type": "code",
    "metadata": {},
    "outputs": [],
-   "source": "from langchain_openai.chat_models import ChatOpenAI\nfrom langchain_core.prompt_values import StringPromptValue\n# lets import a parser for OpenAI\nfrom ragas.cost import get_token_usage_for_openai\n\ngpt4o = ChatOpenAI(model=\"gpt-4o\")\np = StringPromptValue(text=\"hai there\")\nllm_result = gpt4o.generate_prompt([p])\n\nget_token_usage_for_openai(llm_result)"
+   "source": [
+    "from langchain_openai.chat_models import ChatOpenAI\n",
+    "from langchain_core.prompt_values import StringPromptValue\n",
+    "\n",
+    "# lets import a parser for OpenAI\n",
+    "from ragas.cost import get_token_usage_for_openai\n",
+    "\n",
+    "gpt4o = ChatOpenAI(model=\"gpt-4o\")\n",
+    "p = StringPromptValue(text=\"hai there\")\n",
+    "llm_result = gpt4o.generate_prompt([p])\n",
+    "\n",
+    "get_token_usage_for_openai(llm_result)"
+   ]
   },
   {
    "cell_type": "markdown",

docs/howtos/customizations/metrics/cost.ipynb

@@ -37,7 +37,19 @@
    "cell_type": "code",
    "metadata": {},
    "outputs": [],
-   "source": "from langchain_openai.chat_models import ChatOpenAI\nfrom langchain_core.prompt_values import StringPromptValue\n# lets import a parser for OpenAI\nfrom ragas.cost import get_token_usage_for_openai\n\ngpt4o = ChatOpenAI(model=\"gpt-4o\")\np = StringPromptValue(text=\"hai there\")\nllm_result = gpt4o.generate_prompt([p])\n\nget_token_usage_for_openai(llm_result)"
+   "source": [
+    "from langchain_openai.chat_models import ChatOpenAI\n",
+    "from langchain_core.prompt_values import StringPromptValue\n",
+    "\n",
+    "# lets import a parser for OpenAI\n",
+    "from ragas.cost import get_token_usage_for_openai\n",
+    "\n",
+    "gpt4o = ChatOpenAI(model=\"gpt-4o\")\n",
+    "p = StringPromptValue(text=\"hai there\")\n",
+    "llm_result = gpt4o.generate_prompt([p])\n",
+    "\n",
+    "get_token_usage_for_openai(llm_result)"
+   ]
   },
   {
    "cell_type": "markdown",

docs/howtos/integrations/helicone.ipynb

@@ -47,7 +47,29 @@
    "cell_type": "code",
    "metadata": {},
    "outputs": [],
-   "source": "import os\nfrom datasets import Dataset\nfrom ragas import evaluate\nfrom ragas.metrics import faithfulness, answer_relevancy, context_precision\nfrom ragas.integrations.helicone import helicone_config  # import helicone_config\n\n\n# Set up Helicone\nHELICONE_API_KEY = \"your_helicone_api_key_here\"  # Replace with your actual Helicone API key\nhelicone_config.api_key = HELICONE_API_KEY\nos.environ[\"OPENAI_API_KEY\"] = (\n    \"your_openai_api_key_here\"  # Replace with your actual OpenAI API key\n)\n\n# Verify Helicone API key is set\nif HELICONE_API_KEY == \"your_helicone_api_key_here\":\n    raise ValueError(\n        \"Please replace 'your_helicone_api_key_here' with your actual Helicone API key.\"\n    )"
+   "source": [
+    "import os\n",
+    "from datasets import Dataset\n",
+    "from ragas import evaluate\n",
+    "from ragas.metrics import faithfulness, answer_relevancy, context_precision\n",
+    "from ragas.integrations.helicone import helicone_config  # import helicone_config\n",
+    "\n",
+    "\n",
+    "# Set up Helicone\n",
+    "HELICONE_API_KEY = (\n",
+    "    \"your_helicone_api_key_here\"  # Replace with your actual Helicone API key\n",
+    ")\n",
+    "helicone_config.api_key = HELICONE_API_KEY\n",
+    "os.environ[\"OPENAI_API_KEY\"] = (\n",
+    "    \"your_openai_api_key_here\"  # Replace with your actual OpenAI API key\n",
+    ")\n",
+    "\n",
+    "# Verify Helicone API key is set\n",
+    "if HELICONE_API_KEY == \"your_helicone_api_key_here\":\n",
+    "    raise ValueError(\n",
+    "        \"Please replace 'your_helicone_api_key_here' with your actual Helicone API key.\"\n",
+    "    )"
+   ]
   },
   {
    "cell_type": "markdown",

ragas/pyproject.toml

@@ -63,6 +63,7 @@ dev = [
     "haystack-ai",
     "sacrebleu",
     "r2r",
+    "scipy",
 ]
 test = [
     "pytest",

ragas/src/ragas/embeddings/haystack_wrapper.py

@@ -37,10 +37,18 @@ def __init__(
         # Lazy Import of required Haystack components
         try:
             from haystack import AsyncPipeline
-            from haystack.components.embedders.azure_text_embedder import AzureOpenAITextEmbedder
-            from haystack.components.embedders.hugging_face_api_text_embedder import HuggingFaceAPITextEmbedder
-            from haystack.components.embedders.openai_text_embedder import OpenAITextEmbedder
-            from haystack.components.embedders.sentence_transformers_text_embedder import SentenceTransformersTextEmbedder
+            from haystack.components.embedders.azure_text_embedder import (
+                AzureOpenAITextEmbedder,
+            )
+            from haystack.components.embedders.hugging_face_api_text_embedder import (
+                HuggingFaceAPITextEmbedder,
+            )
+            from haystack.components.embedders.openai_text_embedder import (
+                OpenAITextEmbedder,
+            )
+            from haystack.components.embedders.sentence_transformers_text_embedder import (
+                SentenceTransformersTextEmbedder,
+            )
         except ImportError as exc:
             raise ImportError(
                 "Haystack is not installed. Please install it with `pip install haystack-ai`."
@@ -94,10 +102,18 @@ async def aembed_documents(self, texts: t.List[str]) -> t.List[t.List[float]]:
 
     def __repr__(self) -> str:
         try:
-            from haystack.components.embedders.azure_text_embedder import AzureOpenAITextEmbedder
-            from haystack.components.embedders.hugging_face_api_text_embedder import HuggingFaceAPITextEmbedder
-            from haystack.components.embedders.openai_text_embedder import OpenAITextEmbedder
-            from haystack.components.embedders.sentence_transformers_text_embedder import SentenceTransformersTextEmbedder
+            from haystack.components.embedders.azure_text_embedder import (
+                AzureOpenAITextEmbedder,
+            )
+            from haystack.components.embedders.hugging_face_api_text_embedder import (
+                HuggingFaceAPITextEmbedder,
+            )
+            from haystack.components.embedders.openai_text_embedder import (
+                OpenAITextEmbedder,
+            )
+            from haystack.components.embedders.sentence_transformers_text_embedder import (
+                SentenceTransformersTextEmbedder,
+            )
         except ImportError:
             return f"{self.__class__.__name__}(embeddings=Unknown(...))"
 

ragas/src/ragas/llms/haystack_wrapper.py

@@ -39,8 +39,12 @@ def __init__(
         try:
             from haystack import AsyncPipeline
             from haystack.components.generators.azure import AzureOpenAIGenerator
-            from haystack.components.generators.hugging_face_api import HuggingFaceAPIGenerator
-            from haystack.components.generators.hugging_face_local import HuggingFaceLocalGenerator
+            from haystack.components.generators.hugging_face_api import (
+                HuggingFaceAPIGenerator,
+            )
+            from haystack.components.generators.hugging_face_local import (
+                HuggingFaceLocalGenerator,
+            )
             from haystack.components.generators.openai import OpenAIGenerator
         except ImportError as exc:
             raise ImportError(
@@ -115,8 +119,12 @@ async def agenerate_text(
     def __repr__(self) -> str:
         try:
             from haystack.components.generators.azure import AzureOpenAIGenerator
-            from haystack.components.generators.hugging_face_api import HuggingFaceAPIGenerator
-            from haystack.components.generators.hugging_face_local import HuggingFaceLocalGenerator
+            from haystack.components.generators.hugging_face_api import (
+                HuggingFaceAPIGenerator,
+            )
+            from haystack.components.generators.hugging_face_local import (
+                HuggingFaceLocalGenerator,
+            )
             from haystack.components.generators.openai import OpenAIGenerator
         except ImportError:
             return f"{self.__class__.__name__}(llm=Unknown(...))"

ragas/src/ragas/testset/transforms/relationship_builders/cosine.py

@@ -12,59 +12,111 @@ class CosineSimilarityBuilder(RelationshipBuilder):
     property_name: str = "embedding"
     new_property_name: str = "cosine_similarity"
     threshold: float = 0.9
+    block_size: int = 1024
 
-    def _find_similar_embedding_pairs(
-        self, embeddings: np.ndarray, threshold: float
-    ) -> t.List[t.Tuple[int, int, float]]:
-        # Normalize the embeddings
-        normalized = embeddings / np.linalg.norm(embeddings, axis=1)[:, np.newaxis]
+    def _validate_embedding_shapes(self, embeddings: t.List[t.Any]):
+        if not embeddings:
+            raise ValueError(f"No nodes have a valid {self.property_name}")
+        first_len = len(embeddings[0])
+        for idx, emb in enumerate(embeddings):
+            if len(emb) != first_len:
+                raise ValueError(
+                    f"Embedding at index {idx} has length {len(emb)}, expected {first_len}. "
+                    "All embeddings must have the same length."
+                )
 
-        # Calculate cosine similarity matrix
-        similarity_matrix = np.dot(normalized, normalized.T)
-        # Find pairs with similarity >= threshold
-        similar_pairs = np.argwhere(similarity_matrix >= threshold)
+    def _block_cosine_similarity(self, i: np.ndarray, j: np.ndarray):
+        """Calculate cosine similarity matrix between two sets of embeddings."""
+        i_norm = i / np.linalg.norm(i, axis=1, keepdims=True)
+        j_norm = j / np.linalg.norm(j, axis=1, keepdims=True)
+        return np.dot(i_norm, j_norm.T)
 
-        # Filter out self-comparisons and duplicate pairs
-        return [
-            (pair[0], pair[1], similarity_matrix[pair[0], pair[1]])
-            for pair in similar_pairs
-            if pair[0] < pair[1]
-        ]
+    async def _find_similar_embedding_pairs(
+        self, embeddings: np.ndarray, threshold: float, block_size: int = 1024
+    ) -> t.Set[t.Tuple[int, int, float]]:
+        """Sharded computation of cosine similarity to find similar pairs."""
 
-    async def transform(self, kg: KnowledgeGraph) -> t.List[Relationship]:
-        if self.property_name is None:
-            self.property_name = "embedding"
+        def process_block(i: int, j: int) -> t.Set[t.Tuple[int, int, float]]:
+            end_i = min(i + block_size, n_embeddings)
+            end_j = min(j + block_size, n_embeddings)
+            block = self._block_cosine_similarity(
+                embeddings[i:end_i, :], embeddings[j:end_j, :]
+            )
+            similar_idx = np.argwhere(block >= threshold)
+            return {
+                (int(i + ii), int(j + jj), float(block[ii, jj]))
+                for ii, jj in similar_idx
+                if int(i + ii) < int(j + jj)
+            }
+
+        n_embeddings, _dimension = embeddings.shape
+        triplets = set()
 
+        for i in range(0, n_embeddings, block_size):
+            for j in range(i, n_embeddings, block_size):
+                triplets.update(process_block(i, j))
+
+        return triplets
+
+    async def transform(self, kg: KnowledgeGraph) -> t.List[Relationship]:
         embeddings = []
         for node in kg.nodes:
             embedding = node.get_property(self.property_name)
             if embedding is None:
                 raise ValueError(f"Node {node.id} has no {self.property_name}")
             embeddings.append(embedding)
-
-        similar_pairs = self._find_similar_embedding_pairs(
-            np.array(embeddings), self.threshold
+        self._validate_embedding_shapes(embeddings)
+        similar_pairs = await self._find_similar_embedding_pairs(
+            np.array(embeddings), self.threshold, self.block_size
         )
-
         return [
             Relationship(
                 source=kg.nodes[i],
                 target=kg.nodes[j],
-                type="cosine_similarity",
+                type=self.new_property_name,
                 properties={self.new_property_name: similarity_float},
                 bidirectional=True,
             )
             for i, j, similarity_float in similar_pairs
         ]
 
+    def generate_execution_plan(self, kg: KnowledgeGraph) -> t.List[t.Coroutine]:
+        """
+        Generates a coroutine task for finding similar embedding pairs, which can be scheduled/executed by an Executor.
+        """
+        embeddings = []
+        for node in kg.nodes:
+            embedding = node.get_property(self.property_name)
+            if embedding is None:
+                raise ValueError(f"Node {node.id} has no {self.property_name}")
+            embeddings.append(embedding)
+        self._validate_embedding_shapes(embeddings)
+
+        async def find_and_add_relationships():
+            similar_pairs = await self._find_similar_embedding_pairs(
+                np.array(embeddings), self.threshold, self.block_size
+            )
+            for i, j, similarity_float in similar_pairs:
+                rel = Relationship(
+                    source=kg.nodes[i],
+                    target=kg.nodes[j],
+                    type=self.new_property_name,
+                    properties={self.new_property_name: similarity_float},
+                    bidirectional=True,
+                )
+                kg.relationships.append(rel)
+
+        return [find_and_add_relationships()]
+
 
 @dataclass
 class SummaryCosineSimilarityBuilder(CosineSimilarityBuilder):
     property_name: str = "summary_embedding"
     new_property_name: str = "summary_cosine_similarity"
     threshold: float = 0.1
+    block_size: int = 1024
 
-    def filter(self, kg: KnowledgeGraph) -> KnowledgeGraph:
+    def _document_summary_filter(self, kg: KnowledgeGraph) -> KnowledgeGraph:
         """
         Filters the knowledge graph to only include nodes with a summary embedding.
         """
@@ -78,21 +130,22 @@ def filter(self, kg: KnowledgeGraph) -> KnowledgeGraph:
         return KnowledgeGraph(nodes=nodes)
 
     async def transform(self, kg: KnowledgeGraph) -> t.List[Relationship]:
+        filtered_kg = self._document_summary_filter(kg)
         embeddings = [
             node.get_property(self.property_name)
-            for node in kg.nodes
+            for node in filtered_kg.nodes
             if node.get_property(self.property_name) is not None
         ]
         if not embeddings:
             raise ValueError(f"No nodes have a valid {self.property_name}")
-        similar_pairs = self._find_similar_embedding_pairs(
-            np.array(embeddings), self.threshold
+        similar_pairs = await self._find_similar_embedding_pairs(
+            np.array(embeddings), self.threshold, self.block_size
         )
         return [
             Relationship(
-                source=kg.nodes[i],
-                target=kg.nodes[j],
-                type="summary_cosine_similarity",
+                source=filtered_kg.nodes[i],
+                target=filtered_kg.nodes[j],
+                type=self.new_property_name,
                 properties={self.new_property_name: similarity_float},
                 bidirectional=True,
             )