Add new LLM and embedding model provider Novita AI

README.md

@@ -90,7 +90,7 @@ result = query("Write a report about xxx.") # Your question here
 #### LLM Configuration
 
 <pre><code>config.set_provider_config("llm", "(LLMName)", "(Arguments dict)")</code></pre>
-<p>The "LLMName" can be one of the following: ["DeepSeek", "OpenAI", "XAI", "SiliconFlow", "PPIO", "TogetherAI", "Gemini", "Ollama"]</p>
+<p>The "LLMName" can be one of the following: ["DeepSeek", "OpenAI", "XAI", "SiliconFlow", "PPIO", "TogetherAI", "Gemini", "Ollama", "Novita"]</p>
 <p> The "Arguments dict" is a dictionary that contains the necessary arguments for the LLM class.</p>
 
 <details>
@@ -187,7 +187,7 @@ result = query("Write a report about xxx.") # Your question here
 
 #### Embedding Model Configuration
 <pre><code>config.set_provider_config("embedding", "(EmbeddingModelName)", "(Arguments dict)")</code></pre>
-<p>The "EmbeddingModelName" can be one of the following: ["MilvusEmbedding", "OpenAIEmbedding", "VoyageEmbedding", "SiliconflowEmbedding", "PPIOEmbedding"]</p>
+<p>The "EmbeddingModelName" can be one of the following: ["MilvusEmbedding", "OpenAIEmbedding", "VoyageEmbedding", "SiliconflowEmbedding", "PPIOEmbedding", "NovitaEmbedding"]</p>
 <p> The "Arguments dict" is a dictionary that contains the necessary arguments for the embedding model class.</p>
 
 <details>
@@ -221,6 +221,13 @@ result = query("Write a report about xxx.") # Your question here
   <p> You need to install boto3 before running, execute: <code>pip install boto3</code>. More details about Amazon Bedrock: https://docs.aws.amazon.com/bedrock/ </p>
 </details>
 
+<details>
+  <summary>Example (Novita AI embedding)</summary>
+    <p> Make sure you have prepared your Novita AI API KEY as an env variable <code>NOVITA_API_KEY</code>.</p>
+    <pre><code>config.set_provider_config("embedding", "NovitaEmbedding", {"model": "baai/bge-m3"})</code></pre>
+    <p> More details about Novita AI: https://novita.ai/docs/api-reference/model-apis-llm-create-embeddings?utm_source=github_deep-searcher&utm_medium=github_readme&utm_campaign=link </p>
+</details>
+
 <details>
   <summary>Example (Siliconflow embedding)</summary>
     <p> Make sure you have prepared your Siliconflow API KEY as an env variable <code>SILICONFLOW_API_KEY</code>.</p>
