Merge branch 'run-llama:main' into main

Author: mhaligowski

docs/docs/examples/embeddings/textembed.ipynb

@@ -6,7 +6,7 @@
    "source": [
     "# TiDB Vector Store\n",
     "\n",
-    "> [TiDB Cloud](https://tidbcloud.com/), is a comprehensive Database-as-a-Service (DBaaS) solution, that provides dedicated and serverless options. TiDB Serverless is now integrating a built-in vector search into the MySQL landscape. With this enhancement, you can seamlessly develop AI applications using TiDB Serverless without the need for a new database or additional technical stacks. Be among the first to experience it by joining the waitlist for the private beta at https://tidb.cloud/ai.\n",
+    "> [TiDB Cloud](https://www.pingcap.com/tidb-serverless/), is a comprehensive Database-as-a-Service (DBaaS) solution, that provides dedicated and serverless options. TiDB Serverless is now integrating a built-in vector search into the MySQL landscape. With this enhancement, you can seamlessly develop AI applications using TiDB Serverless without the need for a new database or additional technical stacks. Create a free TiDB Serverless cluster and start using the vector search feature at https://pingcap.com/ai.\n",
     "\n",
     "This notebook provides a detailed guide on utilizing the tidb vector search in LlamaIndex."
    ]

docs/docs/examples/vector_stores/TiDBVector.ipynb

@@ -33,6 +33,9 @@
 )
 from llama_index.core.utils import print_text
 from llama_index.core.query_pipeline.components.stateful import BaseStatefulComponent
+import llama_index.core.instrumentation as instrument
+
+dispatcher = instrument.get_dispatcher(__name__)
 
 
 # TODO: Make this (safely) pydantic?
@@ -386,6 +389,7 @@ def set_callback_manager(self, callback_manager: CallbackManager) -> None:
         for module in self.module_dict.values():
             module.set_callback_manager(callback_manager)
 
+    @dispatcher.span
     def run(
         self,
         *args: Any,
@@ -529,6 +533,7 @@ def run_multi_with_intermediates(
             ) as query_event:
                 return self._run_multi(module_input_dict, show_intermediates=True)
 
+    @dispatcher.span
     async def arun(
         self,
         *args: Any,
@@ -725,6 +730,7 @@ def _get_single_result_output(
         else:
             return result_output
 
+    @dispatcher.span
     def _run(
         self,
         *args: Any,
@@ -780,6 +786,7 @@ def _run(
                 intermediates,
             )
 
+    @dispatcher.span
     async def _arun(
         self,
         *args: Any,
@@ -904,6 +911,7 @@ def get_run_state(
             root_key, kwargs = self._get_root_key_and_kwargs(**pipeline_inputs)
             return RunState(self.module_dict, {root_key: kwargs})
 
+    @dispatcher.span
     def _run_multi(
         self, module_input_dict: Dict[str, Any], show_intermediates=False
     ) -> Tuple[Dict[str, Any], Dict[str, ComponentIntermediates]]:
@@ -947,6 +955,7 @@ def _run_multi(
 
         return run_state.result_outputs, run_state.intermediate_outputs
 
+    @dispatcher.span
     async def _arun_multi(
         self, module_input_dict: Dict[str, Any], show_intermediates: bool = False
     ) -> Tuple[Dict[str, Any], Dict[str, ComponentIntermediates]]:

llama-index-core/llama_index/core/query_pipeline/query.py

@@ -0,0 +1,3 @@
+poetry_requirements(
+    name="poetry",
+)

llama-index-integrations/embeddings/llama-index-embeddings-textembed/BUILD

@@ -0,0 +1,17 @@
+GIT_ROOT ?= $(shell git rev-parse --show-toplevel)
+
+help:	## Show all Makefile targets.
+	@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[33m%-30s\033[0m %s\n", $$1, $$2}'
+
+format:	## Run code autoformatters (black).
+	pre-commit install
+	git ls-files | xargs pre-commit run black --files
+
+lint:	## Run linters: pre-commit (black, ruff, codespell) and mypy
+	pre-commit install && git ls-files | xargs pre-commit run --show-diff-on-failure --files
+
+test:	## Run tests via pytest.
+	pytest tests
+
+watch-docs:	## Build and watch documentation.
+	sphinx-autobuild docs/ docs/_build/html --open-browser --watch $(GIT_ROOT)/llama_index/

llama-index-integrations/embeddings/llama-index-embeddings-textembed/Makefile

@@ -0,0 +1,64 @@
+# TextEmbed - Embedding Inference Server
+
+Maintained by Keval Dekivadiya, TextEmbed is licensed under the [Apache-2.0 License](https://opensource.org/licenses/Apache-2.0).
+
+TextEmbed is a high-throughput, low-latency REST API designed for serving vector embeddings. It supports a wide range of sentence-transformer models and frameworks, making it suitable for various applications in natural language processing.
+
+## Features
+
+- **High Throughput & Low Latency**: Designed to handle a large number of requests efficiently.
+- **Flexible Model Support**: Works with various sentence-transformer models.
+- **Scalable**: Easily integrates into larger systems and scales with demand.
+- **Batch Processing**: Supports batch processing for better and faster inference.
+- **OpenAI Compatible REST API Endpoint**: Provides an OpenAI compatible REST API endpoint.
+- **Single Line Command Deployment**: Deploy multiple models via a single command for efficient deployment.
+- **Support for Embedding Formats**: Supports binary, float16, and float32 embeddings formats for faster retrieval.
+
+## Getting Started
+
+### Prerequisites
+
+Ensure you have Python 3.10 or higher installed. You will also need to install the required dependencies.
+
+### Installation via PyPI
+
+Install the required dependencies:
+
+```bash
+pip install -U textembed
+```
+
+### Start the TextEmbed Server
+
+Start the TextEmbed server with your desired models:
+
+```bash
+python -m textembed.server --models sentence-transformers/all-MiniLM-L12-v2 --workers 4 --api-key TextEmbed
+```
+
+### Example Usage with llama-index
+
+Here's a simple example to get you started with llama-index:
+
+```python
+from llama_index.embeddings.textembed import TextEmbedEmbedding
+
+# Initialize the TextEmbedEmbedding class
+embed = TextEmbedEmbedding(
+    model_name="sentence-transformers/all-MiniLM-L12-v2",
+    base_url="http://0.0.0.0:8000/v1",
+    auth_token="TextEmbed",
+)
+
+# Get embeddings for a batch of texts
+embeddings = embed.get_text_embedding_batch(
+    [
+        "It is raining cats and dogs here!",
+        "India has a diverse cultural heritage.",
+    ]
+)
+
+print(embeddings)
+```
+
+For more information, please read the [documentation](https://github.com/kevaldekivadiya2415/textembed/blob/main/docs/setup.md).

llama-index-integrations/embeddings/llama-index-embeddings-textembed/README.md

@@ -0,0 +1 @@
+python_sources()

llama-index-integrations/embeddings/llama-index-embeddings-textembed/llama_index/embeddings/textembed/BUILD

@@ -0,0 +1,3 @@
+from llama_index.embeddings.textembed.base import TextEmbedEmbedding
+
+__all__ = ["TextEmbedEmbedding"]

llama-index-integrations/embeddings/llama-index-embeddings-textembed/llama_index/embeddings/textembed/__init__.py

@@ -0,0 +1,203 @@
+"""TextEmbed: Embedding Inference Server.
+
+TextEmbed offers a high-throughput, low-latency service for generating embeddings using various sentence-transformer models.
+It now also supports image embedding models, providing flexibility and scalability for diverse applications.
+
+Maintained by Keval Dekivadiya, TextEmbed is licensed under Apache-2.0.
+"""
+
+from typing import Callable, List, Optional, Union
+
+import aiohttp
+import requests
+
+from llama_index.core.base.embeddings.base import (
+    DEFAULT_EMBED_BATCH_SIZE,
+    BaseEmbedding,
+)
+from llama_index.core.bridge.pydantic import Field
+from llama_index.core.callbacks import CallbackManager
+
+DEFAULT_URL = "http://0.0.0.0:8000/v1"
+
+
+class TextEmbedEmbedding(BaseEmbedding):
+    """TextEmbedEmbedding is a class for interfacing with the TextEmbed: embedding inference server."""
+
+    base_url: str = Field(
+        default=DEFAULT_URL,
+        description="Base URL for the text embeddings service.",
+    )
+    timeout: float = Field(
+        default=60.0,
+        description="Timeout in seconds for the request.",
+    )
+    auth_token: Optional[Union[str, Callable[[str], str]]] = Field(
+        default=None,
+        description="Authentication token or authentication token generating function for authenticated requests",
+    )
+
+    def __init__(
+        self,
+        model_name: str,
+        base_url: str = DEFAULT_URL,
+        embed_batch_size: int = DEFAULT_EMBED_BATCH_SIZE,
+        timeout: float = 60.0,
+        callback_manager: Optional[CallbackManager] = None,
+        auth_token: Optional[Union[str, Callable[[str], str]]] = None,
+    ):
+        """
+        Initializes the TextEmbedEmbedding object with specified parameters.
+
+        Args:
+            model_name (str): The name of the model to be used for embeddings.
+            base_url (str): The base URL of the embedding service.
+            embed_batch_size (int): The batch size for embedding requests.
+            timeout (float): Timeout for requests.
+            callback_manager (Optional[CallbackManager]): Manager for handling callbacks.
+            auth_token (Optional[Union[str, Callable[[str], str]]]): Authentication token or function for generating it.
+        """
+        super().__init__(
+            base_url=base_url,
+            model_name=model_name,
+            embed_batch_size=embed_batch_size,
+            timeout=timeout,
+            callback_manager=callback_manager,
+            auth_token=auth_token,
+        )
+
+    def _call_api(self, texts: List[str]) -> List[List[float]]:
+        """
+        Calls the TextEmbed API to get embeddings for a list of texts.
+
+        Args:
+            texts (List[str]): A list of texts to get embeddings for.
+
+        Returns:
+            List[List[float]]: A list of embeddings for the input texts.
+
+        Raises:
+            Exception: If the API responds with a status code other than 200.
+        """
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.auth_token}" if self.auth_token else None,
+        }
+        json_data = {"input": texts, "model": self.model_name}
+        with requests.post(
+            f"{self.base_url}/embedding",
+            headers=headers,
+            json=json_data,
+            timeout=self.timeout,
+        ) as response:
+            if response.status_code != 200:
+                raise Exception(
+                    f"TextEmbed responded with an unexpected status message "
+                    f"{response.status_code}: {response.text}"
+                )
+            return [e["embedding"] for e in response.json()["data"]]
+
+    async def _acall_api(self, texts: List[str]) -> List[List[float]]:
+        """
+        Asynchronously calls the TextEmbed API to get embeddings for a list of texts.
+
+        Args:
+            texts (List[str]): A list of texts to get embeddings for.
+
+        Returns:
+            List[List[float]]: A list of embeddings for the input texts.
+
+        Raises:
+            Exception: If the API responds with a status code other than 200.
+        """
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.auth_token}" if self.auth_token else None,
+        }
+        json_data = {"input": texts, "model": self.model_name}
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{self.base_url}/embedding",
+                headers=headers,
+                json=json_data,
+                timeout=self.timeout,
+            ) as response:
+                if response.status != 200:
+                    raise Exception(
+                        f"TextEmbed responded with an unexpected status message "
+                        f"{response.status}: {response.text}"
+                    )
+                data = await response.json()
+                return [e["embedding"] for e in data["data"]]
+
+    def _get_query_embedding(self, query: str) -> List[float]:
+        """
+        Gets the embedding for a single query.
+
+        Args:
+            query (str): The query to get the embedding for.
+
+        Returns:
+            List[float]: The embedding for the query.
+        """
+        return self._call_api([query])[0]
+
+    def _get_text_embedding(self, text: str) -> List[float]:
+        """
+        Gets the embedding for a single text.
+
+        Args:
+            text (str): The text to get the embedding for.
+
+        Returns:
+            List[float]: The embedding for the text.
+        """
+        return self._call_api([text])[0]
+
+    def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
+        """
+        Gets the embeddings for a list of texts.
+
+        Args:
+            texts (List[str]): The texts to get the embeddings for.
+
+        Returns:
+            List[List[float]]: A list of embeddings for the input texts.
+        """
+        return self._call_api(texts)
+
+    async def _aget_query_embedding(self, query: str) -> List[float]:
+        """
+        Asynchronously gets the embedding for a single query.
+
+        Args:
+            query (str): The query to get the embedding for.
+
+        Returns:
+            List[float]: The embedding for the query.
+        """
+        return (await self._acall_api([query]))[0]
+
+    async def _aget_text_embedding(self, text: str) -> List[float]:
+        """
+        Asynchronously gets the embedding for a single text.
+
+        Args:
+            text (str): The text to get the embedding for.
+
+        Returns:
+            List[float]: The embedding for the text.
+        """
+        return (await self._acall_api([text]))[0]
+
+    async def _aget_text_embeddings(self, texts: List[str]) -> List[List[float]]:
+        """
+        Asynchronously gets the embeddings for a list of texts.
+
+        Args:
+            texts (List[str]): The texts to get the embeddings for.
+
+        Returns:
+            List[List[float]]: A list of embeddings for the input texts.
+        """
+        return await self._acall_api(texts)