chore(agents): add query tool to permanent agent fixture

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:
+
+                >>> 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 (

cognite/client/_sync_api/agents/agents.py

@@ -1,6 +1,6 @@
 """
 ===============================================================================
-063c42ab744021733ccbdc455b150b2c
+83ebe251b32d1a7ce1ba064499beebf5
 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:
+
+                >>> 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 (

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",

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["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.

tests/tests_integration/test_api/test_agents.py

@@ -11,6 +11,8 @@
     AskDocumentAgentToolUpsert,
     DataModelInfo,
     Message,
+    QueryAgentToolConfiguration,
+    QueryAgentToolUpsert,
     QueryKnowledgeGraphAgentToolConfiguration,
     QueryKnowledgeGraphAgentToolUpsert,
     QueryTimeSeriesDatapointsAgentToolUpsert,
@@ -59,12 +61,23 @@ 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)
-    if existing and existing.model == agent.model:
-        return existing
     return cognite_client.agents.upsert(agent)