langchain/tests/integration_tests/embeddings/test_deepinfra.py
Yessen Kanapin c66755b661
Add DeepInfra embeddings integration with tests and examples, better exception handling for Deep Infra LLM (#5854)
#### Who can review?

Tag maintainers/contributors who might be interested:
  @hwchase17 - project lead
  - @agola11

---------

Co-authored-by: Yessen Kanapin <yessen@deepinfra.com>
2023-06-07 19:14:30 -07:00

20 lines
626 B
Python

"""Test DeepInfra API wrapper."""
from langchain.embeddings import DeepInfraEmbeddings
def test_deepinfra_call() -> None:
"""Test valid call to DeepInfra."""
deepinfra_emb = DeepInfraEmbeddings(model_id="sentence-transformers/clip-ViT-B-32")
r1 = deepinfra_emb.embed_documents(
[
"Alpha is the first letter of Greek alphabet",
"Beta is the second letter of Greek alphabet",
]
)
assert len(r1) == 2
assert len(r1[0]) == 512
assert len(r1[1]) == 512
r2 = deepinfra_emb.embed_query("What is the third letter of Greek alphabet")
assert len(r2) == 512