feat(agents): add QueryAgentTool (in preview) for the new data modeling query tool type#2569
Merged
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces the QueryAgentTool, QueryAgentToolConfiguration, and QueryAgentToolUpsert classes to the Agents API, enabling flexible queries against data models. The changes include comprehensive documentation examples and unit tests for the new tool types. A review comment identifies a potential TypeError in the QueryAgentToolConfiguration._load method when handling null values for the dataModels key in API responses and provides a suggestion to improve the robustness of the parsing logic.
ks93
force-pushed
the
feat/add-query-agent-tool
branch
from
5085319 to
dda47b6
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #2569 +/- ##
=======================================
Coverage 92.76% 92.76%
=======================================
Files 480 480
Lines 48446 48502 +56
=======================================
+ Hits 44939 44991 +52
- Misses 3507 3511 +4
🚀 New features to boost your workflow:
|
ks93
added
the
waiting-for-risk-review
haakonvt
reviewed
Apr 16, 2026
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
ks93
force-pushed
the
feat/add-query-agent-tool
branch
from
9da3c4d to
6acddf4
Compare
haakonvt
removed
the
waiting-for-risk-review
haakonvt
reviewed
Apr 16, 2026
haakonvt
reviewed
Apr 17, 2026
Contributor
haakonvt
left a comment
haakonvt
left a comment
There was a problem hiding this comment.
Mostly LGTM. Should we add the query tool to the list of tools of the permanent agent in the integration test?
1 task
ks93
changed the base branch from
master
to
test/agent-tools-branch-on-expected-type
ks93
force-pushed
the
feat/add-query-agent-tool
branch
from
5cad072 to
b3d722f
Compare
Contributor
Author
ks93
added a commit
that referenced
this pull request
Apr 18, 2026
Exercises the new QueryAgentTool end-to-end via the shared integration test agent, following review feedback on #2569.
1 task
Base automatically changed from
test/agent-tools-branch-on-expected-type
to
master
April 18, 2026 06:55
haakonvt
reviewed
Apr 20, 2026
|
|
||
| >>> from cognite.client.data_classes.agents import ( | ||
| ... AgentUpsert, | ||
| ... QueryAgentToolUpsert, |
Contributor
There was a problem hiding this comment.
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.
Contributor
Author
ks93
changed the base branch from
master
to
fix/skip-flaky-hosted-extractor-jobs-tests
3 tasks
haakonvt
requested changes
Apr 21, 2026
Contributor
There was a problem hiding this comment.
@ks93 this gives me:
CogniteAPIError: Tool 'explore data' (type: query) requires a runtime version with Query Tool support
(for example 1.1.2-preview) and an experimental runtime. Current runtime version '1.1.3-preview' does
not support this tool. Please upgrade the agent's runtime version to use this tool. | code: 400 |
X-Request-ID: 1874c3ae-adf1-94a8-901c-eff9da9d9e56 | cluster: greenfield | project: forge-sandbox
The sentence Please upgrade the agent's runtime version to use this tool makes no sense 🤣
Base automatically changed from
fix/skip-flaky-hosted-extractor-jobs-tests
to
master
April 21, 2026 19:16
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.
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.
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.
Remove maturity label from docstring examples and surface missing dataModels key instead of silently defaulting to empty dict.
Covered automatically by tests/tests_unit/test_base.py.
Covered by automatic load/dump tests in test_base.py. Keep only test_load_without_configuration for the None config edge case.
Local codegen run during the rebase produced a stale hash; regenerated to match CI.
The current default runtime does not support the Query tool, so the example would fail as written. Pin to 1.1.3-preview until the Query tool lands in the default runtime.
haakonvt
force-pushed
the
feat/add-query-agent-tool
branch
from
56a90b4 to
7a08a59
Compare
haakonvt
force-pushed
the
feat/add-query-agent-tool
branch
from
33c8067 to
11cb8b0
Compare
haakonvt
changed the title
doctrino
approved these changes
Apr 22, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Note
Stacked on #2584 (CI test skip), merge that first.
Summary
"query"agent tool type, which was previously deserialized asUnknownAgentToolQueryAgentTool,QueryAgentToolUpsert,QueryAgentToolConfigurationdata_modelslist (reusingDataModelInfo) with transparent API wrapping indump/_loadTest plan
API not yet public, see
providedAtRuntimeis not yet public.