@@ -440,6 +447,7 @@ nest_asyncio.apply()
 - [Amazon Bedrock](https://docs.aws.amazon.com/bedrock/) (`AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` env variable required)
 - [FastEmbed](https://qdrant.github.io/fastembed/)
 - [PPIO](https://ppinfra.com/model-api/product/llm-api?utm_source=github_deep-searcher) (`PPIO_API_KEY` env variable required)
+- [Novita AI](https://novita.ai/docs/api-reference/model-apis-llm-create-embeddings?utm_source=github_deep-searcher&utm_medium=github_readme&utm_campaign=link) (`NOVITA_API_KEY` env variable required)
 
 ### 🔹 LLM Support
 - [OpenAI](https://platform.openai.com/docs/models) (`OPENAI_API_KEY` env variable required)
@@ -452,6 +460,7 @@ nest_asyncio.apply()
 - [Google Gemini](https://ai.google.dev/gemini-api/docs) (`GEMINI_API_KEY` env variable required)
 - [SambaNova Cloud Inference Service](https://docs.together.ai/docs/introduction) (`SAMBANOVA_API_KEY` env variable required)
 - [Ollama](https://ollama.com/)
+- [Novita AI](https://novita.ai/docs/guides/introduction?utm_source=github_deep-searcher&utm_medium=github_readme&utm_campaign=link) (`NOVITA_API_KEY` env variable required)
 
 ### 🔹 Document Loader
 - Local File

deepsearcher/config.yaml

@@ -41,6 +41,11 @@ provide_settings:
 #      model: "qwq"
 ##      base_url: ""
 
+#    provider: "Novita"
+#    config:
+#      model: "deepseek/deepseek-v3-0324"
+##      api_key: "sk_xxxxxx"  # Uncomment to override the `NOVITA_API_KEY` set in the environment variable
+
   embedding:
     provider: "OpenAIEmbedding"
     config:
@@ -85,6 +90,11 @@ provide_settings:
 #    config:
 #      model: "BAAI/bge-small-en-v1.5"
 
+#    provider: "NovitaEmbedding"
+#    config:
+#      model: "baai/bge-m3"
+##      api_key: "sk_xxxxxx"  # Uncomment to override the `NOVITA_API_KEY` set in the environment variable
+
 
   file_loader:
     provider: "PDFLoader"

deepsearcher/embedding/__init__.py

@@ -3,6 +3,7 @@
 from .gemini_embedding import GeminiEmbedding
 from .glm_embedding import GLMEmbedding
 from .milvus_embedding import MilvusEmbedding
+from .novita_embedding import NovitaEmbedding
 from .ollama_embedding import OllamaEmbedding
 from .openai_embedding import OpenAIEmbedding
 from .ppio_embedding import PPIOEmbedding
@@ -22,4 +23,5 @@
     "GLMEmbedding",
     "OllamaEmbedding",
     "FastEmbedEmbedding",
+    "NovitaEmbedding",
 ]

deepsearcher/embedding/novita_embedding.py

@@ -0,0 +1,128 @@
+import os
+from typing import List, Union
+
+import requests
+
+from deepsearcher.embedding.base import BaseEmbedding
+
+NOVITA_MODEL_DIM_MAP = {
+    "baai/bge-m3": 1024,
+}
+
+NOVITA_EMBEDDING_API = "https://api.novita.ai/v3/openai/embeddings"
+
+
+class NovitaEmbedding(BaseEmbedding):
+    """
+    Novita AI embedding model implementation.
+
+    This class provides an interface to the Novita AI embedding API, which offers
+    various embedding models for text processing.
+
+    For more information, see:
+    https://novita.ai/docs/api-reference/model-apis-llm-create-embeddings
+    """
+
+    def __init__(self, model="baai/bge-m3", batch_size=32, **kwargs):
+        """
+        Initialize the Novita AI embedding model.
+
+        Args:
+            model (str): The model identifier to use for embeddings. Default is "baai/bge-m3".
+            batch_size (int): Maximum number of texts to process in a single batch. Default is 32.
+            **kwargs: Additional keyword arguments.
+                - api_key (str, optional): The Novita AI API key. If not provided,
+                  it will be read from the NOVITA_API_KEY environment variable.
+                - model_name (str, optional): Alternative way to specify the model.
+
+        Raises:
+            RuntimeError: If no API key is provided or found in environment variables.
+        """
+        if "model_name" in kwargs and (not model or model == "baai/bge-m3"):
+            model = kwargs.pop("model_name")
+        self.model = model
+        if "api_key" in kwargs:
+            api_key = kwargs.pop("api_key")
+        else:
+            api_key = os.getenv("NOVITA_API_KEY")
+
+        if not api_key or len(api_key) == 0:
+            raise RuntimeError("api_key is required for NovitaEmbedding")
+        self.api_key = api_key
+        self.batch_size = batch_size
+
+    def embed_query(self, text: str) -> List[float]:
+        """
+        Embed a single query text.
+
+        Args:
+            text (str): The query text to embed.
+
+        Returns:
+            List[float]: A list of floats representing the embedding vector.
+
+        Note:
+            For retrieval cases, this method uses "query" as the input type.
+        """
+        return self._embed_input(text)[0]
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """
+        Embed a list of document texts.
+
+        This method handles batching of document embeddings based on the configured
+        batch size to optimize API calls.
+
+        Args:
+            texts (List[str]): A list of document texts to embed.
+
+        Returns:
+            List[List[float]]: A list of embedding vectors, one for each input text.
+        """
+        # batch embedding
+        if self.batch_size > 0:
+            if len(texts) > self.batch_size:
+                batch_texts = [
+                    texts[i : i + self.batch_size] for i in range(0, len(texts), self.batch_size)
+                ]
+                embeddings = []
+                for batch_text in batch_texts:
+                    batch_embeddings = self._embed_input(batch_text)
+                    embeddings.extend(batch_embeddings)
+                return embeddings
+            return self._embed_input(texts)
+        return [self.embed_query(text) for text in texts]
+
+    def _embed_input(self, input: Union[str, List[str]]) -> List[List[float]]:
+        """
+        Internal method to handle the API call for embedding inputs.
+
+        Args:
+            input (Union[str, List[str]]): Either a single text string or a list of text strings to embed.
+
+        Returns:
+            List[List[float]]: A list of embedding vectors for the input(s).
+
+        Raises:
+            HTTPError: If the API request fails.
+        """
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        payload = {"model": self.model, "input": input, "encoding_format": "float"}
+        response = requests.request("POST", NOVITA_EMBEDDING_API, json=payload, headers=headers)
+        response.raise_for_status()
+        result = response.json()["data"]
+        sorted_results = sorted(result, key=lambda x: x["index"])
+        return [res["embedding"] for res in sorted_results]
+
+    @property
+    def dimension(self) -> int:
+        """
+        Get the dimensionality of the embeddings for the current model.
+
+        Returns:
+            int: The number of dimensions in the embedding vectors.
+        """
+        return NOVITA_MODEL_DIM_MAP[self.model]

deepsearcher/llm/__init__.py

@@ -4,6 +4,7 @@
 from .deepseek import DeepSeek
 from .gemini import Gemini
 from .glm import GLM
+from .novita import Novita
 from .ollama import Ollama
 from .openai_llm import OpenAI
 from .ppio import PPIO
@@ -26,4 +27,5 @@
     "Volcengine",
     "GLM",
     "Bedrock",
+    "Novita",
 ]

deepsearcher/llm/novita.py

@@ -0,0 +1,34 @@
+import os
+from typing import Dict, List
+
+from deepsearcher.llm.base import BaseLLM, ChatResponse
+
+
+class Novita(BaseLLM):
+    """
+    Novita AI API
+    """
+
+    def __init__(self, model: str = "qwen/qwq-32b", **kwargs):
+        from openai import OpenAI as OpenAI_
+
+        self.model = model
+        if "api_key" in kwargs:
+            api_key = kwargs.pop("api_key")
+        else:
+            api_key = os.getenv("NOVITA_API_KEY")
+        if "base_url" in kwargs:
+            base_url = kwargs.pop("base_url")
+        else:
+            base_url = "https://api.novita.ai/v3/openai"
+        self.client = OpenAI_(api_key=api_key, base_url=base_url, **kwargs)
+
+    def chat(self, messages: List[Dict]) -> ChatResponse:
+        completion = self.client.chat.completions.create(
+            model=self.model,
+            messages=messages,
+        )
+        return ChatResponse(
+            content=completion.choices[0].message.content,
+            total_tokens=completion.usage.total_tokens,
+        )

evaluation/eval_config.yaml

@@ -39,6 +39,12 @@ provide_settings:
 #    provider: "Ollama"
 #    config:
 #      model: "qwq"
+##      base_url: ""
+
+#    provider: "Novita"
+#    config:
+#      model: "deepseek/deepseek-v3-0324"
+##      api_key: "xxxx"  # Uncomment to override the `NOVITA_API_KEY` set in the environment variable
 ##      base_url: ""
 
   embedding:
@@ -68,6 +74,11 @@ provide_settings:
 #      model: "BAAI/bge-m3"
 # .    api_key: ""   # Uncomment to override the `SILICONFLOW_API_KEY` set in the environment variable   
 
+#    provider: "NovitaEmbedding"
+#    config:
+#      model: "baai/bge-m3"
+# .    api_key: ""   # Uncomment to override the `NOVITA_API_KEY` set in the environment variable   
+
   file_loader:
 #    provider: "PDFLoader"
 #    config: {}