diff --git a/cognite/client/_api/agents/agents.py b/cognite/client/_api/agents/agents.py index 03869378aa..a842fe8341 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: + + >>> 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..4675e12ec9 100644 --- a/cognite/client/_sync_api/agents/agents.py +++ b/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 ( 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..e8722b258b 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["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_integration/test_api/test_agents.py b/tests/tests_integration/test_api/test_agents.py index f1ec494f35..a7a5fdafcf 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,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) 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..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 @@ -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.""" @@ -70,23 +108,26 @@ 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 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"] 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.""" @@ -105,15 +146,17 @@ class TestAgentToolDump: (ask_document_example, AskDocumentAgentTool), (summarize_document_example, SummarizeDocumentAgentTool), (query_time_series_datapoints_example, QueryTimeSeriesDatapointsAgentTool), + (query_example, QueryAgentTool), (unknown_example, UnknownAgentTool), ], ) 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 isinstance(loaded_tool, UnknownAgentTool): + if expected_type is UnknownAgentTool: assert dumped_tool["type"] == unknown_example["type"] else: assert dumped_tool["type"] == expected_type._type @@ -150,18 +193,27 @@ class TestAgentToolUpsert: (ask_document_example, AskDocumentAgentTool), (summarize_document_example, SummarizeDocumentAgentTool), (query_time_series_datapoints_example, QueryTimeSeriesDatapointsAgentTool), + (query_example, QueryAgentTool), (unknown_example, UnknownAgentTool), ], ) 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 isinstance(loaded_tool, UnknownAgentTool): + if expected_type is UnknownAgentTool: assert dumped_tool["type"] == unknown_example["type"] else: assert dumped_tool["type"] == expected_type._type 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