langchain/tests/integration_tests/embeddings/test_octoai_embeddings.py
Harrison Chase 3ac08c3de4
Harrison/octo ml (#6897)
Co-authored-by: Bassem Yacoube <125713079+AI-Bassem@users.noreply.github.com>
Co-authored-by: Shotaro Kohama <khmshtr28@gmail.com>
Co-authored-by: Rian Dolphin <34861538+rian-dolphin@users.noreply.github.com>
Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>
Co-authored-by: Shashank Deshpande <shashankdeshpande18@gmail.com>
2023-06-28 23:04:11 -07:00

35 lines
1023 B
Python

"""Test octoai embeddings."""
from langchain.embeddings.octoai_embeddings import (
OctoAIEmbeddings,
)
def test_octoai_embedding_documents() -> None:
"""Test octoai embeddings."""
documents = ["foo bar"]
embedding = OctoAIEmbeddings(
endpoint_url="<endpoint_url>",
octoai_api_token="<octoai_api_token>",
embed_instruction="Represent this input: ",
query_instruction="Represent this input: ",
model_kwargs=None,
)
output = embedding.embed_documents(documents)
assert len(output) == 1
assert len(output[0]) == 768
def test_octoai_embedding_query() -> None:
"""Test octoai embeddings."""
document = "foo bar"
embedding = OctoAIEmbeddings(
endpoint_url="<endpoint_url>",
octoai_api_token="<octoai_api_token>",
embed_instruction="Represent this input: ",
query_instruction="Represent this input: ",
model_kwargs=None,
)
output = embedding.embed_query(document)
assert len(output) == 768