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: {}