feat: add labels property to agents SDK (#2372)

Author: andeplane

cognite/client/_api/agents/agents.py

@@ -75,6 +75,7 @@ def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | AgentLi
                 >>> agent = AgentUpsert(
                 ...     external_id="my_agent",
                 ...     name="My Agent",
+                ...     labels=["published"],
                 ...     tools=[find_assets_tool]
                 ... )
                 >>> client.agents.upsert(agents=[agent])
@@ -150,6 +151,7 @@ def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | AgentLi
                 ...     name="My agent",
                 ...     description="An agent with many tools",
                 ...     instructions="You are a helpful assistant that can query knowledge graphs, summarize documents, answer questions about documents, and query time series data points.",
+                ...     labels=["published"],
                 ...     tools=[find_assets_tool, find_files_tool, find_time_series_tool, summarize_tool, ask_doc_tool, ts_tool]
                 ... )
                 >>> client.agents.upsert(agents=[agent])

cognite/client/data_classes/agents/agents.py

@@ -31,13 +31,15 @@ class AgentCore(WriteableCogniteResource["AgentUpsert"]):
         description (str | None): The description of the agent.
         instructions (str | None): Instructions for the agent.
         model (str | None): Name of the language model to use. For example, "azure/gpt-4o", "gcp/gemini-2.0" or "aws/claude-3.5-sonnet".
+        labels (list[str] | None): Labels for the agent. For example, ["published"] to mark an agent as published.
     """
 
     external_id: str
     name: str
     description: str | None = None
     instructions: str | None = None
     model: str | None = None
+    labels: list[str] | None = None
 
 
 class AgentUpsert(AgentCore):
@@ -50,6 +52,7 @@ class AgentUpsert(AgentCore):
         description (str | None): The human readable description of the agent.
         instructions (str | None): Instructions for the agent.
         model (str | None): Name of the language model to use. For example, "azure/gpt-4o", "gcp/gemini-2.0" or "aws/claude-3.5-sonnet".
+        labels (list[str] | None): Labels for the agent. For example, ["published"] to mark an agent as published.
         tools (Sequence[AgentToolUpsert] | None): List of tools for the agent.
 
     """
@@ -63,10 +66,16 @@ def __init__(
         description: str | None = None,
         instructions: str | None = None,
         model: str | None = None,
+        labels: list[str] | None = None,
         tools: Sequence[AgentToolUpsert] | None = None,
     ) -> None:
         super().__init__(
-            external_id=external_id, name=name, description=description, instructions=instructions, model=model
+            external_id=external_id,
+            name=name,
+            description=description,
+            instructions=instructions,
+            model=model,
+            labels=labels,
         )
         self.tools: AgentToolUpsertList | None = AgentToolUpsertList(tools) if tools is not None else None
         # This stores any unknown properties that are not part of the defined fields.
@@ -99,6 +108,7 @@ def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None =
             description=resource.get("description"),
             instructions=resource.get("instructions"),
             model=resource.get("model"),
+            labels=resource.get("labels"),
             tools=tools,
         )
         existing = set(instances.dump(camel_case=True).keys())
@@ -113,9 +123,10 @@ class Agent(AgentCore):
     Args:
         external_id (str): The external ID provided by the client. Must be unique for the resource type.
         name (str): The name of the agent, for use in user interfaces.
-        description (str | None): The human readable description of the agent.
-        instructions (str | None): Instructions for the agent.
-        model (str | None): Name of the language model to use. For example, "azure/gpt-4o", "gcp/gemini-2.0" or "aws/claude-3.5-sonnet".
+        description (str | None): The human readable description of the agent. Always present in API responses.
+        instructions (str | None): Instructions for the agent. Always present in API responses.
+        model (str | None): Name of the language model to use. For example, "azure/gpt-4o", "gcp/gemini-2.0" or "aws/claude-3.5-sonnet". Always present in API responses.
+        labels (list[str] | None): Labels for the agent. For example, ["published"] to mark an agent as published. Always present in API responses.
         tools (Sequence[AgentTool] | None): List of tools for the agent.
         created_time (int | None): The time the agent was created, in milliseconds since Thursday, 1 January 1970 00:00:00 UTC, minus leap seconds.
         last_updated_time (int | None): The time the agent was last updated, in milliseconds since Thursday, 1 January 1970 00:00:00 UTC, minus leap seconds.
@@ -134,14 +145,26 @@ def __init__(
         description: str | None = None,
         instructions: str | None = None,
         model: str | None = None,
+        labels: list[str] | None = None,
         tools: Sequence[AgentTool] | None = None,
         created_time: int | None = None,
         last_updated_time: int | None = None,
         owner_id: str | None = None,
     ) -> None:
         super().__init__(
-            external_id=external_id, name=name, description=description, instructions=instructions, model=model
+            external_id=external_id,
+            name=name,
+            description=description,
+            instructions=instructions,
+            model=model,
+            labels=labels,
         )
