From 0a9b1b5a56d1b379d5f369f8806f24779733370c Mon Sep 17 00:00:00 2001 From: Kelvin Sundli Date: Fri, 17 Apr 2026 16:46:31 -0700 Subject: [PATCH 01/11] test(agents): branch on expected_type in agent tool tests Tests were branching on isinstance(loaded_tool, ...), which is the object under test. A loader regression returning the wrong subtype could silently skip entire assertion blocks and still pass. Branching on the parametrized expected_type keeps the isinstance check as a real assertion. --- .../test_agents/test_agent_tools.py | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py b/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py index a87904e2f8..58a8829474 100644 --- a/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py +++ b/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py @@ -70,23 +70,22 @@ def test_agent_tool_load_returns_correct_subtype(self, tool_data: dict, expected assert loaded_tool.name == tool_data["name"] assert loaded_tool.description == tool_data["description"] - if isinstance(loaded_tool, UnknownAgentTool): + if expected_type is UnknownAgentTool: + assert isinstance(loaded_tool, UnknownAgentTool) assert loaded_tool.type == tool_data["type"] else: assert loaded_tool._type == expected_type._type - # Handle configuration comparison based on tool type if "configuration" in tool_data: - if isinstance(loaded_tool, QueryKnowledgeGraphAgentTool): - # For QueryKnowledgeGraph, we expect a structured configuration object + if expected_type is QueryKnowledgeGraphAgentTool: + assert isinstance(loaded_tool, QueryKnowledgeGraphAgentTool) assert isinstance(loaded_tool.configuration, QueryKnowledgeGraphAgentToolConfiguration) - # Compare by serializing the structured object back to dict assert loaded_tool.configuration.dump(camel_case=True) == tool_data["configuration"] - elif isinstance(loaded_tool, UnknownAgentTool): - # For other tools (like UnknownAgentTool), configuration should be a dict + elif expected_type is UnknownAgentTool: + assert isinstance(loaded_tool, UnknownAgentTool) assert loaded_tool.configuration == tool_data["configuration"] else: - raise TypeError(f"Unhandled tool type in test case: {type(loaded_tool)}") + raise TypeError(f"Unhandled tool type in test case: {expected_type}") def test_unknown_agent_tool_preserves_custom_type(self) -> None: """Test that UnknownAgentTool preserves the original type string.""" @@ -113,7 +112,8 @@ def test_agent_tool_dump_returns_correct_type(self, tool_data: dict, expected_ty loaded_tool = AgentTool._load(tool_data) dumped_tool = loaded_tool.dump(camel_case=True) - if isinstance(loaded_tool, UnknownAgentTool): + if expected_type is UnknownAgentTool: + assert isinstance(loaded_tool, UnknownAgentTool) assert dumped_tool["type"] == unknown_example["type"] else: assert dumped_tool["type"] == expected_type._type @@ -158,7 +158,8 @@ def test_agent_tool_upsert_returns_correct_type(self, tool_data: dict, expected_ loaded_tool = AgentTool._load(tool_data) dumped_tool = loaded_tool.as_write().dump(camel_case=True) - if isinstance(loaded_tool, UnknownAgentTool): + if expected_type is UnknownAgentTool: + assert isinstance(loaded_tool, UnknownAgentTool) assert dumped_tool["type"] == unknown_example["type"] else: assert dumped_tool["type"] == expected_type._type From a4924df8c12203997ed12d92fdc7104be7057bc6 Mon Sep 17 00:00:00 2001 From: Kelvin Sundli Date: Fri, 17 Apr 2026 17:29:43 -0700 Subject: [PATCH 02/11] test(agents): assert loaded_tool type up front in dump and upsert tests Mirrors the pattern already used by test_agent_tool_load_returns_correct_subtype, so a loader regression returning a wrong subtype fails loudly here instead of relying on _type string equality alone. --- .../test_data_classes/test_agents/test_agent_tools.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py b/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py index 58a8829474..d0215314b6 100644 --- a/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py +++ b/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py @@ -110,10 +110,10 @@ class TestAgentToolDump: def test_agent_tool_dump_returns_correct_type(self, tool_data: dict, expected_type: type[AgentTool]) -> None: """Test that AgentTool.dump() returns the correct type.""" loaded_tool = AgentTool._load(tool_data) + assert isinstance(loaded_tool, expected_type) dumped_tool = loaded_tool.dump(camel_case=True) if expected_type is UnknownAgentTool: - assert isinstance(loaded_tool, UnknownAgentTool) assert dumped_tool["type"] == unknown_example["type"] else: assert dumped_tool["type"] == expected_type._type @@ -156,10 +156,10 @@ class TestAgentToolUpsert: def test_agent_tool_upsert_returns_correct_type(self, tool_data: dict, expected_type: type[AgentTool]) -> None: """Test that AgentToolUpsert.dump() returns the correct type.""" loaded_tool = AgentTool._load(tool_data) + assert isinstance(loaded_tool, expected_type) dumped_tool = loaded_tool.as_write().dump(camel_case=True) if expected_type is UnknownAgentTool: - assert isinstance(loaded_tool, UnknownAgentTool) assert dumped_tool["type"] == unknown_example["type"] else: assert dumped_tool["type"] == expected_type._type From f46b65f47ffc1bcac8a119a98152b2c1260f6052 Mon Sep 17 00:00:00 2001 From: Kelvin Sundli Date: Wed, 15 Apr 2026 22:19:51 -0700 Subject: [PATCH 03/11] feat: add QueryAgentTool for the "query" tool type The cog-ai backend and Fusion frontend already support the "query" tool type, but the Python SDK deserialized it as UnknownAgentTool. This adds first-class support with QueryAgentTool, QueryAgentToolUpsert, and QueryAgentToolConfiguration classes. --- cognite/client/_api/agents/agents.py | 31 ++++++ cognite/client/_sync_api/agents/agents.py | 33 +++++- .../client/data_classes/agents/__init__.py | 6 ++ .../client/data_classes/agents/agent_tools.py | 100 ++++++++++++++++++ .../test_agents/test_agent_tools.py | 76 ++++++++++++- 5 files changed, 244 insertions(+), 2 deletions(-) diff --git a/cognite/client/_api/agents/agents.py b/cognite/client/_api/agents/agents.py index 03869378aa..a0214d9183 100644 --- a/cognite/client/_api/agents/agents.py +++ b/cognite/client/_api/agents/agents.py @@ -79,6 +79,37 @@ async def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | A ... ) >>> client.agents.upsert(agents=[agent]) + Create an agent with the query tool (Preview): + + >>> from cognite.client.data_classes.agents import ( + ... AgentUpsert, + ... QueryAgentToolUpsert, + ... QueryAgentToolConfiguration, + ... DataModelInfo, + ... InstanceSpaces, + ... ) + >>> query_tool = QueryAgentToolUpsert( + ... name="explore data", + ... description="Run flexible queries against your data model", + ... configuration=QueryAgentToolConfiguration( + ... data_models=[ + ... DataModelInfo( + ... space="cdf_idm", + ... external_id="CogniteProcessIndustries", + ... version="v1", + ... ) + ... ], + ... instance_spaces=InstanceSpaces(type="all"), + ... ), + ... ) + >>> agent = AgentUpsert( + ... external_id="my_agent", + ... name="My Agent", + ... labels=["published"], + ... tools=[query_tool], + ... ) + >>> client.agents.upsert(agents=[agent]) + Create an agent with multiple different tools: >>> from cognite.client.data_classes.agents import ( diff --git a/cognite/client/_sync_api/agents/agents.py b/cognite/client/_sync_api/agents/agents.py index 01824b8959..48fd6c9a99 100644 --- a/cognite/client/_sync_api/agents/agents.py +++ b/cognite/client/_sync_api/agents/agents.py @@ -1,6 +1,6 @@ """ =============================================================================== -063c42ab744021733ccbdc455b150b2c +d88277929dfc151715f521eb7176622b This file is auto-generated from the Async API modules, - do not edit manually! =============================================================================== """ @@ -77,6 +77,37 @@ def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | AgentLi ... ) >>> client.agents.upsert(agents=[agent]) + Create an agent with the query tool (Preview): + + >>> from cognite.client.data_classes.agents import ( + ... AgentUpsert, + ... QueryAgentToolUpsert, + ... QueryAgentToolConfiguration, + ... DataModelInfo, + ... InstanceSpaces, + ... ) + >>> query_tool = QueryAgentToolUpsert( + ... name="explore data", + ... description="Run flexible queries against your data model", + ... configuration=QueryAgentToolConfiguration( + ... data_models=[ + ... DataModelInfo( + ... space="cdf_idm", + ... external_id="CogniteProcessIndustries", + ... version="v1", + ... ) + ... ], + ... instance_spaces=InstanceSpaces(type="all"), + ... ), + ... ) + >>> agent = AgentUpsert( + ... external_id="my_agent", + ... name="My Agent", + ... labels=["published"], + ... tools=[query_tool], + ... ) + >>> client.agents.upsert(agents=[agent]) + Create an agent with multiple different tools: >>> from cognite.client.data_classes.agents import ( diff --git a/cognite/client/data_classes/agents/__init__.py b/cognite/client/data_classes/agents/__init__.py index f0d39aebc9..9531002d5e 100644 --- a/cognite/client/data_classes/agents/__init__.py +++ b/cognite/client/data_classes/agents/__init__.py @@ -9,6 +9,9 @@ AskDocumentAgentToolUpsert, DataModelInfo, InstanceSpaces, + QueryAgentTool, + QueryAgentToolConfiguration, + QueryAgentToolUpsert, QueryKnowledgeGraphAgentTool, QueryKnowledgeGraphAgentToolConfiguration, QueryKnowledgeGraphAgentToolUpsert, @@ -70,6 +73,9 @@ "Message", "MessageContent", "MessageList", + "QueryAgentTool", + "QueryAgentToolConfiguration", + "QueryAgentToolUpsert", "QueryKnowledgeGraphAgentTool", "QueryKnowledgeGraphAgentToolConfiguration", "QueryKnowledgeGraphAgentToolUpsert", diff --git a/cognite/client/data_classes/agents/agent_tools.py b/cognite/client/data_classes/agents/agent_tools.py index 8957b208cf..a9287563ca 100644 --- a/cognite/client/data_classes/agents/agent_tools.py +++ b/cognite/client/data_classes/agents/agent_tools.py @@ -194,6 +194,43 @@ def as_write(self) -> QueryKnowledgeGraphAgentToolConfiguration: return self +@dataclass +class QueryAgentToolConfiguration(WriteableCogniteResource): + """Configuration for query agent tools. + + Args: + data_models (Sequence[DataModelInfo]): The data models to query. + instance_spaces (InstanceSpaces | None): The instance spaces to query. + """ + + data_models: Sequence[DataModelInfo] + instance_spaces: InstanceSpaces | None = None + + @classmethod + def _load(cls, resource: dict[str, Any]) -> QueryAgentToolConfiguration: + dm_config = resource.get("dataModels", {}) + data_models = [DataModelInfo._load(dm) for dm in dm_config.get("dataModels", [])] + return cls( + data_models=data_models, + instance_spaces=InstanceSpaces._load_if(resource.get("instanceSpaces")), + ) + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + result: dict[str, Any] = {} + key = "dataModels" if camel_case else "data_models" + result[key] = { + "type": "manual", + key: [dm.dump(camel_case=camel_case) for dm in self.data_models], + } + if self.instance_spaces: + key = "instanceSpaces" if camel_case else "instance_spaces" + result[key] = self.instance_spaces.dump(camel_case=camel_case) + return result + + def as_write(self) -> QueryAgentToolConfiguration: + return self + + @dataclass class SummarizeDocumentAgentTool(AgentTool): """Agent tool for summarizing documents. @@ -389,6 +426,69 @@ def _load(cls, resource: dict[str, Any]) -> QueryTimeSeriesDatapointsAgentToolUp ) +@dataclass +class QueryAgentTool(AgentTool): + """Agent tool for running flexible queries against data models. + + Args: + name (str): The name of the agent tool. Used by the agent to decide when to use this tool. + description (str): The description of the agent tool. Used by the agent to decide when to use this tool. + configuration (QueryAgentToolConfiguration | None): The configuration of the query agent tool. + """ + + _type: ClassVar[str] = "query" + configuration: QueryAgentToolConfiguration | None = None + + @classmethod + def _load_tool(cls, resource: dict[str, Any]) -> QueryAgentTool: + return cls( + name=resource["name"], + description=resource["description"], + configuration=QueryAgentToolConfiguration._load_if(resource.get("configuration")), + ) + + def as_write(self) -> QueryAgentToolUpsert: + return QueryAgentToolUpsert( + name=self.name, + description=self.description, + configuration=self.configuration, + ) + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + result = super().dump(camel_case=camel_case) + if self.configuration: + result["configuration"] = self.configuration.dump(camel_case=camel_case) + return result + + +@dataclass +class QueryAgentToolUpsert(AgentToolUpsert): + """Upsert version of query agent tool. + + Args: + name (str): The name of the agent tool. Used by the agent to decide when to use this tool. + description (str): The description of the agent tool. Used by the agent to decide when to use this tool. + configuration (QueryAgentToolConfiguration | None): The configuration of the query agent tool. + """ + + _type: ClassVar[str] = "query" + configuration: QueryAgentToolConfiguration | None = None + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + result = super().dump(camel_case=camel_case) + if self.configuration: + result["configuration"] = self.configuration.dump(camel_case=camel_case) + return result + + @classmethod + def _load(cls, resource: dict[str, Any]) -> QueryAgentToolUpsert: + return cls( + name=resource["name"], + description=resource["description"], + configuration=QueryAgentToolConfiguration._load_if(resource.get("configuration")), + ) + + @dataclass class UnknownAgentTool(AgentTool): """Agent tool for unknown/unrecognized tool types. diff --git a/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py b/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py index d0215314b6..8dd0f6a200 100644 --- a/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py +++ b/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py @@ -5,6 +5,8 @@ from cognite.client.data_classes.agents.agent_tools import ( AgentTool, AskDocumentAgentTool, + QueryAgentTool, + QueryAgentToolConfiguration, QueryKnowledgeGraphAgentTool, QueryKnowledgeGraphAgentToolConfiguration, QueryTimeSeriesDatapointsAgentTool, @@ -42,6 +44,32 @@ "description": "Query the time series datapoints", } +query_example = { + "name": "queryExample", + "type": "query", + "description": "Run flexible queries against your data model", + "configuration": { + "dataModels": { + "type": "manual", + "dataModels": [ + { + "space": "cdf_idm", + "externalId": "CogniteProcessIndustries", + "version": "v1", + "viewExternalIds": ["CogniteAsset"], + } + ], + }, + "instanceSpaces": {"type": "manual", "spaces": ["my_space"]}, + }, +} + +query_no_config_example = { + "name": "queryNoConfigExample", + "type": "query", + "description": "Run flexible queries against your data model", +} + unknown_example = { "name": "unknownExample", "type": "yolo", # This is not a known tool type @@ -58,9 +86,19 @@ class TestAgentToolLoad: (ask_document_example, AskDocumentAgentTool), (summarize_document_example, SummarizeDocumentAgentTool), (query_time_series_datapoints_example, QueryTimeSeriesDatapointsAgentTool), + (query_example, QueryAgentTool), + (query_no_config_example, QueryAgentTool), (unknown_example, UnknownAgentTool), ], - ids=["queryKnowledgeGraph", "askDocument", "summarizeDocument", "queryTimeSeriesDatapoints", "somethingElse"], + ids=[ + "queryKnowledgeGraph", + "askDocument", + "summarizeDocument", + "queryTimeSeriesDatapoints", + "query", + "queryNoConfig", + "somethingElse", + ], ) def test_agent_tool_load_returns_correct_subtype(self, tool_data: dict, expected_type: type[AgentTool]) -> None: """Test that AgentTool._load() returns the correct subtype based on the tool type.""" @@ -81,6 +119,10 @@ def test_agent_tool_load_returns_correct_subtype(self, tool_data: dict, expected assert isinstance(loaded_tool, QueryKnowledgeGraphAgentTool) assert isinstance(loaded_tool.configuration, QueryKnowledgeGraphAgentToolConfiguration) assert loaded_tool.configuration.dump(camel_case=True) == tool_data["configuration"] + elif expected_type is QueryAgentTool: + assert isinstance(loaded_tool, QueryAgentTool) + assert isinstance(loaded_tool.configuration, QueryAgentToolConfiguration) + assert loaded_tool.configuration.dump(camel_case=True) == tool_data["configuration"] elif expected_type is UnknownAgentTool: assert isinstance(loaded_tool, UnknownAgentTool) assert loaded_tool.configuration == tool_data["configuration"] @@ -104,6 +146,7 @@ class TestAgentToolDump: (ask_document_example, AskDocumentAgentTool), (summarize_document_example, SummarizeDocumentAgentTool), (query_time_series_datapoints_example, QueryTimeSeriesDatapointsAgentTool), + (query_example, QueryAgentTool), (unknown_example, UnknownAgentTool), ], ) @@ -150,6 +193,7 @@ class TestAgentToolUpsert: (ask_document_example, AskDocumentAgentTool), (summarize_document_example, SummarizeDocumentAgentTool), (query_time_series_datapoints_example, QueryTimeSeriesDatapointsAgentTool), + (query_example, QueryAgentTool), (unknown_example, UnknownAgentTool), ], ) @@ -166,3 +210,33 @@ def test_agent_tool_upsert_returns_correct_type(self, tool_data: dict, expected_ assert dumped_tool["name"] == tool_data["name"] assert dumped_tool["description"] == tool_data["description"] + + +class TestQueryAgentTool: + def test_load_with_configuration(self) -> None: + loaded = AgentTool._load(query_example) + assert isinstance(loaded, QueryAgentTool) + assert loaded.configuration is not None + assert isinstance(loaded.configuration, QueryAgentToolConfiguration) + assert len(loaded.configuration.data_models) == 1 + assert loaded.configuration.data_models[0].space == "cdf_idm" + assert loaded.configuration.data_models[0].external_id == "CogniteProcessIndustries" + assert loaded.configuration.instance_spaces is not None + assert loaded.configuration.instance_spaces.type == "manual" + assert loaded.configuration.instance_spaces.spaces == ["my_space"] + + def test_load_without_configuration(self) -> None: + loaded = AgentTool._load(query_no_config_example) + assert isinstance(loaded, QueryAgentTool) + assert loaded.configuration is None + + def test_round_trip(self) -> None: + loaded = AgentTool._load(query_example) + dumped = loaded.dump(camel_case=True) + assert dumped == query_example + + def test_as_write_round_trip(self) -> None: + loaded = AgentTool._load(query_example) + write = loaded.as_write() + dumped = write.dump(camel_case=True) + assert dumped == query_example From 7b3fe1b44e699698ae57e7ec80855838cb441b4b Mon Sep 17 00:00:00 2001 From: Kelvin Sundli Date: Wed, 15 Apr 2026 23:17:06 -0700 Subject: [PATCH 04/11] fix: defensive null handling in QueryAgentToolConfiguration._load --- cognite/client/data_classes/agents/agent_tools.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/cognite/client/data_classes/agents/agent_tools.py b/cognite/client/data_classes/agents/agent_tools.py index a9287563ca..40071e3384 100644 --- a/cognite/client/data_classes/agents/agent_tools.py +++ b/cognite/client/data_classes/agents/agent_tools.py @@ -208,8 +208,9 @@ class QueryAgentToolConfiguration(WriteableCogniteResource): @classmethod def _load(cls, resource: dict[str, Any]) -> QueryAgentToolConfiguration: - dm_config = resource.get("dataModels", {}) - data_models = [DataModelInfo._load(dm) for dm in dm_config.get("dataModels", [])] + # API always returns dataModels, but guard against null defensively: + dm_config = resource.get("dataModels") or {} + data_models = [DataModelInfo._load(dm) for dm in (dm_config.get("dataModels") or [])] return cls( data_models=data_models, instance_spaces=InstanceSpaces._load_if(resource.get("instanceSpaces")), From 46cc7d3230a23bc22be51857825c7b80303b2ce7 Mon Sep 17 00:00:00 2001 From: Kelvin Sundli Date: Thu, 16 Apr 2026 00:10:22 -0700 Subject: [PATCH 05/11] fix: address PR feedback on query agent tool Remove maturity label from docstring examples and surface missing dataModels key instead of silently defaulting to empty dict. --- cognite/client/_api/agents/agents.py | 2 +- cognite/client/_sync_api/agents/agents.py | 4 ++-- cognite/client/data_classes/agents/agent_tools.py | 5 ++--- 3 files changed, 5 insertions(+), 6 deletions(-) diff --git a/cognite/client/_api/agents/agents.py b/cognite/client/_api/agents/agents.py index a0214d9183..a842fe8341 100644 --- a/cognite/client/_api/agents/agents.py +++ b/cognite/client/_api/agents/agents.py @@ -79,7 +79,7 @@ async def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | A ... ) >>> client.agents.upsert(agents=[agent]) - Create an agent with the query tool (Preview): + Create an agent with the query tool: >>> from cognite.client.data_classes.agents import ( ... AgentUpsert, diff --git a/cognite/client/_sync_api/agents/agents.py b/cognite/client/_sync_api/agents/agents.py index 48fd6c9a99..f7bfc2a130 100644 --- a/cognite/client/_sync_api/agents/agents.py +++ b/cognite/client/_sync_api/agents/agents.py @@ -1,6 +1,6 @@ """ =============================================================================== -d88277929dfc151715f521eb7176622b +990185c07f95962938122269ec34ae29 This file is auto-generated from the Async API modules, - do not edit manually! =============================================================================== """ @@ -77,7 +77,7 @@ def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | AgentLi ... ) >>> client.agents.upsert(agents=[agent]) - Create an agent with the query tool (Preview): + Create an agent with the query tool: >>> from cognite.client.data_classes.agents import ( ... AgentUpsert, diff --git a/cognite/client/data_classes/agents/agent_tools.py b/cognite/client/data_classes/agents/agent_tools.py index 40071e3384..e8722b258b 100644 --- a/cognite/client/data_classes/agents/agent_tools.py +++ b/cognite/client/data_classes/agents/agent_tools.py @@ -208,9 +208,8 @@ class QueryAgentToolConfiguration(WriteableCogniteResource): @classmethod def _load(cls, resource: dict[str, Any]) -> QueryAgentToolConfiguration: - # API always returns dataModels, but guard against null defensively: - dm_config = resource.get("dataModels") or {} - data_models = [DataModelInfo._load(dm) for dm in (dm_config.get("dataModels") or [])] + dm_config = resource["dataModels"] + data_models = [DataModelInfo._load(dm) for dm in dm_config.get("dataModels", [])] return cls( data_models=data_models, instance_spaces=InstanceSpaces._load_if(resource.get("instanceSpaces")), From 06f067c91648e2a930c080b31422e48c153ee59b Mon Sep 17 00:00:00 2001 From: Kelvin Sundli Date: Thu, 16 Apr 2026 00:13:36 -0700 Subject: [PATCH 06/11] fix(tests): remove redundant round-trip tests for QueryAgentTool Covered automatically by tests/tests_unit/test_base.py. --- .../test_data_classes/test_agents/test_agent_tools.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py b/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py index 8dd0f6a200..b8237b49d4 100644 --- a/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py +++ b/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py @@ -229,14 +229,3 @@ def test_load_without_configuration(self) -> None: loaded = AgentTool._load(query_no_config_example) assert isinstance(loaded, QueryAgentTool) assert loaded.configuration is None - - def test_round_trip(self) -> None: - loaded = AgentTool._load(query_example) - dumped = loaded.dump(camel_case=True) - assert dumped == query_example - - def test_as_write_round_trip(self) -> None: - loaded = AgentTool._load(query_example) - write = loaded.as_write() - dumped = write.dump(camel_case=True) - assert dumped == query_example From b3d722f1f18efc00f855045ab35e2a874c65aeec Mon Sep 17 00:00:00 2001 From: Kelvin Sundli Date: Thu, 16 Apr 2026 00:26:41 -0700 Subject: [PATCH 07/11] fix(tests): remove redundant test_load_with_configuration Covered by automatic load/dump tests in test_base.py. Keep only test_load_without_configuration for the None config edge case. --- .../test_agents/test_agent_tools.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py b/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py index b8237b49d4..358fd4cc5b 100644 --- a/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py +++ b/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py @@ -213,18 +213,6 @@ def test_agent_tool_upsert_returns_correct_type(self, tool_data: dict, expected_ class TestQueryAgentTool: - def test_load_with_configuration(self) -> None: - loaded = AgentTool._load(query_example) - assert isinstance(loaded, QueryAgentTool) - assert loaded.configuration is not None - assert isinstance(loaded.configuration, QueryAgentToolConfiguration) - assert len(loaded.configuration.data_models) == 1 - assert loaded.configuration.data_models[0].space == "cdf_idm" - assert loaded.configuration.data_models[0].external_id == "CogniteProcessIndustries" - assert loaded.configuration.instance_spaces is not None - assert loaded.configuration.instance_spaces.type == "manual" - assert loaded.configuration.instance_spaces.spaces == ["my_space"] - def test_load_without_configuration(self) -> None: loaded = AgentTool._load(query_no_config_example) assert isinstance(loaded, QueryAgentTool) From e0430879106eb26aa8218811ab0bbac59be61b66 Mon Sep 17 00:00:00 2001 From: Kelvin Sundli Date: Fri, 17 Apr 2026 17:46:27 -0700 Subject: [PATCH 08/11] chore(codegen): refresh sync_api hash after rebase Local codegen run during the rebase produced a stale hash; regenerated to match CI. --- cognite/client/_sync_api/agents/agents.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cognite/client/_sync_api/agents/agents.py b/cognite/client/_sync_api/agents/agents.py index f7bfc2a130..4675e12ec9 100644 --- a/cognite/client/_sync_api/agents/agents.py +++ b/cognite/client/_sync_api/agents/agents.py @@ -1,6 +1,6 @@ """ =============================================================================== -990185c07f95962938122269ec34ae29 +83ebe251b32d1a7ce1ba064499beebf5 This file is auto-generated from the Async API modules, - do not edit manually! =============================================================================== """ From ab0a8847d52f293b9d0733b0f798cf88772273e4 Mon Sep 17 00:00:00 2001 From: Kelvin Sundli Date: Fri, 17 Apr 2026 18:00:19 -0700 Subject: [PATCH 09/11] test(agents): add query tool to permanent agent fixture Exercises the new QueryAgentTool end-to-end via the shared integration test agent, following review feedback on #2569. --- .../tests_integration/test_api/test_agents.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/tests/tests_integration/test_api/test_agents.py b/tests/tests_integration/test_api/test_agents.py index f1ec494f35..c147af8022 100644 --- a/tests/tests_integration/test_api/test_agents.py +++ b/tests/tests_integration/test_api/test_agents.py @@ -11,6 +11,8 @@ AskDocumentAgentToolUpsert, DataModelInfo, Message, + QueryAgentToolConfiguration, + QueryAgentToolUpsert, QueryKnowledgeGraphAgentToolConfiguration, QueryKnowledgeGraphAgentToolUpsert, QueryTimeSeriesDatapointsAgentToolUpsert, @@ -59,7 +61,21 @@ def permanent_agent(cognite_client: CogniteClient) -> Agent: instance_spaces=None, version="v2", ), - ) + ), + QueryAgentToolUpsert( + name="query_data", + description="Use this tool to run flexible queries against the data model", + configuration=QueryAgentToolConfiguration( + data_models=[ + DataModelInfo( + space="cdf_cdm", + external_id="CogniteCore", + version="v1", + view_external_ids=["CogniteAsset"], + ) + ], + ), + ), ], ) existing = cognite_client.agents.retrieve(external_ids=agent.external_id, ignore_unknown_ids=True) From 505af12a3c8620da96d57b1fbc6fe342b194c838 Mon Sep 17 00:00:00 2001 From: Kelvin Sundli Date: Fri, 17 Apr 2026 18:05:10 -0700 Subject: [PATCH 10/11] test(agents): pin permanent agent to runtime 1.2.1 and always upsert - Drop the existing-and-same-model reuse check; upsert is idempotent and this avoids silent staleness when anything other than model changes (tools, description, etc.). - Pin runtime_version to 1.2.1 so the fixture exercises the new runtime once it ships. 1.2.1 is not GA yet (expected in 2-4 weeks); until then, this fixture will fail to upsert. --- tests/tests_integration/test_api/test_agents.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/tests_integration/test_api/test_agents.py b/tests/tests_integration/test_api/test_agents.py index c147af8022..378d1d5d91 100644 --- a/tests/tests_integration/test_api/test_agents.py +++ b/tests/tests_integration/test_api/test_agents.py @@ -45,6 +45,10 @@ def permanent_agent(cognite_client: CogniteClient) -> Agent: # Or pick one manually from: # https://docs.cognite.com/cdf/atlas_ai/references/atlas_ai_agent_language_models model="gcp/claude-4.5-haiku", # Note: Should be a fast and cheap like haiku or flash + # TODO: runtime 1.2.1 is not GA yet; expected to ship in 2-4 weeks. Until then this + # fixture will fail to upsert. Leave pinned so the test starts exercising the new + # runtime the moment it's available. + runtime_version="1.2.1", tools=[ QueryKnowledgeGraphAgentToolUpsert( name="find_assets", @@ -78,9 +82,6 @@ def permanent_agent(cognite_client: CogniteClient) -> Agent: ), ], ) - existing = cognite_client.agents.retrieve(external_ids=agent.external_id, ignore_unknown_ids=True) - if existing and existing.model == agent.model: - return existing return cognite_client.agents.upsert(agent) From 7ad46c159c9ff71eca5ce75d6718ac14edc663d0 Mon Sep 17 00:00:00 2001 From: Kelvin Sundli Date: Fri, 17 Apr 2026 18:07:00 -0700 Subject: [PATCH 11/11] test(agents): revert runtime_version pin Keep the fixture on the server default. The underlying query tool needs a runtime that is not yet available; the dependency is tracked in the PR description so the merge can be held until the runtime ships. --- tests/tests_integration/test_api/test_agents.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/tests_integration/test_api/test_agents.py b/tests/tests_integration/test_api/test_agents.py index 378d1d5d91..a7a5fdafcf 100644 --- a/tests/tests_integration/test_api/test_agents.py +++ b/tests/tests_integration/test_api/test_agents.py @@ -45,10 +45,6 @@ def permanent_agent(cognite_client: CogniteClient) -> Agent: # Or pick one manually from: # https://docs.cognite.com/cdf/atlas_ai/references/atlas_ai_agent_language_models model="gcp/claude-4.5-haiku", # Note: Should be a fast and cheap like haiku or flash - # TODO: runtime 1.2.1 is not GA yet; expected to ship in 2-4 weeks. Until then this - # fixture will fail to upsert. Leave pinned so the test starts exercising the new - # runtime the moment it's available. - runtime_version="1.2.1", tools=[ QueryKnowledgeGraphAgentToolUpsert( name="find_assets",