Skip to content

feat(agents): add QueryAgentTool (in preview) for the new data modeling query tool type #2569

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
haakonvt merged 12 commits into from
Apr 22, 2026
Merged

feat(agents): add QueryAgentTool (in preview) for the new data modeling query tool type #2569

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
5e5578c
test(agents): branch on expected_type in agent tool tests
ks93 Apr 17, 2026
2b07993
test(agents): assert loaded_tool type up front in dump and upsert tests
ks93 Apr 18, 2026
a506e4c
feat: add QueryAgentTool for the "query" tool type
ks93 Apr 16, 2026
838c777
fix: defensive null handling in QueryAgentToolConfiguration._load
ks93 Apr 16, 2026
d6c57a9
fix: address PR feedback on query agent tool
ks93 Apr 16, 2026
96d70a7
fix(tests): remove redundant round-trip tests for QueryAgentTool
ks93 Apr 16, 2026
349550b
fix(tests): remove redundant test_load_with_configuration
ks93 Apr 16, 2026
0780663
chore(codegen): refresh sync_api hash after rebase
ks93 Apr 18, 2026
375a55f
docs(agents): pin runtime_version in QueryAgentTool example
ks93 Apr 21, 2026
fb96946
fix wrong runtime_version in example
haakonvt Apr 21, 2026
ce493ab
simplify code example for agents.upsert
haakonvt Apr 21, 2026
11cb8b0
add a todo
haakonvt Apr 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 34 additions & 2 deletions cognite/client/_api/agents/agents.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,39 @@ async def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | A
... labels=["published"],
... tools=[find_assets_tool],
... )
>>> client.agents.upsert(agents=[agent])
>>> client.agents.upsert(agent)

Create an agent with the query tool:

>>> from cognite.client.data_classes.agents import (
... AgentUpsert,
... QueryAgentToolUpsert,
Copy link
Copy Markdown
Contributor

@haakonvt haakonvt Apr 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The current code raises due to "runtime version does not support it". I'm hesitant to merge code examples that don't work yet - I think the API change has to come first.

Copy link
Copy Markdown
Contributor Author

@ks93 ks93 Apr 21, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agree, fixed in cd05430. Filed APF-385 to drop the pinned runtime version once the Query tool lands in the default (non-preview) runtime

... 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"),
Comment thread
ks93 marked this conversation as resolved.
... ),
... )
>>> agent = AgentUpsert(
... external_id="my_agent",
... name="My Agent",
... labels=["published"],
... runtime_version="1.1.2-preview",
... tools=[query_tool],
... )
>>> client.agents.upsert(agent)

Create an agent with multiple different tools:

Expand Down Expand Up @@ -159,7 +191,7 @@ async def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | A
... ts_tool,
... ],
... )
>>> client.agents.upsert(agents=[agent])
>>> client.agents.upsert(agent)


"""
Expand Down
38 changes: 35 additions & 3 deletions cognite/client/_sync_api/agents/agents.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
"""
===============================================================================
063c42ab744021733ccbdc455b150b2c
178ce7222985b04d03174af7b7e0f525
This file is auto-generated from the Async API modules, - do not edit manually!
===============================================================================
"""
Expand Down Expand Up @@ -75,7 +75,39 @@ def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | AgentLi
... labels=["published"],
... tools=[find_assets_tool],
... )
>>> client.agents.upsert(agents=[agent])
>>> client.agents.upsert(agent)

Create an agent with the query tool:

>>> 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"],
... runtime_version="1.1.2-preview",
... tools=[query_tool],
... )
>>> client.agents.upsert(agent)

Create an agent with multiple different tools:

Expand Down Expand Up @@ -157,7 +189,7 @@ def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | AgentLi
... ts_tool,
... ],
... )
>>> client.agents.upsert(agents=[agent])
>>> client.agents.upsert(agent)
"""
return run_sync(self.__async_client.agents.upsert(agents=agents))

Expand Down
6 changes: 6 additions & 0 deletions cognite/client/data_classes/agents/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,9 @@
AskDocumentAgentToolUpsert,
DataModelInfo,
InstanceSpaces,
QueryAgentTool,
QueryAgentToolConfiguration,
QueryAgentToolUpsert,
QueryKnowledgeGraphAgentTool,
QueryKnowledgeGraphAgentToolConfiguration,
QueryKnowledgeGraphAgentToolUpsert,
Expand Down Expand Up @@ -70,6 +73,9 @@
"Message",
"MessageContent",
"MessageList",
"QueryAgentTool",
"QueryAgentToolConfiguration",
"QueryAgentToolUpsert",
"QueryKnowledgeGraphAgentTool",
"QueryKnowledgeGraphAgentToolConfiguration",
"QueryKnowledgeGraphAgentToolUpsert",
Expand Down
102 changes: 102 additions & 0 deletions cognite/client/data_classes/agents/agent_tools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -194,6 +194,45 @@ 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["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")),
)
Comment thread
ks93 marked this conversation as resolved.

def dump(self, camel_case: bool = True) -> dict[str, Any]:
result: dict[str, Any] = {}
key = "dataModels" if camel_case else "data_models"
result[key] = {
# TODO(ks93): Clarify whether "type" can take other values (cf. InstanceSpaces.type);
# if so, make it a real field
"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
Comment thread
ks93 marked this conversation as resolved.


@dataclass
class SummarizeDocumentAgentTool(AgentTool):
"""Agent tool for summarizing documents.
Expand Down Expand Up @@ -389,6 +428,69 @@ def _load(cls, resource: dict[str, Any]) -> QueryTimeSeriesDatapointsAgentToolUp
)


@dataclass
class QueryAgentTool(AgentTool):
"""Agent tool for running flexible queries against data models.
Comment thread
ks93 marked this conversation as resolved.

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.
Expand Down
53 changes: 52 additions & 1 deletion tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
from cognite.client.data_classes.agents.agent_tools import (
AgentTool,
AskDocumentAgentTool,
QueryAgentTool,
QueryAgentToolConfiguration,
QueryKnowledgeGraphAgentTool,
QueryKnowledgeGraphAgentToolConfiguration,
QueryTimeSeriesDatapointsAgentTool,
Expand Down Expand Up @@ -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
Expand All @@ -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."""
Expand All @@ -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"]
Expand All @@ -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),
],
)
Expand Down Expand Up @@ -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),
],
)
Expand All @@ -166,3 +210,10 @@ 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_without_configuration(self) -> None:
loaded = AgentTool._load(query_no_config_example)
assert isinstance(loaded, QueryAgentTool)
assert loaded.configuration is None
Loading