+        # These fields are always present in API responses, but optional when creating.
+        # Force the type to be non-optional for read instances.
+        self.description: str = description  # type: ignore[assignment]
+        self.instructions: str = instructions  # type: ignore[assignment]
+        self.model: str = model  # type: ignore[assignment]
+        self.labels: list[str] = labels  # type: ignore[assignment]
         self.tools: AgentToolList | None = AgentToolList(tools) if tools is not None else None
         self.created_time = created_time
         self.last_updated_time = last_updated_time
@@ -166,6 +189,7 @@ def as_write(self) -> AgentUpsert:
             description=self.description,
             instructions=self.instructions,
             model=self.model,
+            labels=self.labels,
             tools=[tool.as_write() for tool in self.tools] if self.tools else None,
         )
 
@@ -183,6 +207,7 @@ def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None =
             description=resource.get("description"),
             instructions=resource.get("instructions"),
             model=resource.get("model"),
+            labels=resource.get("labels"),
             tools=tools,
             created_time=resource.get("createdTime"),
             last_updated_time=resource.get("lastUpdatedTime"),

tests/tests_unit/test_api/test_agents.py

@@ -25,6 +25,7 @@ def agent_response_body() -> dict:
                 "description": "Description 1",
                 "instructions": "Instructions 1",
                 "model": "vendor/model_1",
+                "labels": ["published"],
                 "tools": [
                     {
                         "name": "tool_1",
@@ -176,3 +177,22 @@ def test_upsert_ensure_retry(
         created = cognite_client.agents.upsert(AgentUpsert(external_id="agent_1", name="Agent 1"))
 
         assert isinstance(created, Agent)
+
+    def test_upsert_with_labels(self, cognite_client: CogniteClient, mock_agent_upsert_response: MagicMock) -> None:
+        agent_write = AgentUpsert(
+            external_id="agent_1",
+            name="Agent 1",
+            labels=["published"],
+        )
+        created_agent = cognite_client.agents.upsert(agent_write)
+        assert isinstance(created_agent, Agent)
+        assert created_agent.external_id == "agent_1"
+        assert created_agent.labels == ["published"]
+
+    def test_retrieve_agent_with_labels(
+        self, cognite_client: CogniteClient, mock_agent_retrieve_response: MagicMock
+    ) -> None:
+        retrieved_agent = cognite_client.agents.retrieve("agent_1")
+        assert isinstance(retrieved_agent, Agent)
+        assert retrieved_agent.external_id == "agent_1"
+        assert retrieved_agent.labels == ["published"]

tests/tests_unit/test_data_classes/test_agents/test_agents.py

@@ -79,12 +79,35 @@ def test_load_dump_maintain_unknown_properties(self) -> None:
             "externalId": "test_agent",
             "name": "Test Agent",
             "unknownProperty": "unknown_value",
-            "labels": ["published"],
         }
         dumped = AgentUpsert._load(agent_data).dump(camel_case=True)
 
         assert dumped == agent_data
 
+    def test_create_agent_with_labels(self) -> None:
+        """Test creating an agent with labels."""
+        agent = AgentUpsert(
+            external_id="test_agent",
+            name="Test Agent",
+            labels=["published"],
+        )
+        assert agent.labels == ["published"]
+
+        dumped = agent.dump(camel_case=True)
+        assert dumped["labels"] == ["published"]
+
+    def test_labels_forward_compatibility(self) -> None:
+        """Test forward compatibility with future label values."""
+        agent = AgentUpsert(
+            external_id="test_agent",
+            name="Test Agent",
+            labels=["published", "charts"],
+        )
+        assert agent.labels == ["published", "charts"]
+
+        dumped = agent.dump(camel_case=True)
+        assert dumped["labels"] == ["published", "charts"]
+
     def test_as_write(self) -> None:
         agent_upsert = AgentUpsert(
             external_id="test_agent",
@@ -185,6 +208,7 @@ def test_as_write(self) -> None:
             description="A test agent",
             instructions="Test instructions",
             model="gpt-4",
+            labels=["published"],
             tools=[SummarizeDocumentAgentTool(name="test_tool", description="A test tool")],
         )
 
@@ -195,10 +219,25 @@ def test_as_write(self) -> None:
         assert write_agent.description == agent.description
         assert write_agent.instructions == agent.instructions
         assert write_agent.model == agent.model
+        assert write_agent.labels == agent.labels
+        assert write_agent.labels == ["published"]
         assert len(write_agent.tools) == 1
         assert isinstance(write_agent.tools[0], AgentToolUpsert)
         assert write_agent.tools[0].name == "test_tool"
 
+    def test_agent_labels_forward_compatibility(self) -> None:
+        """Test forward compatibility with future label values on Agent."""
+        agent = Agent(
+            external_id="test_agent",
+            name="Test Agent",
+            labels=["published", "charts"],
+        )
+        assert agent.labels == ["published", "charts"]
+
+        # Test that as_write() preserves labels
+        write_agent = agent.as_write()
+        assert write_agent.labels == ["published", "charts"]
+
 
 class TestAgentList:
     def test_as_write(self) -> None: