Skip to content
Merged
Changes from 6 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
5 changes: 5 additions & 0 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
version: Literal["v1", "v2"] | str | None = 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
Loading