diff --git a/cognite/client/_api/agents/agents.py b/cognite/client/_api/agents/agents.py index 8805851d3d..60d9c1b6e7 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(Message("Hello, how can you help me?")) + >>> print(response.text) + >>> 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(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/__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..094c054fbb 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,41 @@ 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 + >>> from cognite.client.data_classes.agents import Message + >>> client = CogniteClient() + >>> agent = client.agents.retrieve("my_agent") + >>> session = agent.start_session() + >>> 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 CogniteMissingClientError(self) + + 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..3ce820a20f 100644 --- a/cognite/client/data_classes/agents/chat.py +++ b/cognite/client/data_classes/agents/chat.py @@ -1,6 +1,7 @@ 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 @@ -271,3 +272,70 @@ 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: 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") + ... ]) + """ + # 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..2330b0c5bc --- /dev/null +++ b/tests/tests_unit/test_api/test_agents_session.py @@ -0,0 +1,222 @@ +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, +) +from cognite.client.exceptions import CogniteMissingClientError + + +@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_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(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 + 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_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 + session.chat(Message("Hello")) + assert session.cursor == "cursor_12345" + + # Second interaction - should include cursor from first response + 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" + 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")] + 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(CogniteMissingClientError): + 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(CogniteMissingClientError): + 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"