dataiku-api-client 14.5.1__tar.gz → 14.6.0__tar.gz

{dataiku_api_client-14.5.1 → dataiku_api_client-14.6.0}/HISTORY.txt

@@ -2,6 +2,15 @@ Changelog
 ==========
 
 
+14.6.0 (2026-05-27)
+-------------------
+
+* Initial release for DSS 14.6.0
+
+14.5.2 (2026-05-25)
+-------------------
+
+* Initial release for DSS 14.5.2
 
 14.5.1 (2026-04-27)
 -------------------

{dataiku_api_client-14.5.1/dataiku_api_client.egg-info → dataiku_api_client-14.6.0}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: dataiku-api-client
-Version: 14.5.1
+Version: 14.6.0
 Summary: Python API client for Dataiku APIs
 Home-page: https://www.dataiku.com
-Download-URL: https://github.com/dataiku/dataiku-api-client-python/tarball/14.5.1
+Download-URL: https://github.com/dataiku/dataiku-api-client-python/tarball/14.6.0
 Author: Dataiku
 Author-email: support@dataiku.com
 License: Apache Software License
@@ -36,6 +36,15 @@ Changelog
 ==========
 
 
+14.6.0 (2026-05-27)
+-------------------
+
+* Initial release for DSS 14.6.0
+
+14.5.2 (2026-05-25)
+-------------------
+
+* Initial release for DSS 14.5.2
 
 14.5.1 (2026-04-27)
 -------------------

{dataiku_api_client-14.5.1 → dataiku_api_client-14.6.0/dataiku_api_client.egg-info}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: dataiku-api-client
-Version: 14.5.1
+Version: 14.6.0
 Summary: Python API client for Dataiku APIs
 Home-page: https://www.dataiku.com
-Download-URL: https://github.com/dataiku/dataiku-api-client-python/tarball/14.5.1
+Download-URL: https://github.com/dataiku/dataiku-api-client-python/tarball/14.6.0
 Author: Dataiku
 Author-email: support@dataiku.com
 License: Apache Software License
@@ -36,6 +36,15 @@ Changelog
 ==========
 
 
+14.6.0 (2026-05-27)
+-------------------
+
+* Initial release for DSS 14.6.0
+
+14.5.2 (2026-05-25)
+-------------------
+
+* Initial release for DSS 14.5.2
 
 14.5.1 (2026-04-27)
 -------------------

{dataiku_api_client-14.5.1 → dataiku_api_client-14.6.0}/dataikuapi/dss/admin.py

@@ -2202,6 +2202,7 @@ class DSSCodeEnvContainerConfsBearer(object):
         * PYTHON36_SUPPORT
         * PYTHON37_SUPPORT
         * PYTHON38_SUPPORT
+        * HUGGING_FACE_LOCAL_GPU
 
         :param dict container_runtime_addition: a dict with the container runtime addition definition
         """

{dataiku_api_client-14.5.1 → dataiku_api_client-14.6.0}/dataikuapi/dss/agent.py

@@ -226,6 +226,11 @@ class DSSAgentInteractionLoggingSelection(object):
         self.mode = self.MODE_NONE
 
 
+# Neutral aliases for shared LLM interaction logging concepts.
+DSSLLMInteractionLoggingSettings = DSSAgentInteractionLoggingSettings
+DSSLLMInteractionLoggingSelection = DSSAgentInteractionLoggingSelection
+
+
 class DSSAgentListItem(DSSTaggableObjectListItem):
     """
     An item in a list of agents
@@ -241,7 +246,7 @@ class DSSAgentListItem(DSSTaggableObjectListItem):
     @property
     def project_key(self):
         """
-        :returns: The project        
+        :returns: The project
         :rtype: string
         """
         return self._data["projectKey"]
@@ -263,7 +268,10 @@ class DSSAgentListItem(DSSTaggableObjectListItem):
         return self._data["name"]
 
     def as_llm(self):
-        """Returns this agent as a usable :class:`dataikuapi.dss.llm.DSSLLM` for querying"""
+        """
+        :returns: this agent as a usable :class:`dataikuapi.dss.llm.DSSLLM` for querying
+        :rtype: dataikuapi.dss.llm.DSSLLM
+        """
         return self.client.get_project(self.project_key).get_llm("agent:%s" % self.id)
 
 
@@ -285,14 +293,15 @@ class DSSAgent(object):
         return self._id
 
     def as_llm(self):
