diff --git a/memory_scope/models/base_embedding_model.py b/memory_scope/models/base_embedding_model.py index 6144e2a7..361aeaa3 100644 --- a/memory_scope/models/base_embedding_model.py +++ b/memory_scope/models/base_embedding_model.py @@ -50,4 +50,19 @@ class LLIEmbedding(BaseEmbeddingModel): return results + async def _async_call(self, **kwargs) -> ModelResponse: + """ + :param kwargs: + :return: + """ + results = ModelResponse() + try: + response = self.model.aget_text_embedding_batch(**self.data) + results.raw = response + results.status = True + except Exception as e: + results.details = e + results.status = False + return results + \ No newline at end of file diff --git a/tests/models/test_models_llm_embedding.py b/tests/models/test_models_llm_embedding.py index 9d08bee9..6c380ba6 100644 --- a/tests/models/test_models_llm_embedding.py +++ b/tests/models/test_models_llm_embedding.py @@ -17,8 +17,13 @@ class TestLLIEmbedding(unittest.TestCase): embs = self.emb.call(text=text) def test_batch_embedding(self): - text = ["您吃了吗?", + texts = ["您吃了吗?", "吃了吗您?"] - embs = self.emb.call(text=text) + embs = self.emb.call(text=texts) - \ No newline at end of file + async def test_async_embedding(self): + texts = ["您吃了吗?", + "吃了吗您?"] + # 调用异步函数并等待其结果 + embs = await self.emb.async_call(texts) +