Skip to content

Commit 85fb28e

Browse files
MervinPraisonpraisonai-triage-agent[bot]Copilot
authored
fix: re-export ManagedBackendProtocol and unify retrieve_session schema (fixes #1429) (#1439)
* fix: re-export ManagedBackendProtocol and unify retrieve_session schema (fixes #1429) - Add lazy re-export of ManagedBackendProtocol in praisonaiagents.managed - Create SessionInfo dataclass with unified schema (id, status, title, usage) - Update AnthropicManagedAgent.retrieve_session to use SessionInfo - Update LocalManagedAgent.retrieve_session to use SessionInfo - Add comprehensive unit tests for schema equality - Maintain backward compatibility with existing dict-based access Fixes #1429 Co-authored-by: MervinPraison <MervinPraison@users.noreply.github.com> * test: adapt test_retrieve_session_no_session to unified SessionInfo schema * fix: normalize SessionInfo defaults and local empty-session id Agent-Logs-Url: https://github.com/MervinPraison/PraisonAI/sessions/2663770e-de83-461c-8b12-3d3571e7e9ee Co-authored-by: MervinPraison <454862+MervinPraison@users.noreply.github.com> * fix: resolve SessionInfo schema issues identified by reviewers - Fix SessionInfo.__post_init__ elif chain bug preventing output_tokens initialization - Unify LocalManagedAgent status to 'unknown' (matching Anthropic backend) - Fix type annotation: usage: Optional[Dict[str, int]] = None - Improve getattr() None handling in AnthropicManagedAgent - Update test expectations for unified schema consistency Addresses review feedback from Gemini, CodeRabbit, Qodo, and Copilot. Co-authored-by: Mervin Praison <MervinPraison@users.noreply.github.com> --------- Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com> Co-authored-by: MervinPraison <MervinPraison@users.noreply.github.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: MervinPraison <454862+MervinPraison@users.noreply.github.com>
1 parent 8e308d4 commit 85fb28e

File tree

6 files changed

+336
-17
lines changed

6 files changed

+336
-17
lines changed

src/praisonai-agents/praisonaiagents/managed/__init__.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,4 +41,13 @@
4141
"ComputeConfig",
4242
"InstanceInfo",
4343
"InstanceStatus",
44+
"ManagedBackendProtocol",
4445
]
46+
47+
48+
def __getattr__(name: str):
49+
"""Lazy import ManagedBackendProtocol to keep module lightweight."""
50+
if name == "ManagedBackendProtocol":
51+
from ..agent.protocols import ManagedBackendProtocol
52+
return ManagedBackendProtocol
53+
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
Lines changed: 225 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,225 @@
1+
"""
2+
Tests for unified SessionInfo schema across managed agent backends.
3+
4+
Ensures that AnthropicManagedAgent and LocalManagedAgent return consistent
5+
retrieve_session() schemas as specified in issue #1429.
6+
"""
7+
8+
import pytest
9+
from unittest.mock import Mock, patch
10+
from typing import Dict, Any
11+
12+
# Test SessionInfo dataclass directly
13+
def test_session_info_defaults():
14+
"""Test SessionInfo provides proper defaults for all fields."""
15+
# Import from the wrapper package where SessionInfo is defined
16+
from praisonai.integrations._session_info import SessionInfo
17+
18+
# Test default instance
19+
info = SessionInfo()
20+
assert info.id == ""
21+
assert info.status == "unknown"
22+
assert info.title == ""
23+
assert info.usage == {"input_tokens": 0, "output_tokens": 0}
24+
25+
# Test to_dict() returns expected structure
26+
result = info.to_dict()
27+
expected = {
28+
"id": "",
29+
"status": "unknown",
30+
"title": "",
31+
"usage": {"input_tokens": 0, "output_tokens": 0}
32+
}
33+
assert result == expected
34+
35+
36+
def test_session_info_partial_usage():
37+
"""Test SessionInfo handles partial usage dictionaries properly."""
38+
from praisonai.integrations._session_info import SessionInfo
39+
40+
# Test partial usage (missing output_tokens)
41+
info = SessionInfo(usage={"input_tokens": 100})
42+
assert info.usage["input_tokens"] == 100
43+
assert info.usage["output_tokens"] == 0
44+
45+
# Test partial usage (missing input_tokens)
46+
info = SessionInfo(usage={"output_tokens": 200})
47+
assert info.usage["input_tokens"] == 0
48+
assert info.usage["output_tokens"] == 200
49+
50+
# Test empty usage dict gets both defaults
51+
info = SessionInfo(usage={})
52+
assert info.usage["input_tokens"] == 0
53+
assert info.usage["output_tokens"] == 0
54+
55+
56+
def test_session_info_complete():
57+
"""Test SessionInfo with all fields provided."""
58+
from praisonai.integrations._session_info import SessionInfo
59+
60+
info = SessionInfo(
61+
id="session-123",
62+
status="idle",
63+
title="Test Session",
64+
usage={"input_tokens": 150, "output_tokens": 75}
65+
)
66+
67+
result = info.to_dict()
68+
expected = {
69+
"id": "session-123",
70+
"status": "idle",
71+
"title": "Test Session",
72+
"usage": {"input_tokens": 150, "output_tokens": 75}
73+
}
74+
assert result == expected
75+
76+
77+
# Test schema consistency between backends
78+
@patch('praisonai.integrations.managed_agents.AnthropicManagedAgent._get_client')
79+
def test_anthropic_retrieve_session_schema(mock_get_client):
80+
"""Test AnthropicManagedAgent.retrieve_session returns unified schema."""
81+
from praisonai.integrations.managed_agents import AnthropicManagedAgent
82+
83+
# Mock Anthropic client response
84+
mock_session = Mock()
85+
mock_session.id = "anthropic-session-123"
86+
mock_session.status = "idle"
87+
mock_session.title = "Anthropic Session"
88+
89+
mock_usage = Mock()
90+
mock_usage.input_tokens = 100
91+
mock_usage.output_tokens = 50
92+
mock_session.usage = mock_usage
93+
94+
mock_client = Mock()
95+
mock_client.beta.sessions.retrieve.return_value = mock_session
96+
mock_get_client.return_value = mock_client
97+
98+
# Create agent and test retrieve_session
99+
agent = AnthropicManagedAgent()
100+
agent._session_id = "anthropic-session-123"
101+
102+
result = agent.retrieve_session()
103+
104+
# Verify schema structure
105+
assert isinstance(result, dict)
106+
assert "id" in result
107+
assert "status" in result
108+
assert "title" in result
109+
assert "usage" in result
110+
111+
# Verify field values
112+
assert result["id"] == "anthropic-session-123"
113+
assert result["status"] == "idle"
114+
assert result["title"] == "Anthropic Session"
115+
assert result["usage"] == {"input_tokens": 100, "output_tokens": 50}
116+
117+
118+
def test_local_retrieve_session_schema():
119+
"""Test LocalManagedAgent.retrieve_session returns unified schema."""
120+
from praisonai.integrations.managed_local import LocalManagedAgent
121+
122+
# Create local agent with session
123+
agent = LocalManagedAgent()
124+
agent._session_id = "local-session-456"
125+
agent.total_input_tokens = 200
126+
agent.total_output_tokens = 100
127+
128+
result = agent.retrieve_session()
129+
130+
# Verify schema structure
131+
assert isinstance(result, dict)
132+
assert "id" in result
133+
assert "status" in result
134+
assert "title" in result
135+
assert "usage" in result
136+
137+
# Verify field values
138+
assert result["id"] == "local-session-456"
139+
assert result["status"] == "idle"
140+
assert result["title"] == ""
141+
assert result["usage"] == {"input_tokens": 200, "output_tokens": 100}
142+
143+
144+
def test_schema_equality_between_backends():
145+
"""Test that both backends return schemas with identical structure."""
146+
from praisonai.integrations.managed_local import LocalManagedAgent
147+
148+
# Test LocalManagedAgent (easier to set up)
149+
local_agent = LocalManagedAgent()
150+
local_agent._session_id = "test-session"
151+
local_agent.total_input_tokens = 0
152+
local_agent.total_output_tokens = 0
153+
154+
local_result = local_agent.retrieve_session()
155+
156+
# Both should have identical keys
157+
expected_keys = {"id", "status", "title", "usage"}
158+
assert set(local_result.keys()) == expected_keys
159+
160+
# Usage should be a dict with input_tokens and output_tokens
161+
assert isinstance(local_result["usage"], dict)
162+
assert "input_tokens" in local_result["usage"]
163+
assert "output_tokens" in local_result["usage"]
164+
165+
# All values should be present (no None values)
166+
assert local_result["id"] is not None
167+
assert local_result["status"] is not None
168+
assert local_result["title"] is not None
169+
assert local_result["usage"] is not None
170+
171+
172+
def test_empty_session_anthropic():
173+
"""Test AnthropicManagedAgent.retrieve_session with no session."""
174+
from praisonai.integrations.managed_agents import AnthropicManagedAgent
175+
176+
agent = AnthropicManagedAgent()
177+
agent._session_id = None
178+
179+
result = agent.retrieve_session()
180+
181+
# Should return SessionInfo defaults
182+
expected = {
183+
"id": "",
184+
"status": "unknown",
185+
"title": "",
186+
"usage": {"input_tokens": 0, "output_tokens": 0}
187+
}
188+
assert result == expected
189+
190+
191+
def test_empty_session_local():
192+
"""Test LocalManagedAgent.retrieve_session with no session."""
193+
from praisonai.integrations.managed_local import LocalManagedAgent
194+
195+
agent = LocalManagedAgent()
196+
agent._session_id = ""
197+
agent.total_input_tokens = 0
198+
agent.total_output_tokens = 0
199+
200+
result = agent.retrieve_session()
201+
202+
# Should return unified schema with "unknown" status (matching Anthropic)
203+
assert result["id"] == ""
204+
assert result["status"] == "unknown"
205+
assert result["title"] == ""
206+
assert result["usage"] == {"input_tokens": 0, "output_tokens": 0}
207+
208+
209+
def test_empty_session_local_none_id():
210+
"""Test LocalManagedAgent.retrieve_session normalizes None session id to empty string."""
211+
from praisonai.integrations.managed_local import LocalManagedAgent
212+
213+
agent = LocalManagedAgent()
214+
agent._session_id = None
215+
agent.total_input_tokens = 0
216+
agent.total_output_tokens = 0
217+
218+
result = agent.retrieve_session()
219+
220+
assert result["id"] == ""
221+
assert result["status"] == "none"
222+
223+
224+
if __name__ == "__main__":
225+
pytest.main([__file__])

src/praisonai-agents/tests/unit/test_managed_backend.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -511,11 +511,14 @@ def test_retrieve_session(self):
511511
assert info["usage"]["output_tokens"] == 200
512512

513513
def test_retrieve_session_no_session(self):
514-
"""retrieve_session with no session returns empty dict."""
514+
"""retrieve_session with no session returns SessionInfo defaults (per #1429 unified schema)."""
515515
from praisonai.integrations.managed_agents import ManagedAgent
516516

517517
m = ManagedAgent(api_key="test-key")
518-
assert m.retrieve_session() == {}
518+
info = m.retrieve_session()
519+
assert info["id"] == ""
520+
assert info["status"] == "unknown"
521+
assert info["usage"] == {"input_tokens": 0, "output_tokens": 0}
519522

520523
def test_list_sessions(self):
521524
"""list_sessions should return list of session dicts."""
Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
"""
2+
Session Info dataclass for unified managed agent session schema.
3+
4+
Provides consistent return structure for retrieve_session() across all
5+
managed agent backends (Anthropic, Local, etc.).
6+
"""
7+
8+
from typing import Dict, Any, Optional
9+
from dataclasses import dataclass, asdict
10+
11+
12+
@dataclass
13+
class SessionInfo:
14+
"""Unified session info schema for managed agents.
15+
16+
All fields are always present with sensible defaults to ensure
17+
consistent API across different backend implementations.
18+
"""
19+
20+
id: str = ""
21+
"""Session ID (empty string if no session)"""
22+
23+
status: str = "unknown"
24+
"""Session status (idle, running, error, unknown, etc.)"""
25+
26+
title: str = ""
27+
"""Session title/name (empty if not set)"""
28+
29+
usage: Optional[Dict[str, int]] = None
30+
"""Token usage tracking with input_tokens and output_tokens"""
31+
32+
def __post_init__(self):
33+
"""Ensure usage field has proper defaults."""
34+
if self.id is None:
35+
self.id = ""
36+
if self.status is None:
37+
self.status = "unknown"
38+
if self.title is None:
39+
self.title = ""
40+
41+
if self.usage is None:
42+
self.usage = {"input_tokens": 0, "output_tokens": 0}
43+
else:
44+
# Use independent checks to ensure both tokens are always present
45+
self.usage.setdefault("input_tokens", 0)
46+
self.usage.setdefault("output_tokens", 0)
47+
48+
def to_dict(self) -> Dict[str, Any]:
49+
"""Convert to dictionary for backward compatibility.
50+
51+
Returns the same structure that retrieve_session() used to return,
52+
ensuring existing code continues to work.
53+
"""
54+
return asdict(self)

src/praisonai/praisonai/integrations/managed_agents.py

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -583,22 +583,38 @@ def interrupt(self) -> None:
583583
# retrieve_session — ManagedBackendProtocol
584584
# ------------------------------------------------------------------
585585
def retrieve_session(self) -> Dict[str, Any]:
586-
"""Retrieve current session metadata and usage from the API."""
586+
"""Retrieve current session metadata and usage from the API.
587+
588+
Returns unified SessionInfo schema with all fields always present:
589+
- id: Session ID
590+
- status: Session status (idle, running, error, etc.)
591+
- title: Session title/name
592+
- usage: Token usage with input_tokens and output_tokens
593+
"""
594+
from ._session_info import SessionInfo
595+
587596
if not self._session_id:
588-
return {}
597+
return SessionInfo().to_dict()
598+
589599
client = self._get_client()
590600
sess = client.beta.sessions.retrieve(self._session_id)
591-
result: Dict[str, Any] = {
592-
"id": getattr(sess, "id", self._session_id),
593-
"status": getattr(sess, "status", None),
594-
}
601+
602+
# Extract usage information
595603
usage = getattr(sess, "usage", None)
604+
usage_dict = {}
596605
if usage:
597-
result["usage"] = {
606+
usage_dict = {
598607
"input_tokens": getattr(usage, "input_tokens", 0),
599608
"output_tokens": getattr(usage, "output_tokens", 0),
600609
}
601-
return result
610+
611+
session_info = SessionInfo(
612+
id=getattr(sess, "id", None) or self._session_id or "",
613+
status=getattr(sess, "status", None) or "unknown",
614+
title=getattr(sess, "title", None) or "",
615+
usage=usage_dict if usage_dict else None,
616+
)
617+
return session_info.to_dict()
602618

603619
# ------------------------------------------------------------------
604620
# list_sessions — ManagedBackendProtocol

src/praisonai/praisonai/integrations/managed_local.py

Lines changed: 19 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -880,16 +880,28 @@ def interrupt(self) -> None:
880880
# retrieve_session / list_sessions — ManagedBackendProtocol
881881
# ------------------------------------------------------------------
882882
def retrieve_session(self) -> Dict[str, Any]:
883-
"""Retrieve current session metadata."""
883+
"""Retrieve current session metadata.
884+
885+
Returns unified SessionInfo schema with all fields always present:
886+
- id: Session ID
887+
- status: Session status (idle, running, error, etc.)
888+
- title: Session title/name
889+
- usage: Token usage with input_tokens and output_tokens
890+
"""
891+
from ._session_info import SessionInfo
892+
884893
self._sync_usage()
885-
return {
886-
"id": self._session_id,
887-
"status": "idle" if self._session_id else "none",
888-
"usage": {
894+
895+
session_info = SessionInfo(
896+
id=self._session_id or "",
897+
status="idle" if self._session_id else "unknown",
898+
title="", # Local agent doesn't track titles
899+
usage={
889900
"input_tokens": self.total_input_tokens,
890901
"output_tokens": self.total_output_tokens,
891-
},
892-
}
902+
}
903+
)
904+
return session_info.to_dict()
893905

894906
def list_sessions(self, **kwargs) -> List[Dict[str, Any]]:
895907
"""List all sessions created in this backend instance."""

0 commit comments

Comments
 (0)