-        """Returns this agent as a usable :class:`dataikuapi.dss.llm.DSSLLM` for querying"""
+        """
+        :returns: this agent as a usable :class:`dataikuapi.dss.llm.DSSLLM` for querying
+        :rtype: dataikuapi.dss.llm.DSSLLM
+        """
         return self.client.get_project(self.project_key).get_llm("agent:%s" % self.id)
 
     def get_settings(self):
         """
-        Get the agent's definition
-
-        :return: a handle on the agent definition
+        :return: a handle on the agent's definition
         :rtype: :class:`dataikuapi.dss.agent.DSSAgentSettings`
         """
         settings = self.client._perform_json(
@@ -362,15 +371,26 @@ class DSSAgentSettings(DSSTaggableObjectSettings):
         self._settings = settings
 
     def get_version_ids(self):
+        """
+        List the ids of each version of this agent
+
+        :rtype: list[str]
+        """
         return [v["versionId"] for v in self._settings["versions"]]
 
     @property
     def active_version(self):
-        """Returns the active version of this agent. May return None if no version is declared as active"""
+        """
+        :returns: the active version of this agent, or None if no version is declared as active
+        :rtype: str | None
+        """
         return self._settings.get("activeVersion")
 
     def get_version_settings(self, version_id):
-
+        """
+        :returns: the settings of the given version of this agent
+        :rtype: DSSAgentVersionSettings
+        """
         version_settings = None
         for vs in self._settings["versions"]:
             if vs["versionId"] == version_id:
@@ -387,8 +407,7 @@ class DSSAgentSettings(DSSTaggableObjectSettings):
 
     def get_raw(self):
         """
-        Returns the raw settings of the agent
-        :return: the raw settings of the agent
+        :returns: the raw settings of this agent
         :rtype: dict
         """
         return self._settings
@@ -406,13 +425,18 @@ class DSSAgentVersionSettings(object):
         self._version_settings = version_settings
 
     def get_raw(self):
+        """
+        :returns: the raw settings of this agent version
+        :rtype: dict
+        """
         return self._version_settings
 
     @property
     def llm_id(self):
         """
         Only for Visual Agents
-        :rtype: :class:`str`
+
+        :rtype: str
         """
         if not self._agent_settings.type == "TOOLS_USING_AGENT":
             raise ValueError("Only valid for Simple Visual Agents")
@@ -502,7 +526,7 @@ class DSSAgentVersionSettings(object):
         .. code-block:: python
 
             project = client.get_project("MYPROJECT")
-            project.create_agent_interaction_logging_dataset(
+            project.create_llm_interaction_logging_dataset(
                 "agent_logs",
                 connection_id="filesystem_managed",
                 time_partitioning="DAY",

{dataiku_api_client-14.5.1 → dataiku_api_client-14.6.0}/dataikuapi/dss/agent_review.py

@@ -214,6 +214,32 @@ class DSSAgentReview(object):
     def nb_executions(self, value):
         self.data["nbExecutions"] = value
 
+    @property
+    def hitl_validation_policy(self):
+        """
+        Review-wide default HITL validation policy used when creating a test.
+        Possible values are ``ACCEPT`` or ``REJECT``.
+        :rtype: str
+        """
+        return self.data.get("hitlValidationPolicy")
+
+    @hitl_validation_policy.setter
+    def hitl_validation_policy(self, value):
+        self.data["hitlValidationPolicy"] = value
+
+    @property
+    def hitl_max_validation_turns(self):
+        """
+        Maximum number of validation turns auto-answered during a test run before stopping.
+        Possible values are integers greater than or equal to ``1``.
+        :rtype: int
+        """
+        return self.data.get("hitlMaxValidationTurns")
+
+    @hitl_max_validation_turns.setter
+    def hitl_max_validation_turns(self, value):
+        self.data["hitlMaxValidationTurns"] = value
+
     def get_trait(self, trait_id):
         """
         Get a specific trait by its ID.
@@ -278,13 +304,14 @@ class DSSAgentReview(object):
         )
         return DSSAgentReviewTest(self.dss_client, self.project_key, test)
 
-    def create_test(self, query=None, reference_answer=None, expectations=None):
+    def create_test(self, query=None, reference_answer=None, expectations=None, hitl_validation_policy=None):
         """
         Create a new test for this agent review.
 
         :param str query: Query to test the agent. Optional.
         :param str reference_answer: Reference answer. Optional.
         :param str expectations: Expectations on the agent answer. Optional.
+        :param str hitl_validation_policy: Per-test HITL validation policy. Optional.
         :returns: The created test object.
         :rtype: :class:`DSSAgentReviewTest`
         """
@@ -294,13 +321,14 @@ class DSSAgentReview(object):
             "query": query,
             "referenceAnswer": reference_answer,
             "expectations": expectations,
+            "hitlValidationPolicy": hitl_validation_policy,
         }
         test = self.dss_client._perform_json(
             "POST", "/projects/%s/agent-reviews/tests" % self.project_key, body=body
         )
         return DSSAgentReviewTest(self.dss_client, self.project_key, test)
 
-    def create_tests_from_dataset(self, full_dataset_name, query_column, reference_answer_column=None, expectations_column=None, top_n=None, partitions=None, latest_partitions_n=None):
+    def create_tests_from_dataset(self, full_dataset_name, query_column, reference_answer_column=None, expectations_column=None, top_n=None, partitions=None, latest_partitions_n=None, hitl_validation_policy_column=None):
         """
         Create new tests for this agent review by importing them from a dataset.
 
@@ -311,6 +339,7 @@ class DSSAgentReview(object):
         :param int top_n: Only take the first n rows of the dataset. Optional.
         :param list[str] partitions: For partitioned datasets, only consider the given partitions. Optional.
         :param int latest_partitions_n: For partitioned datasets and if partitions is not set, only consider the latest n partitions. Optional.
+        :param str hitl_validation_policy_column: Name of the column containing per-test HITL validation policy. Optional.
         :returns: A dictionary with keys:
             - "createdTestIds": list of ids of the created tests
             - "error": The error message if any occurred
@@ -324,6 +353,7 @@ class DSSAgentReview(object):
             "queryColumn": query_column,
             "referenceAnswerColumn": reference_answer_column,
             "expectationsColumn": expectations_column,
+            "hitlValidationPolicyColumn": hitl_validation_policy_column,
         }
         if top_n is not None:
             body["samplingMethod"] = "HEAD_SEQUENTIAL"
@@ -587,6 +617,19 @@ class DSSAgentReviewTest(object):
     def expectations(self, value):
         self.data["expectations"] = value
 
+    @property
+    def hitl_validation_policy(self):
+        """
+        Per-test HITL validation policy.
+        Possible values are ``ACCEPT`` or ``REJECT``.
+        :rtype: str
+        """
+        return self.data.get("hitlValidationPolicy")
+
+    @hitl_validation_policy.setter
+    def hitl_validation_policy(self, value):
+        self.data["hitlValidationPolicy"] = value
+
     @property
     def creation_timestamp(self):
         """
@@ -735,6 +778,15 @@ class DSSAgentReviewTestListItem(dict):
         """
         return self.get("expectations")
 
+    @property
+    def hitl_validation_policy(self):
+        """
+        Per-test HITL validation policy.
+        Possible values are ``ACCEPT`` or ``REJECT``.
+        :rtype: str
+        """
+        return self.get("hitlValidationPolicy")
+
 
 class DSSAgentReviewRun(object):
     """

{dataiku_api_client-14.5.1 → dataiku_api_client-14.6.0}/dataikuapi/dss/agent_tool.py

@@ -98,12 +98,22 @@ class DSSAgentTool(object):
         return self.client._perform_empty("DELETE", "/projects/%s/agents/tools/%s" % (self.project_key, self.id))
 
     def as_langchain_structured_tool(self, context = None):
+        """
+        :returns: this tool as a LangChain StructuredTool
+        :rtype: langchain_core.tools.StructuredTool
+        """
         from dataikuapi.dss.langchain.tool import convert_to_langchain_structured_tool
         return convert_to_langchain_structured_tool(self, context)
 
     def run(self, input, context=None, subtool_name=None, memory_fragment=None, tool_validation_responses=None, tool_validation_requests=None):
         """
         Execute a tool call
+
+        :param str input: Text input
+        :param dict context: Additional request context
+        :param str subtool_name: Name of the sub-tool, if applicable (e.g., for a MCP tool)
+        :rtype: dict
+        :returns: The result of running this tool
         """
 
         invocation = {
@@ -201,6 +211,10 @@ class DSSAgentToolSettings(DSSTaggableObjectSettings):
         self._settings = settings
 
     def get_raw(self):
+        """
+        :returns: the raw settings dict for this agent tool
+        :rtype: dict
+        """
         return self._settings
 
     @property
@@ -237,6 +251,10 @@ class DSSVectorStoreSearchAgentToolCreator(DSSAgentToolCreator):
         DSSAgentToolCreator.__init__(self, project, type, name, id)
 
     def with_knowledge_bank(self, kb):
+        """
+        :param kb: Knowledge Bank (object, list item, or identifier) to use in this tool
+        :type kb: DSSKnowledgeBank | DSSKnowledgeBankListItem | str
+        """
         loc = _kb_to_loc(self.project.project_key, kb)
         self.proto["creationParams"]["knowledgeBankRef"] = loc.to_ref(self.project.project_key)
         return self
@@ -246,5 +264,9 @@ class DSSVectorStoreSearchAgentToolSettings(DSSAgentToolSettings):
         DSSAgentToolSettings.__init__(self, agent_tool, settings)
 
     def set_knowledge_bank(self, kb):
+        """
+        :param kb: Knowledge Bank (object, list item, or identifier) to use in this tool
+        :type kb: DSSKnowledgeBank | DSSKnowledgeBankListItem | str
+        """
         loc = _kb_to_loc(self.project.project_key, kb)
         self.settings["params"]["knowledgeBankRef"] = loc.to_ref(self.agent_tool.project_key)

{dataiku_api_client-14.5.1 → dataiku_api_client-14.6.0}/dataikuapi/dss/document_extractor.py

@@ -3,6 +3,8 @@ import copy
 import json
 import warnings
 
+from dataikuapi.dss.llm_utils import get_json_schema_and_parser
+from dataikuapi.dss.recipe import DSSRecipe
 from dataikuapi.utils import _write_response_content_to_file
 
 
@@ -70,6 +72,89 @@ class DocumentExtractor(object):
                                         body=extractor_request)
         return VlmExtractorResponse(ret)
 
+
+    def vlm_extract_fields(self, images, schema=None, llm_id=None, llm_prompt=None, from_recipe=None, strict=None, compatible=None):
+        """
+        Extract specific fields (structured data) from images (typically screenshots of a document's pages) using a vision LLM.
+        Describe expected fields in ``extraction_schema``, or specify an Extract Fields recipe to use its settings.
+
+        :param images: screenshots of the document's pages from which to extract the fields
+        :type images: iterable(:class:`InlineImageRef`) | iterable(:class:`ManagedFolderImageRef`)
+        :param schema: a JSON schema or a Pydantic model class describing the fields to extract.
+                       JSON schema definitions or Pydantic models referencing other models are unsupported.
+        :type schema: str | dict | pydantic.BaseModel class or Python type hint.
+        :param llm_id: Identifier of a vision LLM
+        :type llm_id: str
+        :param llm_prompt: Custom prompt to extract fields from the images
+        :type llm_prompt: str
+        :param from_recipe: Name of a recipe from which to read the other arguments.
+                            Arguments provided explicitly take precedence.
+        :type from_recipe: str
+        :param strict: Whether to strictly enforce the schema. Support varies across models/providers.
+        :type strict: bool
+        :param compatible: Allow DSS to modify the schema in order to increase compatibility, depending on known limitations of the model/provider. Defaults to automatic.
+        :type compatible: bool
+
+        :returns: Extracted fields from images
+        :rtype: :class:`FieldsVlmExtractorResponse`
+        """
+        if from_recipe is not None and (schema is None or llm_id is None or llm_prompt is None or strict is None or compatible is None):
+            recipe_settings = DSSRecipe(self.client, self.project_key, from_recipe).get_settings()
+            recipe_params = recipe_settings.raw_params
+
+            if schema is None:
+                schema = recipe_params.get("extractionSchema")
+            if llm_id is None:
+                llm_id = recipe_params.get("vlmId")
+            if llm_prompt is None:
+                llm_prompt = recipe_params.get("extractionPrompt")
+            if strict is None:
+                strict = recipe_params.get("valueErrorHandling") == "BLANK_VALUE"
+            if compatible is None:
+                compatible = True
+                dku_properties = recipe_settings.get_recipe_raw_definition().get("dkuProperties") or []
+                for prop in dku_properties:
+                    if isinstance(prop, dict) and prop.get("name") == "dku.extractFields.schemaCompatibilityEnhancer" and isinstance(prop.get("value"), bool):
+                        compatible = prop.get("value")
+                        break
+
+        json_schema, parser_method = get_json_schema_and_parser(schema)
+
+        extractor_request = {
+            "settings": {
+                "llmId": llm_id,
+                "llmPrompt": llm_prompt,
+                "extractionSchema" : json.loads(json_schema),
+                "strictSchemaAdherence": strict,
+                "schemaCompatibilityEnhancer": compatible
+            }
+        }
+
+        images = list(images)
+        if not images:
+            raise ValueError("No images provided")
+        if all(isinstance(ir, InlineImageRef) for ir in images):
+            extractor_request["inputs"] = {
+                "imagesRef": {
+                    "type": images[0].type,
+                    "inlineImages": [ir.as_dict() for ir in images]
+                }
+            }
+        elif all(isinstance(ir, ManagedFolderImageRef) for ir in images):
+            extractor_request["inputs"] = {
+                "imagesRef": {
+                    "type": images[0].type,
+                    "managedFolderRef": images[0].managed_folder_ref,
+                    "imagesPaths": [ir.image_path for ir in images]
+                }
+            }
+        else:
+            raise ValueError("Unsupported mix of image types: %s" % set([ir.type for ir in images]))
+
+        ret = self.client._perform_json("POST", "/projects/%s/document-extractors/fields/vlm" % self.project_key,
+                                        body=extractor_request)
+        return FieldsVlmExtractorResponse(ret, parser_method)
+
     def structured_extract(self, document, max_section_depth=6, image_handling_mode='IGNORE', ocr_engine='AUTO', languages="en", llm_id=None, llm_prompt=None,
                            output_managed_folder=None, image_validation=True):
         """
@@ -191,7 +276,8 @@ class DocumentExtractor(object):
 
         return TextExtractorResponse(ret)
 
-    def generate_pages_screenshots(self, document, output_managed_folder=None, offset=0, fetch_size=10, keep_fetched=True):
+    def generate_pages_screenshots(self, document, output_managed_folder=None, offset=0, fetch_size=10, keep_fetched=True,
+                                   page_dpi=None, max_memory_per_document=None):
         """
         Generate per-page screenshots of a document, returning an iterable over the screenshots.
 
@@ -232,12 +318,24 @@ class DocumentExtractor(object):
         :type fetch_size: int
         :param keep_fetched: whether to keep previous screenshots requests within this response object when fetching next ones.
         :type keep_fetched: boolean
+        :param page_dpi: DPI used to render pages if memory allows.
+        :type page_dpi: int
+        :param max_memory_per_document: maximum memory budget in MB used while rendering a document.
+            The effective DPI may be reduced to fit this limit depending on page dimensions.
+        :type max_memory_per_document: int
 
         :returns: An iterable over the result screenshots
         :rtype: :class:`ScreenshotterResponse`
         """
 
-        screenshotter_request = ScreenshotterRequest(document, output_managed_folder, offset, fetch_size)
+        screenshotter_request = ScreenshotterRequest(
+            document,
+            output_managed_folder,
+            offset,
+            fetch_size,
+            page_dpi,
+            max_memory_per_document
+        )
         return ScreenshotterResponse(self.client, self.project_key, screenshotter_request, keep_fetched)
 
 
@@ -269,11 +367,13 @@ class ScreenshotterRequest(object):
 
     """
 
-    def __init__(self, document, output_managed_folder, offset, fetch_size):
+    def __init__(self, document, output_managed_folder, offset, fetch_size, page_dpi=None, max_memory_per_document=None):
         self.document = document
         self.output_managed_folder = output_managed_folder
         self.offset = offset
         self.fetch_size = fetch_size
+        self.page_dpi = page_dpi
+        self.max_memory_per_document = max_memory_per_document # MB
 
     def as_json(self):
         """
@@ -302,6 +402,8 @@ class ScreenshotterRequest(object):
                 "outputManagedFolderRef": self.output_managed_folder,
                 "paginationOffset": self.offset,
                 "paginationSize": self.fetch_size,
+                "pageDPI": self.page_dpi,
+                "maxMemoryPerDocMB": self.max_memory_per_document,
             }
         }
 
@@ -683,7 +785,6 @@ class VlmExtractorResponse(object):
     @property
     def success(self):
         """
-        :returns: The outcome of the extractor request.
         :rtype: bool
         """
         return self._data.get("ok")
@@ -706,6 +807,77 @@ class VlmExtractorResponse(object):
             )
             raise Exception(error_message)
 
