feat: support ark embedding (#449)

* fix: vision embedding

* feat: openai like embedding to ark embedding factory

* feat: header and consts

* fix: comment

Author: zakahan

config.yaml.full

@@ -15,8 +15,8 @@ model:
     api_key: 
   # [optional] for knowledgebase
   embedding:
-    name: doubao-embedding-text-240715
-    dim: 2560
+    name: doubao-embedding-vision-250615
+    dim: 2048
     api_base: https://ark.cn-beijing.volces.com/api/v3/
     api_key:
   video:

veadk/configs/model_configs.py

@@ -23,6 +23,8 @@
     DEFAULT_MODEL_AGENT_API_BASE,
     DEFAULT_MODEL_AGENT_NAME,
     DEFAULT_MODEL_AGENT_PROVIDER,
+    DEFAULT_MODEL_EMBEDDING_NAME,
+    DEFAULT_MODEL_EMBEDDING_DIM,
 )
 
 
@@ -46,13 +48,13 @@ def api_key(self) -> str:
 class EmbeddingModelConfig(BaseSettings):
     model_config = SettingsConfigDict(env_prefix="MODEL_EMBEDDING_")
 
-    name: str = "doubao-embedding-text-240715"
+    name: str = DEFAULT_MODEL_EMBEDDING_NAME
     """Model name for embedding."""
 
-    dim: int = 2560
+    dim: int = DEFAULT_MODEL_EMBEDDING_DIM
     """Embedding dim is different from different models."""
 
-    api_base: str = "https://ark.cn-beijing.volces.com/api/v3/"
+    api_base: str = DEFAULT_MODEL_AGENT_API_BASE
     """The api base of the model for embedding."""
 
     @cached_property

veadk/consts.py

@@ -75,3 +75,7 @@
 
 DEFAULT_NACOS_GROUP = "VEADK_GROUP"
 DEFAULT_NACOS_INSTANCE_NAME = "veadk"
+
+DEFAULT_MODEL_EMBEDDING_NAME = "doubao-embedding-vision-250615"
+DEFAULT_MODEL_EMBEDDING_API_BASE = "https://ark.cn-beijing.volces.com/api/v3/"
+DEFAULT_MODEL_EMBEDDING_DIM = 2048

veadk/knowledgebase/backends/in_memory_backend.py

@@ -14,13 +14,13 @@
 
 from llama_index.core import Document, SimpleDirectoryReader, VectorStoreIndex
 from llama_index.core.schema import BaseNode
-from llama_index.embeddings.openai_like import OpenAILikeEmbedding
 from pydantic import Field
 from typing_extensions import Any, override
 
 from veadk.configs.model_configs import EmbeddingModelConfig, NormalEmbeddingModelConfig
 from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
 from veadk.knowledgebase.backends.utils import get_llama_index_splitter
+from veadk.models.ark_embedding import create_embedding_model
 
 
 class InMemoryKnowledgeBackend(BaseKnowledgebaseBackend):
@@ -39,7 +39,7 @@ class InMemoryKnowledgeBackend(BaseKnowledgebaseBackend):
     )
 
     def model_post_init(self, __context: Any) -> None:
-        self._embed_model = OpenAILikeEmbedding(
+        self._embed_model = create_embedding_model(
             model_name=self.embedding_config.name,
             api_key=self.embedding_config.api_key,
             api_base=self.embedding_config.api_base,

veadk/knowledgebase/backends/opensearch_backend.py

@@ -21,7 +21,6 @@
     VectorStoreIndex,
 )
 from llama_index.core.schema import BaseNode
-from llama_index.embeddings.openai_like import OpenAILikeEmbedding
 from pydantic import Field
 from typing_extensions import Any, override
 
@@ -33,6 +32,7 @@
 )
 from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
 from veadk.knowledgebase.backends.utils import get_llama_index_splitter
+from veadk.models.ark_embedding import create_embedding_model
 from veadk.utils.logger import get_logger
 
 try:
@@ -112,7 +112,7 @@ def model_post_init(self, __context: Any) -> None:
             vector_store=self._vector_store
         )
 
-        self._embed_model = OpenAILikeEmbedding(
+        self._embed_model = create_embedding_model(
             model_name=self.embedding_config.name,
             api_key=self.embedding_config.api_key,
             api_base=self.embedding_config.api_base,

veadk/knowledgebase/backends/redis_backend.py

@@ -19,7 +19,6 @@
     VectorStoreIndex,
 )
 from llama_index.core.schema import BaseNode
-from llama_index.embeddings.openai_like import OpenAILikeEmbedding
 from pydantic import Field
 from typing_extensions import Any, override
 
@@ -28,6 +27,7 @@
 from veadk.configs.model_configs import EmbeddingModelConfig, NormalEmbeddingModelConfig
 from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
 from veadk.knowledgebase.backends.utils import get_llama_index_splitter
+from veadk.models.ark_embedding import create_embedding_model
 
 try:
     from llama_index.vector_stores.redis import RedisVectorStore
@@ -92,7 +92,7 @@ def model_post_init(self, __context: Any) -> None:
             password=self.redis_config.password,
         )
 
-        self._embed_model = OpenAILikeEmbedding(
+        self._embed_model = create_embedding_model(
             model_name=self.embedding_config.name,
             api_key=self.embedding_config.api_key,
             api_base=self.embedding_config.api_base,

veadk/knowledgebase/backends/tos_vector_backend.py

@@ -20,7 +20,6 @@
     SimpleDirectoryReader,
 )
 from llama_index.core.schema import BaseNode
-from llama_index.embeddings.openai_like import OpenAILikeEmbedding
 from pydantic import Field
 from typing_extensions import Any, override
 
@@ -30,6 +29,7 @@
 
 from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend
 from veadk.knowledgebase.backends.utils import get_llama_index_splitter
+from veadk.models.ark_embedding import create_embedding_model
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
@@ -77,7 +77,7 @@ def model_post_init(self, __context: Any) -> None:
         # create_bucket and index if not exist
         self._create_index()
 
-        self._embed_model = OpenAILikeEmbedding(
+        self._embed_model = create_embedding_model(
             model_name=self.embedding_config.name,
             api_key=self.embedding_config.api_key,
             api_base=self.embedding_config.api_base,

veadk/memory/long_term_memory_backends/in_memory_backend.py

@@ -14,7 +14,6 @@
 
 from llama_index.core import Document, VectorStoreIndex
 from llama_index.core.schema import BaseNode
-from llama_index.embeddings.openai_like import OpenAILikeEmbedding
 from pydantic import Field
 from typing_extensions import Any, override
 
@@ -23,14 +22,15 @@
 from veadk.memory.long_term_memory_backends.base_backend import (
     BaseLongTermMemoryBackend,
 )
+from veadk.models.ark_embedding import create_embedding_model
 
 
 class InMemoryLTMBackend(BaseLongTermMemoryBackend):
     embedding_config: EmbeddingModelConfig = Field(default_factory=EmbeddingModelConfig)
     """Embedding model configs"""
 
     def model_post_init(self, __context: Any) -> None:
-        self._embed_model = OpenAILikeEmbedding(
+        self._embed_model = create_embedding_model(
             model_name=self.embedding_config.name,
             api_key=self.embedding_config.api_key,
             api_base=self.embedding_config.api_base,

veadk/memory/long_term_memory_backends/opensearch_backend.py

@@ -16,7 +16,6 @@
 
 from llama_index.core import Document, VectorStoreIndex
 from llama_index.core.schema import BaseNode
-from llama_index.embeddings.openai_like import OpenAILikeEmbedding
 from pydantic import Field
 from typing_extensions import Any, override
 
@@ -30,6 +29,7 @@
 from veadk.memory.long_term_memory_backends.base_backend import (
     BaseLongTermMemoryBackend,
 )
+from veadk.models.ark_embedding import create_embedding_model
 from veadk.utils.logger import get_logger
 
 try:
@@ -55,7 +55,7 @@ class OpensearchLTMBackend(BaseLongTermMemoryBackend):
     """Embedding model configs"""
 
     def model_post_init(self, __context: Any) -> None:
-        self._embed_model = OpenAILikeEmbedding(
+        self._embed_model = create_embedding_model(
             model_name=self.embedding_config.name,
             api_key=self.embedding_config.api_key,
             api_base=self.embedding_config.api_base,

veadk/memory/long_term_memory_backends/redis_backend.py

@@ -14,7 +14,6 @@
 
 from llama_index.core import Document, VectorStoreIndex
 from llama_index.core.schema import BaseNode
-from llama_index.embeddings.openai_like import OpenAILikeEmbedding
 from pydantic import Field
 from typing_extensions import Any, override
 
@@ -25,6 +24,7 @@
 from veadk.memory.long_term_memory_backends.base_backend import (
     BaseLongTermMemoryBackend,
 )
+from veadk.models.ark_embedding import create_embedding_model
 from veadk.utils.logger import get_logger
 
 try:
@@ -51,7 +51,7 @@ class RedisLTMBackend(BaseLongTermMemoryBackend):
     """Embedding model configs"""
 
     def model_post_init(self, __context: Any) -> None:
-        self._embed_model = OpenAILikeEmbedding(
+        self._embed_model = create_embedding_model(
             model_name=self.embedding_config.name,
             api_key=self.embedding_config.api_key,
             api_base=self.embedding_config.api_base,