my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

my-project-public — gitlawb

#1	import os
#2	from typing import Optional, Union
#3
#4	from chromadb import EmbeddingFunction, Embeddings
#5
#6	from embedchain.config import BaseEmbedderConfig
#7	from embedchain.embedder.base import BaseEmbedder
#8	from embedchain.models import VectorDimensions
#9
#10
#11	class MistralAIEmbeddingFunction(EmbeddingFunction):
#12	def __init__(self, config: BaseEmbedderConfig) -> None:
#13	super().__init__()
#14	try:
#15	from langchain_mistralai import MistralAIEmbeddings
#16	except ModuleNotFoundError:
#17	raise ModuleNotFoundError(
#18	"The required dependencies for MistralAI are not installed."
#19	'Please install with `pip install --upgrade "embedchain[mistralai]"`'
#20	) from None
#21	self.config = config
#22	api_key = self.config.api_key or os.getenv("MISTRAL_API_KEY")
#23	self.client = MistralAIEmbeddings(mistral_api_key=api_key)
#24	self.client.model = self.config.model
#25
#26	def __call__(self, input: Union[list[str], str]) -> Embeddings:
#27	if isinstance(input, str):
#28	input_ = [input]
#29	else:
#30	input_ = input
#31	response = self.client.embed_documents(input_)
#32	return response
#33
#34
#35	class MistralAIEmbedder(BaseEmbedder):
#36	def __init__(self, config: Optional[BaseEmbedderConfig] = None):
#37	super().__init__(config)
#38
#39	if self.config.model is None:
#40	self.config.model = "mistral-embed"
#41
#42	embedding_fn = MistralAIEmbeddingFunction(config=self.config)
#43	self.set_embedding_fn(embedding_fn=embedding_fn)
#44
#45	vector_dimension = self.config.vector_dimension or VectorDimensions.MISTRAL_AI.value
#46	self.set_vector_dimension(vector_dimension=vector_dimension)
#47

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public