From af866a18739fd0de15a831178e4e2eb3996180c6 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 19 Aug 2025 22:01:45 +0000 Subject: [PATCH 1/4] Add AgentSession for simplified multi-turn agent conversations Co-authored-by: kelvin.sundli --- cognite/client/_api/agents/agents.py | 38 ++- .../client/data_classes/agents/__init__.py | 2 + cognite/client/data_classes/agents/agents.py | 34 +++ cognite/client/data_classes/agents/chat.py | 76 +++++ .../test_api/test_agents_session.py | 261 ++++++++++++++++++ 5 files changed, 410 insertions(+), 1 deletion(-) create mode 100644 tests/tests_unit/test_api/test_agents_session.py diff --git a/cognite/client/_api/agents/agents.py b/cognite/client/_api/agents/agents.py index 8805851d3d..1143976767 100644 --- a/cognite/client/_api/agents/agents.py +++ b/cognite/client/_api/agents/agents.py @@ -5,7 +5,7 @@ from cognite.client._api_client import APIClient from cognite.client.data_classes.agents import Agent, AgentList, AgentUpsert -from cognite.client.data_classes.agents.chat import AgentChatResponse, Message, MessageList +from cognite.client.data_classes.agents.chat import AgentChatResponse, AgentSession, Message, MessageList from cognite.client.utils._experimental import FeaturePreviewWarning from cognite.client.utils._identifier import IdentifierSequence from cognite.client.utils.useful_types import SequenceNotStr @@ -313,3 +313,39 @@ def chat( ) return AgentChatResponse._load(response.json(), cognite_client=self._cognite_client) + + def start_session(self, agent_id: str, cursor: str | None = None) -> AgentSession: + """Start a new chat session with an agent. + + This creates a session object that automatically manages cursor state across + multiple chat interactions, providing a more convenient interface for + multi-turn conversations. + + Args: + agent_id (str): External ID that uniquely identifies the agent. + cursor (str | None): Optional cursor to continue an existing conversation. + If None, starts a fresh session. + + Returns: + AgentSession: A session object for chatting with the agent. + + Examples: + + Start a new session and have a conversation: + + >>> from cognite.client import CogniteClient + >>> from cognite.client.data_classes.agents import Message + >>> client = CogniteClient() + >>> session = client.agents.start_session("my_agent") + >>> response = session.chat("Hello, how can you help me?") + >>> print(response.text) + >>> followup = session.chat("Tell me more about that") + >>> print(followup.text) + + Continue an existing session with a cursor: + + >>> session = client.agents.start_session("my_agent", cursor="existing_cursor_123") + >>> response = session.chat("Continue our previous conversation") + """ + self._warnings.warn() + return AgentSession(agent_id=agent_id, cognite_client=self._cognite_client, cursor=cursor) diff --git a/cognite/client/data_classes/agents/__init__.py b/cognite/client/data_classes/agents/__init__.py index a9931b6240..b269b83665 100644 --- a/cognite/client/data_classes/agents/__init__.py +++ b/cognite/client/data_classes/agents/__init__.py @@ -26,6 +26,7 @@ AgentMessage, AgentMessageList, AgentReasoningItem, + AgentSession, Message, MessageContent, MessageList, @@ -41,6 +42,7 @@ "AgentMessage", "AgentMessageList", "AgentReasoningItem", + "AgentSession", "AgentTool", "AgentToolList", "AgentToolUpsert", diff --git a/cognite/client/data_classes/agents/agents.py b/cognite/client/data_classes/agents/agents.py index 43d1113b35..bdfe7f39c8 100644 --- a/cognite/client/data_classes/agents/agents.py +++ b/cognite/client/data_classes/agents/agents.py @@ -6,6 +6,7 @@ if TYPE_CHECKING: from cognite.client import CogniteClient + from cognite.client.data_classes.agents.chat import AgentSession from cognite.client.data_classes._base import ( CogniteResourceList, @@ -192,6 +193,39 @@ def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = instance._unknown_properties = {key: value for key, value in resource.items() if key not in existing} return instance + def start_session(self, cursor: str | None = None) -> AgentSession: + """Start a new chat session with this agent. + + This creates a session object that automatically manages cursor state across + multiple chat interactions, providing a more convenient interface for + multi-turn conversations. + + Args: + cursor (str | None): Optional cursor to continue an existing conversation. + If None, starts a fresh session. + + Returns: + AgentSession: A session object for chatting with this agent. + + Examples: + + Start a session from an agent instance: + + >>> from cognite.client import CogniteClient + >>> client = CogniteClient() + >>> agent = client.agents.retrieve("my_agent") + >>> session = agent.start_session() + >>> response = session.chat("Hello!") + >>> print(response.text) + """ + # Import here to avoid circular imports + from cognite.client.data_classes.agents.chat import AgentSession + + if not hasattr(self, '_cognite_client') or self._cognite_client is None: + raise ValueError("Agent instance must have a cognite_client to start a session") + + return AgentSession(agent_id=self.external_id, cognite_client=self._cognite_client, cursor=cursor) + class AgentUpsertList(CogniteResourceList[AgentUpsert], ExternalIDTransformerMixin): _RESOURCE = AgentUpsert diff --git a/cognite/client/data_classes/agents/chat.py b/cognite/client/data_classes/agents/chat.py index 1a8654ae2d..d4059760e9 100644 --- a/cognite/client/data_classes/agents/chat.py +++ b/cognite/client/data_classes/agents/chat.py @@ -3,6 +3,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar, Literal +from collections.abc import Sequence from cognite.client.data_classes._base import CogniteObject, CogniteResource, CogniteResourceList from cognite.client.utils._text import convert_all_keys_to_camel_case @@ -271,3 +272,78 @@ def _load(cls, data: dict[str, Any], cognite_client: CogniteClient | None = None ) return instance + + +class AgentSession: + """A session for chatting with an agent that manages cursor state automatically. + + This class provides a higher-level interface for chatting with agents by automatically + managing the conversation cursor state across multiple chat interactions. + + Args: + agent_id (str): The external ID of the agent to chat with. + cognite_client (CogniteClient): The Cognite client instance. + cursor (str | None): Optional initial cursor for continuing an existing session. + """ + + def __init__(self, agent_id: str, cognite_client: CogniteClient, cursor: str | None = None) -> None: + self.agent_id = agent_id + self._cognite_client = cognite_client + self._cursor = cursor + + def chat(self, messages: str | Message | Sequence[Message]) -> AgentChatResponse: + """Send a message or messages to the agent and return the response. + + The cursor state is automatically managed and updated after each interaction. + + Args: + messages (str | Message | Sequence[Message]): The message(s) to send to the agent. + If a string is provided, it will be converted to a Message. + + Returns: + AgentChatResponse: The response from the agent. + + Examples: + + Simple string message: + + >>> session = client.agents.start_session("my_agent") + >>> response = session.chat("Hello, how can you help me?") + >>> print(response.text) + + Follow-up message (cursor automatically managed): + + >>> followup = session.chat("Tell me more about that") + >>> print(followup.text) + + Multiple messages at once: + + >>> response = session.chat([ + ... Message("Find the temperature sensors"), + ... Message("Show me their recent data") + ... ]) + """ + # Convert string to Message if needed + if isinstance(messages, str): + messages = Message(messages) + + # Call the underlying agents.chat method with current cursor + response = self._cognite_client.agents.chat( + agent_id=self.agent_id, + messages=messages, + cursor=self._cursor + ) + + # Update the cursor for the next interaction + self._cursor = response.cursor + + return response + + @property + def cursor(self) -> str | None: + """Get the current cursor state for this session. + + Returns: + str | None: The current cursor, or None if no conversation has taken place yet. + """ + return self._cursor diff --git a/tests/tests_unit/test_api/test_agents_session.py b/tests/tests_unit/test_api/test_agents_session.py new file mode 100644 index 0000000000..1c36ec1bd9 --- /dev/null +++ b/tests/tests_unit/test_api/test_agents_session.py @@ -0,0 +1,261 @@ +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +from cognite.client import CogniteClient +from cognite.client.data_classes.agents import Agent, AgentSession, Message +from cognite.client.data_classes.agents.chat import ( + AgentChatResponse, + AgentMessage, + TextContent, +) + + +@pytest.fixture +def chat_response_body() -> dict: + return { + "agentId": "my_agent", + "response": { + "cursor": "cursor_12345", + "messages": [ + { + "content": { + "text": "Hello! How can I help you today?", + "type": "text", + }, + "role": "agent", + } + ], + "type": "result", + }, + } + + +@pytest.fixture +def followup_response_body() -> dict: + return { + "agentId": "my_agent", + "response": { + "cursor": "cursor_67890", + "messages": [ + { + "content": { + "text": "I can help you with data analysis, finding assets, and more!", + "type": "text", + }, + "role": "agent", + } + ], + "type": "result", + }, + } + + +class TestAgentSession: + def test_start_session_from_agents_api(self, cognite_client: CogniteClient) -> None: + """Test starting a session from the agents API.""" + session = cognite_client.agents.start_session("my_agent") + + assert isinstance(session, AgentSession) + assert session.agent_id == "my_agent" + assert session.cursor is None + assert session._cognite_client is cognite_client + + def test_start_session_with_cursor(self, cognite_client: CogniteClient) -> None: + """Test starting a session with an existing cursor.""" + initial_cursor = "existing_cursor_123" + session = cognite_client.agents.start_session("my_agent", cursor=initial_cursor) + + assert session.agent_id == "my_agent" + assert session.cursor == initial_cursor + + def test_session_chat_with_string(self, cognite_client: CogniteClient, chat_response_body: dict) -> None: + """Test chatting with a string message.""" + # Mock the API response + cognite_client.agents._post = MagicMock(return_value=MagicMock(json=lambda: chat_response_body)) + + session = cognite_client.agents.start_session("my_agent") + response = session.chat("Hello") + + # Verify the underlying chat method was called correctly + cognite_client.agents._post.assert_called_once() + call_args = cognite_client.agents._post.call_args + assert call_args[1]["json"]["agentId"] == "my_agent" + assert call_args[1]["json"]["messages"][0]["content"]["text"] == "Hello" + assert call_args[1]["json"]["messages"][0]["role"] == "user" + assert "cursor" not in call_args[1]["json"] # No cursor on first message + + # Verify response + assert isinstance(response, AgentChatResponse) + assert response.text == "Hello! How can I help you today?" + assert session.cursor == "cursor_12345" # Cursor should be updated + + def test_session_chat_with_message_object(self, cognite_client: CogniteClient, chat_response_body: dict) -> None: + """Test chatting with a Message object.""" + cognite_client.agents._post = MagicMock(return_value=MagicMock(json=lambda: chat_response_body)) + + session = cognite_client.agents.start_session("my_agent") + message = Message("Hello agent") + response = session.chat(message) + + # Verify the request + call_args = cognite_client.agents._post.call_args + assert call_args[1]["json"]["messages"][0]["content"]["text"] == "Hello agent" + + # Verify response and cursor update + assert response.text == "Hello! How can I help you today?" + assert session.cursor == "cursor_12345" + + def test_session_cursor_management( + self, + cognite_client: CogniteClient, + chat_response_body: dict, + followup_response_body: dict + ) -> None: + """Test that cursor is automatically managed across multiple interactions.""" + # Set up mock responses for two consecutive calls + responses = [ + MagicMock(json=lambda: chat_response_body), + MagicMock(json=lambda: followup_response_body), + ] + cognite_client.agents._post = MagicMock(side_effect=responses) + + session = cognite_client.agents.start_session("my_agent") + + # First interaction + response1 = session.chat("Hello") + assert session.cursor == "cursor_12345" + + # Second interaction - should include cursor from first response + response2 = session.chat("Tell me more") + assert session.cursor == "cursor_67890" # Updated to new cursor + + # Verify the second call included the cursor from the first response + second_call_args = cognite_client.agents._post.call_args_list[1] + assert second_call_args[1]["json"]["cursor"] == "cursor_12345" + assert second_call_args[1]["json"]["messages"][0]["content"]["text"] == "Tell me more" + + def test_session_chat_with_multiple_messages(self, cognite_client: CogniteClient, chat_response_body: dict) -> None: + """Test chatting with multiple messages at once.""" + cognite_client.agents._post = MagicMock(return_value=MagicMock(json=lambda: chat_response_body)) + + session = cognite_client.agents.start_session("my_agent") + messages = [ + Message("Find temperature sensors"), + Message("Show me their recent data") + ] + response = session.chat(messages) + + # Verify multiple messages were sent + call_args = cognite_client.agents._post.call_args + assert len(call_args[1]["json"]["messages"]) == 2 + assert call_args[1]["json"]["messages"][0]["content"]["text"] == "Find temperature sensors" + assert call_args[1]["json"]["messages"][1]["content"]["text"] == "Show me their recent data" + + +class TestAgentStartSession: + def test_agent_start_session(self, cognite_client: CogniteClient) -> None: + """Test starting a session from an Agent instance.""" + # Create an agent instance with a cognite client + agent = Agent( + external_id="my_agent", + name="My Agent", + description="Test agent" + ) + agent._cognite_client = cognite_client + + session = agent.start_session() + + assert isinstance(session, AgentSession) + assert session.agent_id == "my_agent" + assert session.cursor is None + assert session._cognite_client is cognite_client + + def test_agent_start_session_with_cursor(self, cognite_client: CogniteClient) -> None: + """Test starting a session from an Agent instance with a cursor.""" + agent = Agent( + external_id="my_agent", + name="My Agent" + ) + agent._cognite_client = cognite_client + + initial_cursor = "existing_cursor_456" + session = agent.start_session(cursor=initial_cursor) + + assert session.agent_id == "my_agent" + assert session.cursor == initial_cursor + + def test_agent_start_session_without_client_raises_error(self) -> None: + """Test that starting a session without a cognite client raises an error.""" + agent = Agent( + external_id="my_agent", + name="My Agent" + ) + # Don't set _cognite_client + + with pytest.raises(ValueError, match="Agent instance must have a cognite_client to start a session"): + agent.start_session() + + def test_agent_start_session_with_none_client_raises_error(self) -> None: + """Test that starting a session with None cognite client raises an error.""" + agent = Agent( + external_id="my_agent", + name="My Agent" + ) + agent._cognite_client = None + + with pytest.raises(ValueError, match="Agent instance must have a cognite_client to start a session"): + agent.start_session() + + +class TestIntegrationWorkflow: + def test_complete_workflow( + self, + cognite_client: CogniteClient, + chat_response_body: dict, + followup_response_body: dict + ) -> None: + """Test the complete workflow as described in the user requirements.""" + # Set up mock responses + responses = [ + MagicMock(json=lambda: chat_response_body), + MagicMock(json=lambda: followup_response_body), + ] + cognite_client.agents._post = MagicMock(side_effect=responses) + + # Mock the retrieve method to return an agent with cognite_client + mock_agent = Agent(external_id="my-agent", name="Test Agent") + mock_agent._cognite_client = cognite_client + cognite_client.agents.retrieve = MagicMock(return_value=mock_agent) + + # Test the workflow from the user requirements + + # 1. Start session from agents API + session = cognite_client.agents.start_session(agent_id="my-agent", cursor=None) + + # 2. First chat + resp = session.chat(Message("hello")) + assert resp.text == "Hello! How can I help you today?" + + # 3. Follow-up question with automatic cursor management + followup_response = session.chat(Message("nice to meet you")) + assert followup_response.text == "I can help you with data analysis, finding assets, and more!" + + # 4. Test starting session from agent instance + agent = cognite_client.agents.retrieve("my-agent") + agent_session = agent.start_session() # also with optional cursor + + assert isinstance(agent_session, AgentSession) + assert agent_session.agent_id == "my-agent" + + # Verify that cursors were managed correctly + first_call_args = cognite_client.agents._post.call_args_list[0] + second_call_args = cognite_client.agents._post.call_args_list[1] + + # First call should have no cursor + assert "cursor" not in first_call_args[1]["json"] + + # Second call should have cursor from first response + assert second_call_args[1]["json"]["cursor"] == "cursor_12345" \ No newline at end of file From 4cf4e70334d48f1276bee4e95f61fc33db75c3dc Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 19 Aug 2025 22:22:52 +0000 Subject: [PATCH 2/4] Refactor agent chat to require Message object and improve error handling Co-authored-by: kelvin.sundli --- cognite/client/data_classes/agents/agents.py | 6 ++-- cognite/client/data_classes/agents/chat.py | 16 ++++------ .../test_api/test_agents_session.py | 31 ++++++------------- 3 files changed, 19 insertions(+), 34 deletions(-) diff --git a/cognite/client/data_classes/agents/agents.py b/cognite/client/data_classes/agents/agents.py index bdfe7f39c8..d538648592 100644 --- a/cognite/client/data_classes/agents/agents.py +++ b/cognite/client/data_classes/agents/agents.py @@ -212,17 +212,19 @@ def start_session(self, cursor: str | None = None) -> AgentSession: Start a session from an agent instance: >>> from cognite.client import CogniteClient + >>> from cognite.client.data_classes.agents import Message >>> client = CogniteClient() >>> agent = client.agents.retrieve("my_agent") >>> session = agent.start_session() - >>> response = session.chat("Hello!") + >>> response = session.chat(Message("Hello!")) >>> print(response.text) """ # Import here to avoid circular imports from cognite.client.data_classes.agents.chat import AgentSession + from cognite.client.exceptions import CogniteMissingClientError if not hasattr(self, '_cognite_client') or self._cognite_client is None: - raise ValueError("Agent instance must have a cognite_client to start a session") + raise CogniteMissingClientError(self) return AgentSession(agent_id=self.external_id, cognite_client=self._cognite_client, cursor=cursor) diff --git a/cognite/client/data_classes/agents/chat.py b/cognite/client/data_classes/agents/chat.py index d4059760e9..1b9990a58c 100644 --- a/cognite/client/data_classes/agents/chat.py +++ b/cognite/client/data_classes/agents/chat.py @@ -291,29 +291,29 @@ def __init__(self, agent_id: str, cognite_client: CogniteClient, cursor: str | N self._cognite_client = cognite_client self._cursor = cursor - def chat(self, messages: str | Message | Sequence[Message]) -> AgentChatResponse: + def chat(self, messages: Message | Sequence[Message]) -> AgentChatResponse: """Send a message or messages to the agent and return the response. The cursor state is automatically managed and updated after each interaction. Args: - messages (str | Message | Sequence[Message]): The message(s) to send to the agent. - If a string is provided, it will be converted to a Message. + messages (Message | Sequence[Message]): The message(s) to send to the agent. Returns: AgentChatResponse: The response from the agent. Examples: - Simple string message: + Simple message: + >>> from cognite.client.data_classes.agents import Message >>> session = client.agents.start_session("my_agent") - >>> response = session.chat("Hello, how can you help me?") + >>> response = session.chat(Message("Hello, how can you help me?")) >>> print(response.text) Follow-up message (cursor automatically managed): - >>> followup = session.chat("Tell me more about that") + >>> followup = session.chat(Message("Tell me more about that")) >>> print(followup.text) Multiple messages at once: @@ -323,10 +323,6 @@ def chat(self, messages: str | Message | Sequence[Message]) -> AgentChatResponse ... Message("Show me their recent data") ... ]) """ - # Convert string to Message if needed - if isinstance(messages, str): - messages = Message(messages) - # Call the underlying agents.chat method with current cursor response = self._cognite_client.agents.chat( agent_id=self.agent_id, diff --git a/tests/tests_unit/test_api/test_agents_session.py b/tests/tests_unit/test_api/test_agents_session.py index 1c36ec1bd9..9e5654c141 100644 --- a/tests/tests_unit/test_api/test_agents_session.py +++ b/tests/tests_unit/test_api/test_agents_session.py @@ -11,6 +11,7 @@ AgentMessage, TextContent, ) +from cognite.client.exceptions import CogniteMissingClientError @pytest.fixture @@ -71,13 +72,13 @@ def test_start_session_with_cursor(self, cognite_client: CogniteClient) -> None: assert session.agent_id == "my_agent" assert session.cursor == initial_cursor - def test_session_chat_with_string(self, cognite_client: CogniteClient, chat_response_body: dict) -> None: - """Test chatting with a string message.""" + def test_session_chat_with_message(self, cognite_client: CogniteClient, chat_response_body: dict) -> None: + """Test chatting with a Message object.""" # Mock the API response cognite_client.agents._post = MagicMock(return_value=MagicMock(json=lambda: chat_response_body)) session = cognite_client.agents.start_session("my_agent") - response = session.chat("Hello") + response = session.chat(Message("Hello")) # Verify the underlying chat method was called correctly cognite_client.agents._post.assert_called_once() @@ -92,21 +93,7 @@ def test_session_chat_with_string(self, cognite_client: CogniteClient, chat_resp assert response.text == "Hello! How can I help you today?" assert session.cursor == "cursor_12345" # Cursor should be updated - def test_session_chat_with_message_object(self, cognite_client: CogniteClient, chat_response_body: dict) -> None: - """Test chatting with a Message object.""" - cognite_client.agents._post = MagicMock(return_value=MagicMock(json=lambda: chat_response_body)) - - session = cognite_client.agents.start_session("my_agent") - message = Message("Hello agent") - response = session.chat(message) - - # Verify the request - call_args = cognite_client.agents._post.call_args - assert call_args[1]["json"]["messages"][0]["content"]["text"] == "Hello agent" - - # Verify response and cursor update - assert response.text == "Hello! How can I help you today?" - assert session.cursor == "cursor_12345" + def test_session_cursor_management( self, @@ -125,11 +112,11 @@ def test_session_cursor_management( session = cognite_client.agents.start_session("my_agent") # First interaction - response1 = session.chat("Hello") + response1 = session.chat(Message("Hello")) assert session.cursor == "cursor_12345" # Second interaction - should include cursor from first response - response2 = session.chat("Tell me more") + response2 = session.chat(Message("Tell me more")) assert session.cursor == "cursor_67890" # Updated to new cursor # Verify the second call included the cursor from the first response @@ -195,7 +182,7 @@ def test_agent_start_session_without_client_raises_error(self) -> None: ) # Don't set _cognite_client - with pytest.raises(ValueError, match="Agent instance must have a cognite_client to start a session"): + with pytest.raises(CogniteMissingClientError): agent.start_session() def test_agent_start_session_with_none_client_raises_error(self) -> None: @@ -206,7 +193,7 @@ def test_agent_start_session_with_none_client_raises_error(self) -> None: ) agent._cognite_client = None - with pytest.raises(ValueError, match="Agent instance must have a cognite_client to start a session"): + with pytest.raises(CogniteMissingClientError): agent.start_session() From f76da10d56db0951a9be0b0a81a6c7515e814752 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 20 Aug 2025 04:31:20 +0000 Subject: [PATCH 3/4] Update agent session chat method to require Message object Co-authored-by: kelvin.sundli --- cognite/client/_api/agents/agents.py | 24 +++---- cognite/client/data_classes/agents/agents.py | 16 ++--- cognite/client/data_classes/agents/chat.py | 34 ++++----- .../test_api/test_agents_session.py | 72 +++++++++---------- 4 files changed, 73 insertions(+), 73 deletions(-) diff --git a/cognite/client/_api/agents/agents.py b/cognite/client/_api/agents/agents.py index 1143976767..60d9c1b6e7 100644 --- a/cognite/client/_api/agents/agents.py +++ b/cognite/client/_api/agents/agents.py @@ -316,36 +316,36 @@ def chat( def start_session(self, agent_id: str, cursor: str | None = None) -> AgentSession: """Start a new chat session with an agent. - + This creates a session object that automatically manages cursor state across - multiple chat interactions, providing a more convenient interface for + multiple chat interactions, providing a more convenient interface for multi-turn conversations. - + Args: agent_id (str): External ID that uniquely identifies the agent. cursor (str | None): Optional cursor to continue an existing conversation. If None, starts a fresh session. - + Returns: AgentSession: A session object for chatting with the agent. - + Examples: - + Start a new session and have a conversation: - + >>> from cognite.client import CogniteClient >>> from cognite.client.data_classes.agents import Message >>> client = CogniteClient() >>> session = client.agents.start_session("my_agent") - >>> response = session.chat("Hello, how can you help me?") + >>> response = session.chat(Message("Hello, how can you help me?")) >>> print(response.text) - >>> followup = session.chat("Tell me more about that") + >>> followup = session.chat(Message("Tell me more about that")) >>> print(followup.text) - + Continue an existing session with a cursor: - + >>> session = client.agents.start_session("my_agent", cursor="existing_cursor_123") - >>> response = session.chat("Continue our previous conversation") + >>> response = session.chat(Message("Continue our previous conversation")) """ self._warnings.warn() return AgentSession(agent_id=agent_id, cognite_client=self._cognite_client, cursor=cursor) diff --git a/cognite/client/data_classes/agents/agents.py b/cognite/client/data_classes/agents/agents.py index d538648592..faa01d7e1d 100644 --- a/cognite/client/data_classes/agents/agents.py +++ b/cognite/client/data_classes/agents/agents.py @@ -195,22 +195,22 @@ def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = def start_session(self, cursor: str | None = None) -> AgentSession: """Start a new chat session with this agent. - + This creates a session object that automatically manages cursor state across - multiple chat interactions, providing a more convenient interface for + multiple chat interactions, providing a more convenient interface for multi-turn conversations. - + Args: cursor (str | None): Optional cursor to continue an existing conversation. If None, starts a fresh session. - + Returns: AgentSession: A session object for chatting with this agent. - + Examples: - + Start a session from an agent instance: - + >>> from cognite.client import CogniteClient >>> from cognite.client.data_classes.agents import Message >>> client = CogniteClient() @@ -222,7 +222,7 @@ def start_session(self, cursor: str | None = None) -> AgentSession: # Import here to avoid circular imports from cognite.client.data_classes.agents.chat import AgentSession from cognite.client.exceptions import CogniteMissingClientError - + if not hasattr(self, '_cognite_client') or self._cognite_client is None: raise CogniteMissingClientError(self) diff --git a/cognite/client/data_classes/agents/chat.py b/cognite/client/data_classes/agents/chat.py index 1b9990a58c..9653b53644 100644 --- a/cognite/client/data_classes/agents/chat.py +++ b/cognite/client/data_classes/agents/chat.py @@ -276,48 +276,48 @@ def _load(cls, data: dict[str, Any], cognite_client: CogniteClient | None = None class AgentSession: """A session for chatting with an agent that manages cursor state automatically. - + This class provides a higher-level interface for chatting with agents by automatically managing the conversation cursor state across multiple chat interactions. - + Args: agent_id (str): The external ID of the agent to chat with. cognite_client (CogniteClient): The Cognite client instance. cursor (str | None): Optional initial cursor for continuing an existing session. """ - + def __init__(self, agent_id: str, cognite_client: CogniteClient, cursor: str | None = None) -> None: self.agent_id = agent_id self._cognite_client = cognite_client self._cursor = cursor - + def chat(self, messages: Message | Sequence[Message]) -> AgentChatResponse: """Send a message or messages to the agent and return the response. - + The cursor state is automatically managed and updated after each interaction. - + Args: messages (Message | Sequence[Message]): The message(s) to send to the agent. - + Returns: AgentChatResponse: The response from the agent. - + Examples: - + Simple message: - + >>> from cognite.client.data_classes.agents import Message >>> session = client.agents.start_session("my_agent") >>> response = session.chat(Message("Hello, how can you help me?")) >>> print(response.text) - + Follow-up message (cursor automatically managed): - + >>> followup = session.chat(Message("Tell me more about that")) >>> print(followup.text) - + Multiple messages at once: - + >>> response = session.chat([ ... Message("Find the temperature sensors"), ... Message("Show me their recent data") @@ -332,13 +332,13 @@ def chat(self, messages: Message | Sequence[Message]) -> AgentChatResponse: # Update the cursor for the next interaction self._cursor = response.cursor - + return response - + @property def cursor(self) -> str | None: """Get the current cursor state for this session. - + Returns: str | None: The current cursor, or None if no conversation has taken place yet. """ diff --git a/tests/tests_unit/test_api/test_agents_session.py b/tests/tests_unit/test_api/test_agents_session.py index 9e5654c141..eea48803eb 100644 --- a/tests/tests_unit/test_api/test_agents_session.py +++ b/tests/tests_unit/test_api/test_agents_session.py @@ -58,7 +58,7 @@ class TestAgentSession: def test_start_session_from_agents_api(self, cognite_client: CogniteClient) -> None: """Test starting a session from the agents API.""" session = cognite_client.agents.start_session("my_agent") - + assert isinstance(session, AgentSession) assert session.agent_id == "my_agent" assert session.cursor is None @@ -68,7 +68,7 @@ def test_start_session_with_cursor(self, cognite_client: CogniteClient) -> None: """Test starting a session with an existing cursor.""" initial_cursor = "existing_cursor_123" session = cognite_client.agents.start_session("my_agent", cursor=initial_cursor) - + assert session.agent_id == "my_agent" assert session.cursor == initial_cursor @@ -76,10 +76,10 @@ def test_session_chat_with_message(self, cognite_client: CogniteClient, chat_res """Test chatting with a Message object.""" # Mock the API response cognite_client.agents._post = MagicMock(return_value=MagicMock(json=lambda: chat_response_body)) - + session = cognite_client.agents.start_session("my_agent") response = session.chat(Message("Hello")) - + # Verify the underlying chat method was called correctly cognite_client.agents._post.assert_called_once() call_args = cognite_client.agents._post.call_args @@ -87,7 +87,7 @@ def test_session_chat_with_message(self, cognite_client: CogniteClient, chat_res assert call_args[1]["json"]["messages"][0]["content"]["text"] == "Hello" assert call_args[1]["json"]["messages"][0]["role"] == "user" assert "cursor" not in call_args[1]["json"] # No cursor on first message - + # Verify response assert isinstance(response, AgentChatResponse) assert response.text == "Hello! How can I help you today?" @@ -96,9 +96,9 @@ def test_session_chat_with_message(self, cognite_client: CogniteClient, chat_res def test_session_cursor_management( - self, - cognite_client: CogniteClient, - chat_response_body: dict, + self, + cognite_client: CogniteClient, + chat_response_body: dict, followup_response_body: dict ) -> None: """Test that cursor is automatically managed across multiple interactions.""" @@ -108,17 +108,17 @@ def test_session_cursor_management( MagicMock(json=lambda: followup_response_body), ] cognite_client.agents._post = MagicMock(side_effect=responses) - + session = cognite_client.agents.start_session("my_agent") - + # First interaction - response1 = session.chat(Message("Hello")) + session.chat(Message("Hello")) assert session.cursor == "cursor_12345" - + # Second interaction - should include cursor from first response - response2 = session.chat(Message("Tell me more")) + session.chat(Message("Tell me more")) assert session.cursor == "cursor_67890" # Updated to new cursor - + # Verify the second call included the cursor from the first response second_call_args = cognite_client.agents._post.call_args_list[1] assert second_call_args[1]["json"]["cursor"] == "cursor_12345" @@ -127,14 +127,14 @@ def test_session_cursor_management( def test_session_chat_with_multiple_messages(self, cognite_client: CogniteClient, chat_response_body: dict) -> None: """Test chatting with multiple messages at once.""" cognite_client.agents._post = MagicMock(return_value=MagicMock(json=lambda: chat_response_body)) - + session = cognite_client.agents.start_session("my_agent") messages = [ Message("Find temperature sensors"), Message("Show me their recent data") ] - response = session.chat(messages) - + session.chat(messages) + # Verify multiple messages were sent call_args = cognite_client.agents._post.call_args assert len(call_args[1]["json"]["messages"]) == 2 @@ -152,9 +152,9 @@ def test_agent_start_session(self, cognite_client: CogniteClient) -> None: description="Test agent" ) agent._cognite_client = cognite_client - + session = agent.start_session() - + assert isinstance(session, AgentSession) assert session.agent_id == "my_agent" assert session.cursor is None @@ -167,10 +167,10 @@ def test_agent_start_session_with_cursor(self, cognite_client: CogniteClient) -> name="My Agent" ) agent._cognite_client = cognite_client - + initial_cursor = "existing_cursor_456" session = agent.start_session(cursor=initial_cursor) - + assert session.agent_id == "my_agent" assert session.cursor == initial_cursor @@ -181,7 +181,7 @@ def test_agent_start_session_without_client_raises_error(self) -> None: name="My Agent" ) # Don't set _cognite_client - + with pytest.raises(CogniteMissingClientError): agent.start_session() @@ -192,16 +192,16 @@ def test_agent_start_session_with_none_client_raises_error(self) -> None: name="My Agent" ) agent._cognite_client = None - + with pytest.raises(CogniteMissingClientError): agent.start_session() class TestIntegrationWorkflow: def test_complete_workflow( - self, - cognite_client: CogniteClient, - chat_response_body: dict, + self, + cognite_client: CogniteClient, + chat_response_body: dict, followup_response_body: dict ) -> None: """Test the complete workflow as described in the user requirements.""" @@ -211,38 +211,38 @@ def test_complete_workflow( MagicMock(json=lambda: followup_response_body), ] cognite_client.agents._post = MagicMock(side_effect=responses) - + # Mock the retrieve method to return an agent with cognite_client mock_agent = Agent(external_id="my-agent", name="Test Agent") mock_agent._cognite_client = cognite_client cognite_client.agents.retrieve = MagicMock(return_value=mock_agent) - + # Test the workflow from the user requirements - + # 1. Start session from agents API session = cognite_client.agents.start_session(agent_id="my-agent", cursor=None) - + # 2. First chat resp = session.chat(Message("hello")) assert resp.text == "Hello! How can I help you today?" - + # 3. Follow-up question with automatic cursor management followup_response = session.chat(Message("nice to meet you")) assert followup_response.text == "I can help you with data analysis, finding assets, and more!" - + # 4. Test starting session from agent instance agent = cognite_client.agents.retrieve("my-agent") agent_session = agent.start_session() # also with optional cursor - + assert isinstance(agent_session, AgentSession) assert agent_session.agent_id == "my-agent" - + # Verify that cursors were managed correctly first_call_args = cognite_client.agents._post.call_args_list[0] second_call_args = cognite_client.agents._post.call_args_list[1] - + # First call should have no cursor assert "cursor" not in first_call_args[1]["json"] - + # Second call should have cursor from first response assert second_call_args[1]["json"]["cursor"] == "cursor_12345" \ No newline at end of file From a969924b16a29987b900972adf35ff141f8a17f8 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 22 Aug 2025 17:04:33 +0000 Subject: [PATCH 4/4] Refactor agent session code: improve imports, formatting, and remove redundancy Co-authored-by: kelvin.sundli --- cognite/client/data_classes/agents/agents.py | 4 +- cognite/client/data_classes/agents/chat.py | 10 ++--- .../test_api/test_agents_session.py | 42 ++++--------------- 3 files changed, 13 insertions(+), 43 deletions(-) diff --git a/cognite/client/data_classes/agents/agents.py b/cognite/client/data_classes/agents/agents.py index faa01d7e1d..094c054fbb 100644 --- a/cognite/client/data_classes/agents/agents.py +++ b/cognite/client/data_classes/agents/agents.py @@ -223,9 +223,9 @@ def start_session(self, cursor: str | None = None) -> AgentSession: from cognite.client.data_classes.agents.chat import AgentSession from cognite.client.exceptions import CogniteMissingClientError - if not hasattr(self, '_cognite_client') or self._cognite_client is None: + if not hasattr(self, "_cognite_client") or self._cognite_client is None: raise CogniteMissingClientError(self) - + return AgentSession(agent_id=self.external_id, cognite_client=self._cognite_client, cursor=cursor) diff --git a/cognite/client/data_classes/agents/chat.py b/cognite/client/data_classes/agents/chat.py index 9653b53644..3ce820a20f 100644 --- a/cognite/client/data_classes/agents/chat.py +++ b/cognite/client/data_classes/agents/chat.py @@ -1,9 +1,9 @@ from __future__ import annotations from abc import ABC, abstractmethod +from collections.abc import Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar, Literal -from collections.abc import Sequence from cognite.client.data_classes._base import CogniteObject, CogniteResource, CogniteResourceList from cognite.client.utils._text import convert_all_keys_to_camel_case @@ -324,12 +324,8 @@ def chat(self, messages: Message | Sequence[Message]) -> AgentChatResponse: ... ]) """ # Call the underlying agents.chat method with current cursor - response = self._cognite_client.agents.chat( - agent_id=self.agent_id, - messages=messages, - cursor=self._cursor - ) - + response = self._cognite_client.agents.chat(agent_id=self.agent_id, messages=messages, cursor=self._cursor) + # Update the cursor for the next interaction self._cursor = response.cursor diff --git a/tests/tests_unit/test_api/test_agents_session.py b/tests/tests_unit/test_api/test_agents_session.py index eea48803eb..2330b0c5bc 100644 --- a/tests/tests_unit/test_api/test_agents_session.py +++ b/tests/tests_unit/test_api/test_agents_session.py @@ -8,8 +8,6 @@ from cognite.client.data_classes.agents import Agent, AgentSession, Message from cognite.client.data_classes.agents.chat import ( AgentChatResponse, - AgentMessage, - TextContent, ) from cognite.client.exceptions import CogniteMissingClientError @@ -93,13 +91,8 @@ def test_session_chat_with_message(self, cognite_client: CogniteClient, chat_res assert response.text == "Hello! How can I help you today?" assert session.cursor == "cursor_12345" # Cursor should be updated - - def test_session_cursor_management( - self, - cognite_client: CogniteClient, - chat_response_body: dict, - followup_response_body: dict + self, cognite_client: CogniteClient, chat_response_body: dict, followup_response_body: dict ) -> None: """Test that cursor is automatically managed across multiple interactions.""" # Set up mock responses for two consecutive calls @@ -129,10 +122,7 @@ def test_session_chat_with_multiple_messages(self, cognite_client: CogniteClient cognite_client.agents._post = MagicMock(return_value=MagicMock(json=lambda: chat_response_body)) session = cognite_client.agents.start_session("my_agent") - messages = [ - Message("Find temperature sensors"), - Message("Show me their recent data") - ] + messages = [Message("Find temperature sensors"), Message("Show me their recent data")] session.chat(messages) # Verify multiple messages were sent @@ -146,11 +136,7 @@ class TestAgentStartSession: def test_agent_start_session(self, cognite_client: CogniteClient) -> None: """Test starting a session from an Agent instance.""" # Create an agent instance with a cognite client - agent = Agent( - external_id="my_agent", - name="My Agent", - description="Test agent" - ) + agent = Agent(external_id="my_agent", name="My Agent", description="Test agent") agent._cognite_client = cognite_client session = agent.start_session() @@ -162,10 +148,7 @@ def test_agent_start_session(self, cognite_client: CogniteClient) -> None: def test_agent_start_session_with_cursor(self, cognite_client: CogniteClient) -> None: """Test starting a session from an Agent instance with a cursor.""" - agent = Agent( - external_id="my_agent", - name="My Agent" - ) + agent = Agent(external_id="my_agent", name="My Agent") agent._cognite_client = cognite_client initial_cursor = "existing_cursor_456" @@ -176,10 +159,7 @@ def test_agent_start_session_with_cursor(self, cognite_client: CogniteClient) -> def test_agent_start_session_without_client_raises_error(self) -> None: """Test that starting a session without a cognite client raises an error.""" - agent = Agent( - external_id="my_agent", - name="My Agent" - ) + agent = Agent(external_id="my_agent", name="My Agent") # Don't set _cognite_client with pytest.raises(CogniteMissingClientError): @@ -187,10 +167,7 @@ def test_agent_start_session_without_client_raises_error(self) -> None: def test_agent_start_session_with_none_client_raises_error(self) -> None: """Test that starting a session with None cognite client raises an error.""" - agent = Agent( - external_id="my_agent", - name="My Agent" - ) + agent = Agent(external_id="my_agent", name="My Agent") agent._cognite_client = None with pytest.raises(CogniteMissingClientError): @@ -199,10 +176,7 @@ def test_agent_start_session_with_none_client_raises_error(self) -> None: class TestIntegrationWorkflow: def test_complete_workflow( - self, - cognite_client: CogniteClient, - chat_response_body: dict, - followup_response_body: dict + self, cognite_client: CogniteClient, chat_response_body: dict, followup_response_body: dict ) -> None: """Test the complete workflow as described in the user requirements.""" # Set up mock responses @@ -245,4 +219,4 @@ def test_complete_workflow( assert "cursor" not in first_call_args[1]["json"] # Second call should have cursor from first response - assert second_call_args[1]["json"]["cursor"] == "cursor_12345" \ No newline at end of file + assert second_call_args[1]["json"]["cursor"] == "cursor_12345"