+class FieldsVlmExtractorResponse(object):
+    """
+    A handle to interact with a VLM fields extraction result.
+
+    .. important::
+        Do not create this class directly; use :meth:`~DocumentExtractor.vlm_extract_fields`
+    """
+
+    def __init__(self, data, response_parser=None):
+        self._data = data
+        self._response_parser = response_parser
+        self._parsed_fields = None
+
+    def get_raw(self):
+        return self._data
+
+    @property
+    def success(self):
+        """
+        :rtype: bool
+        """
+        return self._data.get("status", "nok") != "nok"
+
+    @property
+    def fields(self):
+        """
+        Fields extracted from the original document.
+        Follows the structure of the extraction schema, has only the fields that abide by it.
+
+        :returns: extracted fields.
+        :rtype: dict
+        """
+        self._fail_unless_success()
+        return self._data["curatedFields"]
+
+    @property
+    def parsed_fields(self):
+        """
+        Fields extracted from the original document.
+        Follows the structure of the extraction schema, has only the fields that abide by it.
+        Only available for extraction schema given as a Pydantic model or using Python type hint.
+
+        :returns: extracted fields deserialized into a Pydantic model instance.
+        :rtype: pydantic.BaseModel
+        """
+        self._fail_unless_success()
+        if self._parsed_fields is None and self._data["curatedFields"] is not None:
+            if not self._response_parser:
+                raise Exception("Require extraction schema to be given as a Pydantic model or using Python type hint.")
+            else:
+                self._parsed_fields = self._response_parser(json.dumps(self._data["curatedFields"]))
+        return self._parsed_fields
+
+    @property
+    def invalid_fields(self):
+        """
+        Fields in the extraction schema that the Vision LLM could not extract.
+        Follows the structure/hierarchy of the extraction schema, but has only the incorrect or missing fields.
+
+        :rtype: dict
+        """
+        self._fail_unless_success()
+        return self._data["invalidFields"]
+
+    def _fail_unless_success(self):
+        if not self.success:
+            error_message = "Document failed to be extracted - request failed: {}".format(
+                self._data.get("errorMessage", "An unknown error occurred")
+            )
+            raise Exception(error_message)
+
 
 class InputRef(object):
     def as_dict(self):

