Skip to content
Merged
Show file tree
Hide file tree
Changes from 12 commits
Commits
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
9 changes: 7 additions & 2 deletions cognite/client/data_classes/agents/agent_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -166,10 +166,12 @@ class QueryKnowledgeGraphAgentToolConfiguration(WriteableCogniteResource):
Args:
data_models (Sequence[DataModelInfo] | None): The data models and views to query.
instance_spaces (InstanceSpaces | None): The instance spaces to query.
version (str | None): The version of the query generation strategy to use. A higher number does not necessarily mean a better query. Supported values are "v1" and "v2".
"""

data_models: Sequence[DataModelInfo] | None = None
instance_spaces: InstanceSpaces | None = None
data_models: Sequence[DataModelInfo] | None
instance_spaces: InstanceSpaces | None
version: Literal["v1", "v2"] | str | None

@classmethod
def _load(
Expand All @@ -186,6 +188,7 @@ def _load(
return cls(
data_models=data_models,
instance_spaces=instance_spaces,
version=resource.get("version"),
)

def dump(self, camel_case: bool = True) -> dict[str, Any]:
Expand All @@ -196,6 +199,8 @@ def dump(self, camel_case: bool = True) -> dict[str, Any]:
if self.instance_spaces:
key = "instanceSpaces" if camel_case else "instance_spaces"
result[key] = self.instance_spaces.dump(camel_case=camel_case)
if self.version:
result["version"] = self.version
return result

def as_write(self) -> QueryKnowledgeGraphAgentToolConfiguration:
Expand Down
8 changes: 6 additions & 2 deletions tests/tests_integration/test_api/test_agents.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,9 @@ def permanent_agent(cognite_client: CogniteClient) -> Agent:
version="v1",
view_external_ids=["CogniteAsset"],
)
]
],
instance_spaces=None,
version="v2",
),
)
],
Expand Down Expand Up @@ -80,7 +82,9 @@ def test_create_retrieve_update_delete_agent(self, cognite_client: CogniteClient
version="v1",
view_external_ids=["CogniteAsset"],
)
]
],
instance_spaces=None,
version="v2",
),
),
SummarizeDocumentAgentToolUpsert(
Expand Down
1 change: 1 addition & 0 deletions tests/tests_unit/test_api/test_agents.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,7 @@ def test_upsert_full(self, cognite_client: CogniteClient, mock_agent_upsert_resp
)
],
instance_spaces=InstanceSpaces(type="all"),
version="v2",
),
)
],
Expand Down