codekingpro/portable-devtools

1from typing import Dict, Any, Type, Set
2from chromadb.api.types import (
3    EmbeddingFunction,
4    DefaultEmbeddingFunction,
5    SparseEmbeddingFunction,
6)
7
8# Import all embedding functions
9from chromadb.utils.embedding_functions.cohere_embedding_function import (
10    CohereEmbeddingFunction,
11)
12from chromadb.utils.embedding_functions.openai_embedding_function import (
13    OpenAIEmbeddingFunction,
14)
15from chromadb.utils.embedding_functions.huggingface_embedding_function import (
16    HuggingFaceEmbeddingFunction,
17    HuggingFaceEmbeddingServer,
18)
19from chromadb.utils.embedding_functions.sentence_transformer_embedding_function import (
20    SentenceTransformerEmbeddingFunction,
21)
22from chromadb.utils.embedding_functions.google_embedding_function import (
23    GooglePalmEmbeddingFunction,
24    GoogleGenerativeAiEmbeddingFunction,
25    GoogleVertexEmbeddingFunction,
26    GoogleGeminiEmbeddingFunction,
27    GoogleGenaiEmbeddingFunction,  # Backward compatibility alias
28)
29from chromadb.utils.embedding_functions.ollama_embedding_function import (
30    OllamaEmbeddingFunction,
31)
32from chromadb.utils.embedding_functions.instructor_embedding_function import (
33    InstructorEmbeddingFunction,
34)
35from chromadb.utils.embedding_functions.jina_embedding_function import (
36    JinaEmbeddingFunction,
37    JinaQueryConfig,
38)
39from chromadb.utils.embedding_functions.voyageai_embedding_function import (
40    VoyageAIEmbeddingFunction,
41)
42from chromadb.utils.embedding_functions.onnx_mini_lm_l6_v2 import ONNXMiniLM_L6_V2
43from chromadb.utils.embedding_functions.open_clip_embedding_function import (
44    OpenCLIPEmbeddingFunction,
45)
46from chromadb.utils.embedding_functions.roboflow_embedding_function import (
47    RoboflowEmbeddingFunction,
48)
49from chromadb.utils.embedding_functions.text2vec_embedding_function import (
50    Text2VecEmbeddingFunction,
51)
52from chromadb.utils.embedding_functions.amazon_bedrock_embedding_function import (
53    AmazonBedrockEmbeddingFunction,
54)
55from chromadb.utils.embedding_functions.chroma_langchain_embedding_function import (
56    ChromaLangchainEmbeddingFunction,
57)
58from chromadb.utils.embedding_functions.baseten_embedding_function import (
59    BasetenEmbeddingFunction,
60)
61from chromadb.utils.embedding_functions.cloudflare_workers_ai_embedding_function import (
62    CloudflareWorkersAIEmbeddingFunction,
63)
64from chromadb.utils.embedding_functions.together_ai_embedding_function import (
65    TogetherAIEmbeddingFunction,
66)
67from chromadb.utils.embedding_functions.mistral_embedding_function import (
68    MistralEmbeddingFunction,
69)
70from chromadb.utils.embedding_functions.morph_embedding_function import (
71    MorphEmbeddingFunction,
72)
73from chromadb.utils.embedding_functions.nomic_embedding_function import (
74    NomicEmbeddingFunction,
75    NomicQueryConfig,
76)
77from chromadb.utils.embedding_functions.huggingface_sparse_embedding_function import (
78    HuggingFaceSparseEmbeddingFunction,
79)
80from chromadb.utils.embedding_functions.fastembed_sparse_embedding_function import (
81    FastembedSparseEmbeddingFunction,
82)
83from chromadb.utils.embedding_functions.bm25_embedding_function import (
84    Bm25EmbeddingFunction,
85)
86from chromadb.utils.embedding_functions.chroma_cloud_qwen_embedding_function import (
87    ChromaCloudQwenEmbeddingFunction,
88)
89from chromadb.utils.embedding_functions.chroma_cloud_splade_embedding_function import (
90    ChromaCloudSpladeEmbeddingFunction,
91)
92from chromadb.utils.embedding_functions.chroma_bm25_embedding_function import (
93    ChromaBm25EmbeddingFunction,
94)
95from chromadb.utils.embedding_functions.perplexity_embedding_function import (
96    PerplexityEmbeddingFunction,
97)
98
99
100# Get all the class names for backward compatibility
101_all_classes: Set[str] = {
102    "CohereEmbeddingFunction",
103    "OpenAIEmbeddingFunction",
104    "HuggingFaceEmbeddingFunction",
105    "HuggingFaceEmbeddingServer",
106    "SentenceTransformerEmbeddingFunction",
107    "GooglePalmEmbeddingFunction",
108    "GoogleGenerativeAiEmbeddingFunction",
109    "GoogleVertexEmbeddingFunction",
110    "GoogleGeminiEmbeddingFunction",
111    "GoogleGenaiEmbeddingFunction",  # Backward compatibility alias
112    "OllamaEmbeddingFunction",
113    "InstructorEmbeddingFunction",
114    "JinaEmbeddingFunction",
115    "MistralEmbeddingFunction",
116    "MorphEmbeddingFunction",
117    "NomicEmbeddingFunction",
118    "VoyageAIEmbeddingFunction",
119    "ONNXMiniLM_L6_V2",
120    "OpenCLIPEmbeddingFunction",
121    "RoboflowEmbeddingFunction",
122    "Text2VecEmbeddingFunction",
123    "AmazonBedrockEmbeddingFunction",
124    "ChromaLangchainEmbeddingFunction",
125    "BasetenEmbeddingFunction",
126    "CloudflareWorkersAIEmbeddingFunction",
127    "TogetherAIEmbeddingFunction",
128    "DefaultEmbeddingFunction",
129    "HuggingFaceSparseEmbeddingFunction",
130    "FastembedSparseEmbeddingFunction",
131    "Bm25EmbeddingFunction",
132    "ChromaCloudQwenEmbeddingFunction",
133    "ChromaCloudSpladeEmbeddingFunction",
134    "ChromaBm25EmbeddingFunction",
135    "PerplexityEmbeddingFunction"
136}
137
138
139def get_builtins() -> Set[str]:
140    return _all_classes
141
142
143# Dictionary of supported embedding functions
144known_embedding_functions: Dict[str, Type[EmbeddingFunction]] = {  # type: ignore
145    "cohere": CohereEmbeddingFunction,
146    "openai": OpenAIEmbeddingFunction,
147    "huggingface": HuggingFaceEmbeddingFunction,
148    "huggingface_server": HuggingFaceEmbeddingServer,
149    "sentence_transformer": SentenceTransformerEmbeddingFunction,
150    "google_palm": GooglePalmEmbeddingFunction,
151    "google_generative_ai": GoogleGenerativeAiEmbeddingFunction,
152    "google_vertex": GoogleVertexEmbeddingFunction,
153    "google_gemini": GoogleGeminiEmbeddingFunction,
154    "google_genai": GoogleGeminiEmbeddingFunction,  # Backward compatibility alias
155    "ollama": OllamaEmbeddingFunction,
156    "instructor": InstructorEmbeddingFunction,
157    "jina": JinaEmbeddingFunction,
158    "mistral": MistralEmbeddingFunction,
159    "morph": MorphEmbeddingFunction,
160    "nomic": NomicEmbeddingFunction,
161    "voyageai": VoyageAIEmbeddingFunction,
162    "onnx_mini_lm_l6_v2": ONNXMiniLM_L6_V2,
163    "open_clip": OpenCLIPEmbeddingFunction,
164    "roboflow": RoboflowEmbeddingFunction,
165    "text2vec": Text2VecEmbeddingFunction,
166    "amazon_bedrock": AmazonBedrockEmbeddingFunction,
167    "chroma_langchain": ChromaLangchainEmbeddingFunction,
168    "baseten": BasetenEmbeddingFunction,
169    "default": DefaultEmbeddingFunction,
170    "cloudflare_workers_ai": CloudflareWorkersAIEmbeddingFunction,
171    "together_ai": TogetherAIEmbeddingFunction,
172    "chroma-cloud-qwen": ChromaCloudQwenEmbeddingFunction,
173    "perplexity": PerplexityEmbeddingFunction,
174}
175
176sparse_known_embedding_functions: Dict[str, Type[SparseEmbeddingFunction]] = {  # type: ignore
177    "huggingface_sparse": HuggingFaceSparseEmbeddingFunction,
178    "fastembed_sparse": FastembedSparseEmbeddingFunction,
179    "bm25": Bm25EmbeddingFunction,
180    "chroma-cloud-splade": ChromaCloudSpladeEmbeddingFunction,
181    "chroma_bm25": ChromaBm25EmbeddingFunction,
182}
183
184
185def register_embedding_function(ef_class=None):  # type: ignore
186    """Register a custom embedding function.
187
188    Can be used as a decorator:
189        @register_embedding_function
190        class MyEmbedding(EmbeddingFunction):
191            @classmethod
192            def name(cls): return "my_embedding"
193
194    Or directly:
195        register_embedding_function(MyEmbedding)
196
197    Args:
198        ef_class: The embedding function class to register.
199    """
200
201    def _register(cls):  # type: ignore
202        try:
203            name = cls.name()
204            known_embedding_functions[name] = cls
205        except Exception as e:
206            raise ValueError(f"Failed to register embedding function: {e}")
207        return cls  # Return the class unchanged
208
209    # If called with a class, register it immediately
210    if ef_class is not None:
211        return _register(ef_class)  # type: ignore
212
213    # If called without arguments, return a decorator
214    return _register
215
216
217def register_sparse_embedding_function(ef_class=None):  # type: ignore
218    """Register a custom sparse embedding function.
219
220    Can be used as a decorator:
221        @register_sparse_embedding_function
222        class MySparseEmbeddingFunction(SparseEmbeddingFunction):
223            @classmethod
224            def name(cls): return "my_sparse_embedding"
225    """
226
227    def _register(cls):  # type: ignore
228        try:
229            name = cls.name()
230            sparse_known_embedding_functions[name] = cls
231        except Exception as e:
232            raise ValueError(f"Failed to register sparse embedding function: {e}")
233        return cls  # Return the class unchanged
234
235    if ef_class is not None:
236        return _register(ef_class)  # type: ignore
237
238    return _register
239
240
241# Function to convert config to embedding function
242def config_to_embedding_function(config: Dict[str, Any]) -> EmbeddingFunction:  # type: ignore
243    """Convert a config dictionary to an embedding function.
244
245    Args:
246        config: The config dictionary.
247
248    Returns:
249        The embedding function.
250    """
251    if "name" not in config:
252        raise ValueError("Config must contain a 'name' field.")
253
254    name = config["name"]
255    if name not in known_embedding_functions:
256        raise ValueError(f"Unsupported embedding function: {name}")
257
258    ef_config = config.get("config", {})
259
260    if known_embedding_functions[name] is None:
261        raise ValueError(f"Unsupported embedding function: {name}")
262
263    return known_embedding_functions[name].build_from_config(ef_config)
264
265
266__all__ = [
267    "EmbeddingFunction",
268    "DefaultEmbeddingFunction",
269    "CohereEmbeddingFunction",
270    "OpenAIEmbeddingFunction",
271    "BasetenEmbeddingFunction",
272    "CloudflareWorkersAIEmbeddingFunction",
273    "HuggingFaceEmbeddingFunction",
274    "HuggingFaceEmbeddingServer",
275    "SentenceTransformerEmbeddingFunction",
276    "GooglePalmEmbeddingFunction",
277    "GoogleGenerativeAiEmbeddingFunction",
278    "GoogleVertexEmbeddingFunction",
279    "GoogleGeminiEmbeddingFunction",
280    "GoogleGenaiEmbeddingFunction",  # Backward compatibility alias
281    "OllamaEmbeddingFunction",
282    "InstructorEmbeddingFunction",
283    "JinaEmbeddingFunction",
284    "JinaQueryConfig",
285    "MistralEmbeddingFunction",
286    "MorphEmbeddingFunction",
287    "NomicEmbeddingFunction",
288    "NomicQueryConfig",
289    "VoyageAIEmbeddingFunction",
290    "ONNXMiniLM_L6_V2",
291    "OpenCLIPEmbeddingFunction",
292    "RoboflowEmbeddingFunction",
293    "Text2VecEmbeddingFunction",
294    "AmazonBedrockEmbeddingFunction",
295    "ChromaLangchainEmbeddingFunction",
296    "TogetherAIEmbeddingFunction",
297    "HuggingFaceSparseEmbeddingFunction",
298    "FastembedSparseEmbeddingFunction",
299    "Bm25EmbeddingFunction",
300    "ChromaCloudQwenEmbeddingFunction",
301    "ChromaCloudSpladeEmbeddingFunction",
302    "ChromaBm25EmbeddingFunction",
303    "PerplexityEmbeddingFunction",
304    "register_embedding_function",
305    "config_to_embedding_function",
306    "known_embedding_functions",
307]
308