feat: add QueryAgentTool for the query tool type#2569
feat: add QueryAgentTool for the query tool type#2569ks93 wants to merge 5 commits intofeat/agents/add-runtime-versionfrom
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
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 @@
## feat/agents/add-runtime-version #2569 +/- ##
===================================================================
+ Coverage 92.72% 92.76% +0.04%
===================================================================
Files 480 480
Lines 48314 48369 +55
===================================================================
+ Hits 44798 44871 +73
+ Misses 3516 3498 -18
🚀 New features to boost your workflow:
|
ks93
added
the
waiting-for-risk-review
| 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 [])] |
There was a problem hiding this comment.
Is the payload really ["dataModels"]["dataModels"]?
There was a problem hiding this comment.
Yes — it's a discriminated union in the API. The outer dataModels is an object with a type field ("manual" or "providedAtRuntime" (internal for now)), and when manual, it contains an inner dataModels array:
"configuration": {
"dataModels": {
"type": "manual",
"dataModels": [
{
"space": "cdf_idm",
"externalId": "CogniteProcessIndustries",
"version": "v1",
"viewExternalIds": ["CogniteAsset"]
}
]
},
"instanceSpaces": {"type": "manual", "spaces": ["my_space"]}
}Confirmed in cog-ai — see e.g. QueryToolConfigurationDTO → QueryToolManualDataModelsDTO → list[QueryToolDataModelItemDTO].
There was a problem hiding this comment.
@haakonvt how would you shape this instead?
| def as_write(self) -> QueryAgentToolConfiguration: | ||
| return self |
There was a problem hiding this comment.
The load method implies this is the response version, not the write version: "# API always returns dataModels, but guard against null defensively:"
|
|
||
| @dataclass | ||
| class QueryAgentTool(AgentTool): | ||
| """Agent tool for running flexible queries against data models. |
There was a problem hiding this comment.
Can you give me a link to the documentation for this tool? Would make it simpler to verify e.g. required fields and so on.
There was a problem hiding this comment.
Documentation? :hide:
It's not yet publicly documented. Adding screenshots of local docs in PR desc.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
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.
ks93
force-pushed
the
feat/add-query-agent-tool
branch
from
9da3c4d to
6acddf4
Compare
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.
haakonvt
removed
the
waiting-for-risk-review
2 participants
Note
Stacked on #2568 — merge that first.
Summary
"query"agent tool type, which was previously deserialized asUnknownAgentToolQueryAgentTool,QueryAgentToolUpsert,QueryAgentToolConfigurationdata_modelslist (reusingDataModelInfo) with transparent API wrapping indump/_loadTest plan
_AGENT_TOOL_CLS_BY_TYPEcontains"query": QueryAgentToolAPI not yet public, see
providedAtRuntimeis not yet public.