From 20023ed0a940279fc8c48f6aa79d61dcdff8986e Mon Sep 17 00:00:00 2001 From: Matthew Farrellee Date: Tue, 17 Dec 2024 09:56:03 -0500 Subject: [PATCH] fix format & lint --- .../embeddings.py | 8 ++--- .../integration_tests/test_embeddings.py | 32 ++++++++++++++----- 2 files changed, 28 insertions(+), 12 deletions(-) diff --git a/libs/ai-endpoints/langchain_nvidia_ai_endpoints/embeddings.py b/libs/ai-endpoints/langchain_nvidia_ai_endpoints/embeddings.py index bfb46ce6..466570ef 100644 --- a/libs/ai-endpoints/langchain_nvidia_ai_endpoints/embeddings.py +++ b/libs/ai-endpoints/langchain_nvidia_ai_endpoints/embeddings.py @@ -1,4 +1,4 @@ -from typing import Any, List, Literal, Optional +from typing import Any, Dict, List, Literal, Optional from langchain_core.embeddings import Embeddings from langchain_core.outputs.llm_result import LLMResult @@ -74,8 +74,8 @@ def __init__(self, **kwargs: Any): trucate (str): "NONE", "START", "END", truncate input text if it exceeds the model's context length. Default is "NONE", which raises an error if an input is too long. - dimensions (int): The number of dimensions for the embeddings. This parameter - is not supported by all models. + dimensions (int): The number of dimensions for the embeddings. This + parameter is not supported by all models. API Key: - The recommended way to provide the API key is through the `NVIDIA_API_KEY` @@ -135,7 +135,7 @@ def _embed( # truncate: "NONE" | "START" | "END" -- default "NONE", error raised if # an input is too long # dimensions: int -- not supported by all models - payload = { + payload: Dict[str, Any] = { "input": texts, "model": self.model, "encoding_format": "float", diff --git a/libs/ai-endpoints/tests/integration_tests/test_embeddings.py b/libs/ai-endpoints/tests/integration_tests/test_embeddings.py index 841767f4..b2cf374d 100644 --- a/libs/ai-endpoints/tests/integration_tests/test_embeddings.py +++ b/libs/ai-endpoints/tests/integration_tests/test_embeddings.py @@ -98,7 +98,9 @@ def test_embed_documents_truncate( @pytest.mark.parametrize("dimensions", [32, 64, 128, 2048]) -def test_embed_query_with_dimensions(embedding_model: str, mode: dict, dimensions: int) -> None: +def test_embed_query_with_dimensions( + embedding_model: str, mode: dict, dimensions: int +) -> None: if embedding_model != "nvidia/llama-3.2-nv-embedqa-1b-v2": pytest.skip("Model does not support custom dimensions.") query = "foo bar" @@ -107,7 +109,9 @@ def test_embed_query_with_dimensions(embedding_model: str, mode: dict, dimension @pytest.mark.parametrize("dimensions", [32, 64, 128, 2048]) -def test_embed_documents_with_dimensions(embedding_model: str, mode: dict, dimensions: int) -> None: +def test_embed_documents_with_dimensions( + embedding_model: str, mode: dict, dimensions: int +) -> None: if embedding_model != "nvidia/llama-3.2-nv-embedqa-1b-v2": pytest.skip("Model does not support custom dimensions.") documents = ["foo bar", "bar foo"] @@ -118,7 +122,9 @@ def test_embed_documents_with_dimensions(embedding_model: str, mode: dict, dimen @pytest.mark.parametrize("dimensions", [102400]) -def test_embed_query_with_large_dimensions(embedding_model: str, mode: dict, dimensions: int) -> None: +def test_embed_query_with_large_dimensions( + embedding_model: str, mode: dict, dimensions: int +) -> None: if embedding_model != "nvidia/llama-3.2-nv-embedqa-1b-v2": pytest.skip("Model does not support custom dimensions.") query = "foo bar" @@ -127,7 +133,9 @@ def test_embed_query_with_large_dimensions(embedding_model: str, mode: dict, dim @pytest.mark.parametrize("dimensions", [102400]) -def test_embed_documents_with_large_dimensions(embedding_model: str, mode: dict, dimensions: int) -> None: +def test_embed_documents_with_large_dimensions( + embedding_model: str, mode: dict, dimensions: int +) -> None: if embedding_model != "nvidia/llama-3.2-nv-embedqa-1b-v2": pytest.skip("Model does not support custom dimensions.") documents = ["foo bar", "bar foo"] @@ -138,22 +146,30 @@ def test_embed_documents_with_large_dimensions(embedding_model: str, mode: dict, @pytest.mark.parametrize("dimensions", [-1]) -def test_embed_query_invalid_dimensions(embedding_model: str, mode: dict, dimensions: int) -> None: +def test_embed_query_invalid_dimensions( + embedding_model: str, mode: dict, dimensions: int +) -> None: if embedding_model != "nvidia/llama-3.2-nv-embedqa-1b-v2": pytest.skip("Model does not support custom dimensions.") query = "foo bar" with pytest.raises(Exception) as exc: - NVIDIAEmbeddings(model=embedding_model, dimensions=dimensions, **mode).embed_query(query) + NVIDIAEmbeddings( + model=embedding_model, dimensions=dimensions, **mode + ).embed_query(query) assert "400" in str(exc.value) @pytest.mark.parametrize("dimensions", [-1]) -def test_embed_documents_invalid_dimensions(embedding_model: str, mode: dict, dimensions: int) -> None: +def test_embed_documents_invalid_dimensions( + embedding_model: str, mode: dict, dimensions: int +) -> None: if embedding_model != "nvidia/llama-3.2-nv-embedqa-1b-v2": pytest.skip("Model does not support custom dimensions.") documents = ["foo bar", "bar foo"] with pytest.raises(Exception) as exc: - NVIDIAEmbeddings(model=embedding_model, dimensions=dimensions, **mode).embed_documents(documents) + NVIDIAEmbeddings( + model=embedding_model, dimensions=dimensions, **mode + ).embed_documents(documents) assert "400" in str(exc.value)