diff --git a/README.md b/README.md
index 48bb0d18..64b6b29b 100644
--- a/README.md
+++ b/README.md
@@ -90,7 +90,7 @@ result = query("Write a report about xxx.") # Your question here
#### LLM Configuration
config.set_provider_config("llm", "(LLMName)", "(Arguments dict)")
-The "LLMName" can be one of the following: ["DeepSeek", "OpenAI", "XAI", "SiliconFlow", "PPIO", "TogetherAI", "Gemini", "Ollama"]
+The "LLMName" can be one of the following: ["DeepSeek", "OpenAI", "XAI", "SiliconFlow", "PPIO", "TogetherAI", "Gemini", "Ollama", "Novita"]
The "Arguments dict" is a dictionary that contains the necessary arguments for the LLM class.
@@ -187,7 +187,7 @@ result = query("Write a report about xxx.") # Your question here
#### Embedding Model Configuration
config.set_provider_config("embedding", "(EmbeddingModelName)", "(Arguments dict)")
-The "EmbeddingModelName" can be one of the following: ["MilvusEmbedding", "OpenAIEmbedding", "VoyageEmbedding", "SiliconflowEmbedding", "PPIOEmbedding"]
+The "EmbeddingModelName" can be one of the following: ["MilvusEmbedding", "OpenAIEmbedding", "VoyageEmbedding", "SiliconflowEmbedding", "PPIOEmbedding", "NovitaEmbedding"]
The "Arguments dict" is a dictionary that contains the necessary arguments for the embedding model class.
@@ -221,6 +221,13 @@ result = query("Write a report about xxx.") # Your question here
You need to install boto3 before running, execute: pip install boto3. More details about Amazon Bedrock: https://docs.aws.amazon.com/bedrock/
+
+ Example (Novita AI embedding)
+ Make sure you have prepared your Novita AI API KEY as an env variable NOVITA_API_KEY.
+ config.set_provider_config("embedding", "NovitaEmbedding", {"model": "baai/bge-m3"})
+ 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
+
+
Example (Siliconflow embedding)
Make sure you have prepared your Siliconflow API KEY as an env variable SILICONFLOW_API_KEY.
@@ -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
diff --git a/deepsearcher/config.yaml b/deepsearcher/config.yaml
index 9a6539f2..bfd06320 100644
--- a/deepsearcher/config.yaml
+++ b/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"
diff --git a/deepsearcher/embedding/__init__.py b/deepsearcher/embedding/__init__.py
index 1fb960f3..faf5695b 100644
--- a/deepsearcher/embedding/__init__.py
+++ b/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",
]
diff --git a/deepsearcher/embedding/novita_embedding.py b/deepsearcher/embedding/novita_embedding.py
new file mode 100644
index 00000000..b54de16b
--- /dev/null
+++ b/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]
diff --git a/deepsearcher/llm/__init__.py b/deepsearcher/llm/__init__.py
index 62926527..fc0b055b 100644
--- a/deepsearcher/llm/__init__.py
+++ b/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",
]
diff --git a/deepsearcher/llm/novita.py b/deepsearcher/llm/novita.py
new file mode 100644
index 00000000..dd4b08c5
--- /dev/null
+++ b/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,
+ )
diff --git a/evaluation/eval_config.yaml b/evaluation/eval_config.yaml
index 2b17c50f..c2537406 100644
--- a/evaluation/eval_config.yaml
+++ b/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: {}