test_jina_ai.py

importjson
importos
importsys
fromdatetimeimportdatetime
fromunittest.mockimportAsyncMock

sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path


frombase_rerank_unit_testsimportBaseLLMRerankTest
importlitellm


classTestJinaAI(BaseLLMRerankTest):
defget_custom_llm_provider(self) ->litellm.LlmProviders:
returnlitellm.LlmProviders.JINA_AI

defget_base_rerank_call_args(self) ->dict:
return {
"model": "jina_ai/jina-reranker-v2-base-multilingual",
 }


deftest_jina_ai_embedding():
litellm.embedding(
model="jina_ai/jina-embeddings-v3",
input=["a"],
task="separation",
dimensions=1024,
 )