{dataiku_api_client-14.5.1 → dataiku_api_client-14.6.0}/dataikuapi/dss/knowledgebank.py

@@ -245,7 +245,7 @@ class DSSKnowledgeBankSettings(DSSTaggableObjectSettings):
     @property
     def vector_store_type(self):
         """
-        Returns the type of storage backing the vector store (could be CHROMA, PINECONE, ELASTICSEARCH, AZURE_AI_SEARCH, VERTEX_AI_GCS_BASED, FAISS, QDRANT_LOCAL)
+        Returns the type of storage backing the vector store (could be CHROMA, PINECONE, ELASTICSEARCH, AZURE_AI_SEARCH, SNOWFLAKE_CORTEX_SEARCH, VERTEX_AI_GCS_BASED, FAISS, QDRANT_LOCAL)
 
         :rtype: str
         """

{dataiku_api_client-14.5.1 → dataiku_api_client-14.6.0}/dataikuapi/dss/langchain/tool.py

@@ -2,11 +2,8 @@ import textwrap
 import warnings
 from typing import Optional, Callable, Literal, Type, Any, Dict
 
-try:
-    from langchain_core.tools import StructuredTool
-except ModuleNotFoundError:
-    from langchain.tools import StructuredTool
-
+from langchain_core.tools import StructuredTool, create_schema_from_function
+from langchain_core.tools.structured import _filter_schema_args
 from pydantic import BaseModel
 
 from dataiku.langchain.dku_tracer import dku_span_builder_for_callbacks
@@ -98,12 +95,12 @@ class DKUStructuredTool(StructuredTool):
         name = name or source_function.__name__
         if args_schema is None and infer_schema:
             # schema name is appended within function
-            args_schema = create_schema_from_function(  # noqa: F821
+            args_schema = create_schema_from_function(
                 name,
                 source_function,
                 parse_docstring=parse_docstring,
                 error_on_invalid_docstring=error_on_invalid_docstring,
-                filter_args=_filter_schema_args(source_function),  # noqa: F821
+                filter_args=_filter_schema_args(source_function),
             )
         description_ = description
         if description is None and not parse_docstring: