From 94f71cf735d525a48f660259f1170e09044332a6 Mon Sep 17 00:00:00 2001 From: Prithvi Kannan Date: Mon, 9 Dec 2024 10:50:13 -0800 Subject: [PATCH] ruff Signed-off-by: Prithvi Kannan --- integrations/langchain/src/databricks_langchain/chat_models.py | 3 ++- integrations/langchain/src/databricks_langchain/embeddings.py | 3 ++- .../langchain/src/databricks_langchain/vectorstores.py | 1 + 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/integrations/langchain/src/databricks_langchain/chat_models.py b/integrations/langchain/src/databricks_langchain/chat_models.py index 490aa11..1dbe33c 100644 --- a/integrations/langchain/src/databricks_langchain/chat_models.py +++ b/integrations/langchain/src/databricks_langchain/chat_models.py @@ -50,10 +50,11 @@ from langchain_core.tools import BaseTool from langchain_core.utils.function_calling import convert_to_openai_tool from langchain_core.utils.pydantic import is_basemodel_subclass -from databricks_langchain.utils import get_deployment_client from mlflow.deployments import BaseDeploymentClient # type: ignore from pydantic import BaseModel, Field +from databricks_langchain.utils import get_deployment_client + logger = logging.getLogger(__name__) diff --git a/integrations/langchain/src/databricks_langchain/embeddings.py b/integrations/langchain/src/databricks_langchain/embeddings.py index 49e04fd..421e45c 100644 --- a/integrations/langchain/src/databricks_langchain/embeddings.py +++ b/integrations/langchain/src/databricks_langchain/embeddings.py @@ -1,9 +1,10 @@ from typing import Any, Dict, Iterator, List from langchain_core.embeddings import Embeddings -from databricks_langchain.utils import get_deployment_client from pydantic import BaseModel, PrivateAttr +from databricks_langchain.utils import get_deployment_client + class DatabricksEmbeddings(Embeddings, BaseModel): """Databricks embedding model integration. diff --git a/integrations/langchain/src/databricks_langchain/vectorstores.py b/integrations/langchain/src/databricks_langchain/vectorstores.py index 0e7372b..e67f315 100644 --- a/integrations/langchain/src/databricks_langchain/vectorstores.py +++ b/integrations/langchain/src/databricks_langchain/vectorstores.py @@ -22,6 +22,7 @@ from langchain_core.documents import Document from langchain_core.embeddings import Embeddings from langchain_core.vectorstores import VST, VectorStore + from databricks_langchain.utils import maximal_marginal_relevance logger = logging.getLogger(__name__)