chore: import upstream snapshot with attribution

knowledge_storm/encoder.py

@@ -0,0 +1,178 @@
+import os
+import numpy as np
+
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import List, Tuple, Union, Optional, Dict, Literal
+from pathlib import Path
+
+try:
+    import warnings
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=UserWarning)
+        if "LITELLM_LOCAL_MODEL_COST_MAP" not in os.environ:
+            os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
+        import litellm
+
+        litellm.drop_params = True
+        litellm.telemetry = False
+
+    from litellm.caching.caching import Cache
+
+    disk_cache_dir = os.path.join(Path.home(), ".storm_local_cache")
+    litellm.cache = Cache(disk_cache_dir=disk_cache_dir, type="disk")
+
+except ImportError:
+
+    class LitellmPlaceholder:
+        def __getattr__(self, _):
+            raise ImportError(
+                "The LiteLLM package is not installed. Run `pip install litellm`."
+            )
+
+    litellm = LitellmPlaceholder()
+
+
+class Encoder:
+    """
+    A wrapper class for the LiteLLM embedding model, designed to handle embedding
+    generation tasks efficiently. It supports parallel processing and local caching of
+    embedding results for improved performance.
+
+    The Encoder utilizes the LiteLLM library to interact with various embedding models,
+    such as OpenAI and Azure embeddings. Users can specify the desired encoder type and
+    provide relevant API credentials during initialization.
+
+    Features:
+        - Support for multiple embedding models (e.g., OpenAI, Azure).
+        - Parallel processing for faster embedding generation.
+        - Local disk caching to store and reuse embedding results.
+        - Total token usage tracking for cost monitoring.
+
+    Note:
+        Refer to the LiteLLM documentation for details on supported embedding models:
+        https://docs.litellm.ai/docs/embedding/supported_embedding
+    """
+
+    def __init__(
+        self,
+        encoder_type: Optional[str] = None,
+        api_key: Optional[str] = None,
+        api_base: Optional[str] = None,
+        api_version: Optional[str] = None,
+    ):
+        """
+        Initializes the Encoder with the appropriate embedding model.
+
+        Args:
+            encoder_type (Optional[str]): Type of encoder ('openai', 'azure', etc.).
+            api_key (Optional[str]): API key for the encoder service.
+            api_base (Optional[str]): API base URL for the encoder service.
+            api_version (Optional[str]): API version for the encoder service.
+        """
+        self.embedding_model_name = None
+        self.kargs = {}
+        self.total_token_usage = 0
+
+        # Initialize the appropriate embedding model
+        encoder_type = encoder_type or os.getenv("ENCODER_API_TYPE")
+        if not encoder_type:
+            raise ValueError("ENCODER_API_TYPE environment variable is not set.")
+
+        if encoder_type.lower() == "openai":
+            self.embedding_model_name = "text-embedding-3-small"
+            self.kargs = {"api_key": api_key or os.getenv("OPENAI_API_KEY")}
+        elif encoder_type.lower() == "azure":
+            self.embedding_model_name = "azure/text-embedding-3-small"
+            self.kargs = {
+                "api_key": api_key or os.getenv("AZURE_API_KEY"),
+                "api_base": api_base or os.getenv("AZURE_API_BASE"),
+                "api_version": api_version or os.getenv("AZURE_API_VERSION"),
+            }
+        else:
+            raise ValueError(
+                f"Unsupported ENCODER_API_TYPE '{encoder_type}'. Supported types are 'openai', 'azure', 'together'."
+            )
+
+    def get_total_token_usage(self, reset: bool = False) -> int:
+        """
+        Retrieves the total token usage.
+
+        Args:
+            reset (bool): If True, resets the total token usage counter after retrieval.
+
+        Returns:
+            int: The total number of tokens used.
+        """
+        token_usage = self.total_token_usage
+        if reset:
+            self.total_token_usage = 0
+        return token_usage
+
+    def encode(self, texts: Union[str, List[str]], max_workers: int = 5) -> np.ndarray:
+        """
+        Public method to get embeddings for the given texts.
+
+        Args:
+            texts (Union[str, List[str]]): A single text string or a list of text strings to embed.
+
+        Returns:
+            np.ndarray: The array of embeddings.
+        """
+        return self._get_text_embeddings(texts, max_workers=max_workers)
+
+    def _get_single_text_embedding(self, text):
+        response = litellm.embedding(
+            model=self.embedding_model_name, input=text, caching=True, **self.kargs
+        )
+        embedding = response.data[0]["embedding"]
+        token_usage = response.get("usage", {}).get("total_tokens", 0)
+        return text, embedding, token_usage
+
+    def _get_text_embeddings(
+        self,
+        texts: Union[str, List[str]],
+        max_workers: int = 5,
+    ) -> Tuple[np.ndarray, int]:
+        """
+        Get text embeddings using OpenAI's text-embedding-3-small model.
+
+        Args:
+            texts (Union[str, List[str]]): A single text string or a list of text strings to embed.
+            max_workers (int): The maximum number of workers for parallel processing.
+            api_key (str): The API key for accessing OpenAI's services.
+            embedding_cache (Optional[Dict[str, np.ndarray]]): A cache to store previously computed embeddings.
+
+        Returns:
+            Tuple[np.ndarray, int]: The 2D array of embeddings and the total token usage.
+        """
+
+        if isinstance(texts, str):
+            _, embedding, tokens = self._get_single_text_embedding(texts)
+            self.total_token_usage += tokens
+            return np.array(embedding)
+
+        embeddings = []
+        total_tokens = 0
+
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            futures = {
+                executor.submit(self._get_single_text_embedding, text): text
+                for text in texts
+            }
+
+            for future in as_completed(futures):
+                try:
+                    text, embedding, tokens = future.result()
+                    embeddings.append((text, embedding, tokens))
+                    total_tokens += tokens
+                except Exception as e:
+                    print(f"An error occurred for text: {futures[future]}")
+                    print(e)
+
+        # Sort results to match the order of the input texts
+        embeddings.sort(key=lambda x: texts.index(x[0]))
+        embeddings = [result[1] for result in embeddings]
+        self.total_token_usage += total_tokens
+
+        return np.array(embeddings)