mirror of
https://github.com/hwchase17/langchain
synced 2024-11-10 01:10:59 +00:00
682d21c3de
Thank you for contributing to LangChain!
- [x] **PR title**: "langchain-ibm: Add support for ibm-watsonx-ai new
major version"
- [x] **PR message**:
- **Description:** Add support for ibm-watsonx-ai new major version
- **Dependencies:** `ibm_watsonx_ai`
- [x] **Add tests and docs**: If you're adding a new integration, please
include
1. a test for the integration, preferably unit tests that do not rely on
network access,
2. an example notebook showing its use. It lives in
`docs/docs/integrations` directory.
- [x] **Lint and test**: Run `make format`, `make lint` and `make test`
from the root of the package(s) you've modified. See contribution
guidelines for more: https://python.langchain.com/docs/contributing/
Co-authored-by: Erick Friis <erick@langchain.dev>
69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
"""Test WatsonxEmbeddings.
|
|
|
|
You'll need to set WATSONX_APIKEY and WATSONX_PROJECT_ID environment variables.
|
|
"""
|
|
|
|
import os
|
|
|
|
from ibm_watsonx_ai import APIClient # type: ignore
|
|
from ibm_watsonx_ai.metanames import EmbedTextParamsMetaNames # type: ignore
|
|
|
|
from langchain_ibm import WatsonxEmbeddings
|
|
|
|
WX_APIKEY = os.environ.get("WATSONX_APIKEY", "")
|
|
WX_PROJECT_ID = os.environ.get("WATSONX_PROJECT_ID", "")
|
|
|
|
URL = "https://us-south.ml.cloud.ibm.com"
|
|
MODEL_ID = "ibm/slate-125m-english-rtrvr"
|
|
|
|
DOCUMENTS = ["What is a generative ai?", "What is a loan and how does it works?"]
|
|
|
|
|
|
def test_01_generate_embed_documents() -> None:
|
|
watsonx_embedding = WatsonxEmbeddings(
|
|
model_id=MODEL_ID, url=URL, project_id=WX_PROJECT_ID
|
|
)
|
|
generate_embedding = watsonx_embedding.embed_documents(texts=DOCUMENTS)
|
|
assert len(generate_embedding) == len(DOCUMENTS)
|
|
assert all(isinstance(el, float) for el in generate_embedding[0])
|
|
|
|
|
|
def test_02_generate_embed_query() -> None:
|
|
watsonx_embedding = WatsonxEmbeddings(
|
|
model_id=MODEL_ID,
|
|
url=URL,
|
|
project_id=WX_PROJECT_ID,
|
|
)
|
|
generate_embedding = watsonx_embedding.embed_query(text=DOCUMENTS[0])
|
|
assert isinstance(generate_embedding, list) and isinstance(
|
|
generate_embedding[0], float
|
|
)
|
|
|
|
|
|
def test_03_generate_embed_documents_with_param() -> None:
|
|
embed_params = {
|
|
EmbedTextParamsMetaNames.TRUNCATE_INPUT_TOKENS: 3,
|
|
}
|
|
watsonx_embedding = WatsonxEmbeddings(
|
|
model_id=MODEL_ID, url=URL, project_id=WX_PROJECT_ID, params=embed_params
|
|
)
|
|
generate_embedding = watsonx_embedding.embed_documents(texts=DOCUMENTS)
|
|
assert len(generate_embedding) == len(DOCUMENTS)
|
|
assert all(isinstance(el, float) for el in generate_embedding[0])
|
|
|
|
|
|
def test_10_generate_embed_query_with_client_initialization() -> None:
|
|
watsonx_client = APIClient(
|
|
credentials={
|
|
"url": URL,
|
|
"apikey": WX_APIKEY,
|
|
}
|
|
)
|
|
watsonx_embedding = WatsonxEmbeddings(
|
|
model_id=MODEL_ID, project_id=WX_PROJECT_ID, watsonx_client=watsonx_client
|
|
)
|
|
generate_embedding = watsonx_embedding.embed_query(text=DOCUMENTS[0])
|
|
assert isinstance(generate_embedding, list) and isinstance(
|
|
generate_embedding[0], float
|
|
)
|