From 56a7e17cbc96f0ee93d460374f3639343c3b9cef Mon Sep 17 00:00:00 2001 From: James Briggs Date: Thu, 10 Oct 2024 23:06:07 +0200 Subject: [PATCH] fix: move embeds type to dunder --- semantic_router/encoders/cohere.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/semantic_router/encoders/cohere.py b/semantic_router/encoders/cohere.py index 6be77845..006b6605 100644 --- a/semantic_router/encoders/cohere.py +++ b/semantic_router/encoders/cohere.py @@ -9,6 +9,7 @@ class CohereEncoder(BaseEncoder): _client: Any = PrivateAttr() + _embed_type: Any = PrivateAttr() type: str = "cohere" input_type: Optional[str] = "search_query" @@ -40,7 +41,7 @@ def _initialize_client(self, cohere_api_key: Optional[str] = None): try: import cohere from cohere.types.embed_response import EmbeddingsByTypeEmbedResponse - self.EmbeddingsByTypeEmbedResponse = EmbeddingsByTypeEmbedResponse + self._embed_type = EmbeddingsByTypeEmbedResponse except ImportError: raise ImportError( "Please install Cohere to use CohereEncoder. " @@ -66,7 +67,7 @@ def __call__(self, docs: List[str]) -> List[List[float]]: texts=docs, input_type=self.input_type, model=self.name ) # Check for unsupported type. - if isinstance(embeds, self.EmbeddingsByTypeEmbedResponse): + if isinstance(embeds, self._embed_type): raise NotImplementedError( "Handling of EmbedByTypeResponseEmbeddings is not implemented." )