wehub-resource-sync
1818 lines
65 KiB
Python
1818 lines
65 KiB
Python
"""Tests for AdvancedSQLiteSession functionality."""
|
|
|
|
import asyncio
|
|
import json
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import Any, cast
|
|
|
|
import pytest
|
|
|
|
pytest.importorskip("sqlalchemy") # Skip tests if SQLAlchemy is not installed
|
|
from openai.types.responses.response_usage import InputTokensDetails, OutputTokensDetails
|
|
|
|
from agents import Agent, Runner, TResponseInputItem, function_tool
|
|
from agents.extensions.memory import AdvancedSQLiteSession
|
|
from agents.result import RunResult
|
|
from agents.run_context import RunContextWrapper
|
|
from agents.usage import Usage
|
|
from tests.fake_model import FakeModel
|
|
from tests.test_responses import get_text_message
|
|
|
|
# Mark all tests in this file as asyncio
|
|
pytestmark = pytest.mark.asyncio
|
|
|
|
|
|
@function_tool
|
|
async def test_tool(query: str) -> str:
|
|
"""A test tool for testing tool call tracking."""
|
|
return f"Tool result for: {query}"
|
|
|
|
|
|
@pytest.fixture
|
|
def agent() -> Agent:
|
|
"""Fixture for a basic agent with a fake model."""
|
|
return Agent(name="test", model=FakeModel(), tools=[test_tool])
|
|
|
|
|
|
@pytest.fixture
|
|
def usage_data() -> Usage:
|
|
"""Fixture for test usage data."""
|
|
return Usage(
|
|
requests=1,
|
|
input_tokens=50,
|
|
output_tokens=30,
|
|
total_tokens=80,
|
|
input_tokens_details=InputTokensDetails.model_validate(
|
|
{"cache_write_tokens": 0, "cached_tokens": 10}
|
|
),
|
|
output_tokens_details=OutputTokensDetails(reasoning_tokens=5),
|
|
)
|
|
|
|
|
|
def create_mock_run_result(usage: Usage | None = None, agent: Agent | None = None) -> RunResult:
|
|
"""Helper function to create a mock RunResult for testing."""
|
|
if agent is None:
|
|
agent = Agent(name="test", model=FakeModel())
|
|
|
|
if usage is None:
|
|
usage = Usage(
|
|
requests=1,
|
|
input_tokens=50,
|
|
output_tokens=30,
|
|
total_tokens=80,
|
|
input_tokens_details=InputTokensDetails.model_validate(
|
|
{"cache_write_tokens": 0, "cached_tokens": 10}
|
|
),
|
|
output_tokens_details=OutputTokensDetails(reasoning_tokens=5),
|
|
)
|
|
|
|
context_wrapper = RunContextWrapper(context=None, usage=usage)
|
|
|
|
return RunResult(
|
|
input="test input",
|
|
new_items=[],
|
|
raw_responses=[],
|
|
final_output="test output",
|
|
input_guardrail_results=[],
|
|
output_guardrail_results=[],
|
|
tool_input_guardrail_results=[],
|
|
tool_output_guardrail_results=[],
|
|
context_wrapper=context_wrapper,
|
|
_last_agent=agent,
|
|
interruptions=[],
|
|
)
|
|
|
|
|
|
class FailingOnceStructureMetadataSession(AdvancedSQLiteSession):
|
|
"""Advanced session test double that fails the next structure metadata write."""
|
|
|
|
def __init__(self, **kwargs: Any):
|
|
super().__init__(**kwargs)
|
|
self.fail_structure_metadata_once = True
|
|
|
|
def _insert_structure_metadata(
|
|
self,
|
|
conn: Any,
|
|
items: list[TResponseInputItem],
|
|
) -> None:
|
|
if self.fail_structure_metadata_once:
|
|
self.fail_structure_metadata_once = False
|
|
raise RuntimeError("structure metadata failed")
|
|
super()._insert_structure_metadata(conn, items)
|
|
|
|
|
|
class PartiallyFailingStructureMetadataSession(AdvancedSQLiteSession):
|
|
"""Advanced session test double that fails after writing one structure row."""
|
|
|
|
def _insert_structure_metadata(
|
|
self,
|
|
conn: Any,
|
|
items: list[TResponseInputItem],
|
|
) -> None:
|
|
cursor = conn.execute(
|
|
f"SELECT id FROM {self.messages_table} WHERE session_id = ? ORDER BY id ASC LIMIT 1",
|
|
(self.session_id,),
|
|
)
|
|
row = cursor.fetchone()
|
|
if row is None:
|
|
raise RuntimeError("no inserted message id found")
|
|
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO message_structure
|
|
(session_id, message_id, branch_id, message_type, sequence_number,
|
|
user_turn_number, branch_turn_number, tool_name)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(self.session_id, row[0], self._current_branch_id, "user", 1, 1, 1, None),
|
|
)
|
|
raise RuntimeError("structure metadata failed after partial write")
|
|
|
|
|
|
async def test_advanced_session_basic_functionality(agent: Agent):
|
|
"""Test basic AdvancedSQLiteSession functionality."""
|
|
session_id = "advanced_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Test basic session operations work
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi there!"},
|
|
]
|
|
await session.add_items(items)
|
|
|
|
# Get items and verify
|
|
retrieved = await session.get_items()
|
|
assert len(retrieved) == 2
|
|
assert retrieved[0].get("content") == "Hello"
|
|
assert retrieved[1].get("content") == "Hi there!"
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_advanced_session_respects_custom_table_names():
|
|
"""AdvancedSQLiteSession should consistently use configured table names."""
|
|
session = AdvancedSQLiteSession(
|
|
session_id="advanced_custom_tables",
|
|
create_tables=True,
|
|
sessions_table="custom_agent_sessions",
|
|
messages_table="custom_agent_messages",
|
|
)
|
|
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi there!"},
|
|
{"role": "user", "content": "Let's do some math"},
|
|
{"role": "assistant", "content": "Sure"},
|
|
]
|
|
await session.add_items(items)
|
|
|
|
assert await session.get_items() == items
|
|
|
|
conversation_turns = await session.get_conversation_turns()
|
|
assert [turn["turn"] for turn in conversation_turns] == [1, 2]
|
|
|
|
matching_turns = await session.find_turns_by_content("math")
|
|
assert [turn["turn"] for turn in matching_turns] == [2]
|
|
|
|
conn = session._get_connection()
|
|
structure_foreign_keys = {
|
|
row[2] for row in conn.execute("PRAGMA foreign_key_list(message_structure)").fetchall()
|
|
}
|
|
usage_foreign_keys = {
|
|
row[2] for row in conn.execute("PRAGMA foreign_key_list(turn_usage)").fetchall()
|
|
}
|
|
assert structure_foreign_keys == {
|
|
session.messages_table,
|
|
session.sessions_table,
|
|
}
|
|
assert usage_foreign_keys == {session.sessions_table}
|
|
|
|
branch_name = await session.create_branch_from_turn(2, "custom_branch")
|
|
assert branch_name == "custom_branch"
|
|
assert await session.get_items() == items[:2]
|
|
assert await session.get_items(branch_id="main") == items
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_add_items_rolls_back_messages_when_structure_metadata_fails():
|
|
"""Failed structure metadata writes should not leave invisible message rows."""
|
|
session = FailingOnceStructureMetadataSession(
|
|
session_id="advanced_add_items_rollback",
|
|
create_tables=True,
|
|
)
|
|
items: list[TResponseInputItem] = [{"role": "user", "content": "not saved"}]
|
|
|
|
try:
|
|
with pytest.raises(RuntimeError, match="structure metadata failed"):
|
|
await session.add_items(items)
|
|
|
|
assert await session.get_items() == []
|
|
|
|
with session._locked_connection() as conn:
|
|
message_count = conn.execute(
|
|
f"SELECT COUNT(*) FROM {session.messages_table} WHERE session_id = ?",
|
|
(session.session_id,),
|
|
).fetchone()[0]
|
|
structure_count = conn.execute(
|
|
"SELECT COUNT(*) FROM message_structure WHERE session_id = ?",
|
|
(session.session_id,),
|
|
).fetchone()[0]
|
|
|
|
assert message_count == 0
|
|
assert structure_count == 0
|
|
finally:
|
|
session.close()
|
|
|
|
|
|
async def test_add_items_can_retry_after_structure_metadata_failure():
|
|
"""Retrying after a metadata failure should persist the batch exactly once."""
|
|
session = FailingOnceStructureMetadataSession(
|
|
session_id="advanced_add_items_retry",
|
|
create_tables=True,
|
|
)
|
|
items: list[TResponseInputItem] = [{"role": "user", "content": "saved once"}]
|
|
|
|
try:
|
|
with pytest.raises(RuntimeError, match="structure metadata failed"):
|
|
await session.add_items(items)
|
|
|
|
await session.add_items(items)
|
|
|
|
assert await session.get_items() == items
|
|
|
|
with session._locked_connection() as conn:
|
|
message_count = conn.execute(
|
|
f"SELECT COUNT(*) FROM {session.messages_table} WHERE session_id = ?",
|
|
(session.session_id,),
|
|
).fetchone()[0]
|
|
structure_count = conn.execute(
|
|
"SELECT COUNT(*) FROM message_structure WHERE session_id = ?",
|
|
(session.session_id,),
|
|
).fetchone()[0]
|
|
|
|
assert message_count == 1
|
|
assert structure_count == 1
|
|
finally:
|
|
session.close()
|
|
|
|
|
|
async def test_add_items_failure_preserves_existing_history():
|
|
"""A failed batch should not roll back or hide previously committed messages."""
|
|
session = FailingOnceStructureMetadataSession(
|
|
session_id="advanced_add_items_existing_history",
|
|
create_tables=True,
|
|
)
|
|
existing_items: list[TResponseInputItem] = [{"role": "user", "content": "already saved"}]
|
|
failed_items: list[TResponseInputItem] = [{"role": "assistant", "content": "not saved"}]
|
|
|
|
try:
|
|
session.fail_structure_metadata_once = False
|
|
await session.add_items(existing_items)
|
|
|
|
session.fail_structure_metadata_once = True
|
|
with pytest.raises(RuntimeError, match="structure metadata failed"):
|
|
await session.add_items(failed_items)
|
|
|
|
assert await session.get_items() == existing_items
|
|
|
|
with session._locked_connection() as conn:
|
|
message_count = conn.execute(
|
|
f"SELECT COUNT(*) FROM {session.messages_table} WHERE session_id = ?",
|
|
(session.session_id,),
|
|
).fetchone()[0]
|
|
structure_count = conn.execute(
|
|
"SELECT COUNT(*) FROM message_structure WHERE session_id = ?",
|
|
(session.session_id,),
|
|
).fetchone()[0]
|
|
|
|
assert message_count == 1
|
|
assert structure_count == 1
|
|
finally:
|
|
session.close()
|
|
|
|
|
|
async def test_add_items_rolls_back_partial_structure_metadata_write():
|
|
"""Partial metadata writes should roll back with the message rows in the same batch."""
|
|
session = PartiallyFailingStructureMetadataSession(
|
|
session_id="advanced_add_items_partial_metadata",
|
|
create_tables=True,
|
|
)
|
|
items: list[TResponseInputItem] = [{"role": "user", "content": "not saved"}]
|
|
|
|
try:
|
|
with pytest.raises(RuntimeError, match="structure metadata failed after partial write"):
|
|
await session.add_items(items)
|
|
|
|
assert await session.get_items() == []
|
|
|
|
with session._locked_connection() as conn:
|
|
message_count = conn.execute(
|
|
f"SELECT COUNT(*) FROM {session.messages_table} WHERE session_id = ?",
|
|
(session.session_id,),
|
|
).fetchone()[0]
|
|
structure_count = conn.execute(
|
|
"SELECT COUNT(*) FROM message_structure WHERE session_id = ?",
|
|
(session.session_id,),
|
|
).fetchone()[0]
|
|
|
|
assert message_count == 0
|
|
assert structure_count == 0
|
|
finally:
|
|
session.close()
|
|
|
|
|
|
async def test_message_structure_tracking(agent: Agent):
|
|
"""Test that message structure is properly tracked."""
|
|
session_id = "structure_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Add various types of messages
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "What's 2+2?"},
|
|
{"type": "function_call", "name": "calculator", "arguments": '{"expression": "2+2"}'}, # type: ignore
|
|
{"type": "function_call_output", "output": "4"}, # type: ignore
|
|
{"role": "assistant", "content": "The answer is 4"},
|
|
{"type": "reasoning", "summary": [{"text": "Simple math", "type": "summary_text"}]}, # type: ignore
|
|
]
|
|
await session.add_items(items)
|
|
|
|
# Get conversation structure
|
|
conversation_turns = await session.get_conversation_by_turns()
|
|
assert len(conversation_turns) == 1 # Should be one user turn
|
|
|
|
turn_1_items = conversation_turns[1]
|
|
assert len(turn_1_items) == 5
|
|
|
|
# Verify item types are classified correctly
|
|
item_types = [item["type"] for item in turn_1_items]
|
|
assert "user" in item_types
|
|
assert "function_call" in item_types
|
|
assert "function_call_output" in item_types
|
|
assert "assistant" in item_types
|
|
assert "reasoning" in item_types
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_tool_usage_tracking(agent: Agent):
|
|
"""Test tool usage tracking functionality."""
|
|
session_id = "tools_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Add items with tool calls
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Search for cats"},
|
|
{"type": "function_call", "name": "web_search", "arguments": '{"query": "cats"}'}, # type: ignore
|
|
{"type": "function_call_output", "output": "Found cat information"}, # type: ignore
|
|
{"type": "function_call", "name": "calculator", "arguments": '{"expression": "1+1"}'}, # type: ignore
|
|
{"type": "function_call_output", "output": "2"}, # type: ignore
|
|
{"role": "assistant", "content": "I found information about cats and calculated 1+1=2"},
|
|
]
|
|
await session.add_items(items)
|
|
|
|
# Get tool usage
|
|
tool_usage = await session.get_tool_usage()
|
|
assert len(tool_usage) == 2 # Two different tools used
|
|
|
|
tool_names = {usage[0] for usage in tool_usage}
|
|
assert "web_search" in tool_names
|
|
assert "calculator" in tool_names
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_tool_usage_tracking_preserves_namespaces_and_tool_search(agent: Agent):
|
|
"""Tool usage should retain namespaces and count tool_search calls once."""
|
|
session_id = "tools_namespace_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Look up the same account in multiple systems"},
|
|
{
|
|
"type": "function_call",
|
|
"name": "lookup_account",
|
|
"namespace": "crm",
|
|
"arguments": '{"account_id": "acct_123"}',
|
|
"call_id": "crm-call",
|
|
},
|
|
{
|
|
"type": "function_call",
|
|
"name": "lookup_account",
|
|
"namespace": "billing",
|
|
"arguments": '{"account_id": "acct_123"}',
|
|
"call_id": "billing-call",
|
|
},
|
|
{
|
|
"type": "tool_search_call",
|
|
"id": "tsc_memory",
|
|
"arguments": {"paths": ["crm"], "query": "lookup_account"},
|
|
"execution": "server",
|
|
"status": "completed",
|
|
},
|
|
cast(
|
|
TResponseInputItem,
|
|
{
|
|
"type": "tool_search_output",
|
|
"id": "tso_memory",
|
|
"execution": "server",
|
|
"status": "completed",
|
|
"tools": [
|
|
{
|
|
"type": "function",
|
|
"name": "lookup_account",
|
|
"description": "Look up an account.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"account_id": {
|
|
"type": "string",
|
|
}
|
|
},
|
|
"required": ["account_id"],
|
|
},
|
|
"defer_loading": True,
|
|
}
|
|
],
|
|
},
|
|
),
|
|
]
|
|
await session.add_items(items)
|
|
|
|
usage_by_tool = {tool_name: count for tool_name, count, _turn in await session.get_tool_usage()}
|
|
|
|
assert usage_by_tool["crm.lookup_account"] == 1
|
|
assert usage_by_tool["billing.lookup_account"] == 1
|
|
assert usage_by_tool["tool_search"] == 1
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_tool_usage_tracking_counts_tool_search_output_without_matching_call(
|
|
agent: Agent,
|
|
) -> None:
|
|
"""Tool-search output-only histories should still report one tool_search usage."""
|
|
session_id = "tools_tool_search_output_only_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Look up customer_42"},
|
|
cast(
|
|
TResponseInputItem,
|
|
{
|
|
"type": "tool_search_output",
|
|
"id": "tso_memory_only",
|
|
"execution": "server",
|
|
"status": "completed",
|
|
"tools": [
|
|
{
|
|
"type": "function",
|
|
"name": "lookup_account",
|
|
"description": "Look up an account.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"account_id": {
|
|
"type": "string",
|
|
}
|
|
},
|
|
"required": ["account_id"],
|
|
},
|
|
}
|
|
],
|
|
},
|
|
),
|
|
]
|
|
await session.add_items(items)
|
|
|
|
usage_by_tool = {tool_name: count for tool_name, count, _turn in await session.get_tool_usage()}
|
|
|
|
assert usage_by_tool["tool_search"] == 1
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_tool_usage_tracking_uses_bare_name_for_deferred_top_level_calls(agent: Agent):
|
|
"""Deferred top-level tool calls should not retain synthetic namespace aliases."""
|
|
session_id = "tools_deferred_top_level_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "What is the weather?"},
|
|
{
|
|
"type": "function_call",
|
|
"name": "get_weather",
|
|
"arguments": '{"city": "Tokyo"}',
|
|
"call_id": "weather-call",
|
|
},
|
|
{
|
|
"type": "function_call",
|
|
"name": "get_weather",
|
|
"namespace": "get_weather",
|
|
"arguments": '{"city": "Osaka"}',
|
|
"call_id": "weather-call-2",
|
|
},
|
|
]
|
|
await session.add_items(items)
|
|
|
|
usage_by_tool = {tool_name: count for tool_name, count, _turn in await session.get_tool_usage()}
|
|
|
|
assert usage_by_tool["get_weather"] == 2
|
|
assert "get_weather.get_weather" not in usage_by_tool
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_tool_usage_tracking_collapses_reserved_same_name_namespace_shape(
|
|
agent: Agent,
|
|
):
|
|
"""Reserved same-name namespace wire shapes should collapse to the bare tool name."""
|
|
session_id = "tools_deferred_top_level_namespace_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "What is the weather?"},
|
|
{
|
|
"type": "function_call",
|
|
"name": "lookup_account",
|
|
"namespace": "lookup_account",
|
|
"arguments": '{"account_id": "acct_123"}',
|
|
"call_id": "lookup-call",
|
|
},
|
|
]
|
|
await session.add_items(items)
|
|
|
|
usage_by_tool = {tool_name: count for tool_name, count, _turn in await session.get_tool_usage()}
|
|
|
|
assert usage_by_tool["lookup_account"] == 1
|
|
assert "lookup_account.lookup_account" not in usage_by_tool
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_branching_functionality(agent: Agent):
|
|
"""Test branching functionality - create, switch, and delete branches."""
|
|
session_id = "branching_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Add multiple turns to main branch
|
|
turn_1_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "First question"},
|
|
{"role": "assistant", "content": "First answer"},
|
|
]
|
|
await session.add_items(turn_1_items)
|
|
|
|
turn_2_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Second question"},
|
|
{"role": "assistant", "content": "Second answer"},
|
|
]
|
|
await session.add_items(turn_2_items)
|
|
|
|
turn_3_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Third question"},
|
|
{"role": "assistant", "content": "Third answer"},
|
|
]
|
|
await session.add_items(turn_3_items)
|
|
|
|
# Verify all items are in main branch
|
|
all_items = await session.get_items()
|
|
assert len(all_items) == 6
|
|
|
|
# Create a branch from turn 2
|
|
branch_name = await session.create_branch_from_turn(2, "test_branch")
|
|
assert branch_name == "test_branch"
|
|
|
|
# Verify we're now on the new branch
|
|
assert session._current_branch_id == "test_branch"
|
|
|
|
# Verify the branch has the same content up to turn 2 (copies messages before turn 2)
|
|
branch_items = await session.get_items()
|
|
assert len(branch_items) == 2 # Only first turn items (before turn 2)
|
|
assert branch_items[0].get("content") == "First question"
|
|
assert branch_items[1].get("content") == "First answer"
|
|
|
|
# Switch back to main branch
|
|
await session.switch_to_branch("main")
|
|
assert session._current_branch_id == "main"
|
|
|
|
# Verify main branch still has all items
|
|
main_items = await session.get_items()
|
|
assert len(main_items) == 6
|
|
|
|
# List branches
|
|
branches = await session.list_branches()
|
|
assert len(branches) == 2
|
|
branch_ids = [b["branch_id"] for b in branches]
|
|
assert "main" in branch_ids
|
|
assert "test_branch" in branch_ids
|
|
|
|
# Delete the test branch
|
|
await session.delete_branch("test_branch")
|
|
|
|
# Verify branch is deleted
|
|
branches_after_delete = await session.list_branches()
|
|
assert len(branches_after_delete) == 1
|
|
assert branches_after_delete[0]["branch_id"] == "main"
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_delete_branch_removes_branch_only_messages():
|
|
"""Deleting a branch should not leave unreferenced branch-only messages behind."""
|
|
session_id = "branch_delete_cleanup_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
main_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "First question"},
|
|
{"role": "assistant", "content": "First answer"},
|
|
{"role": "user", "content": "Second question"},
|
|
{"role": "assistant", "content": "Second answer"},
|
|
]
|
|
await session.add_items(main_items)
|
|
|
|
await session.create_branch_from_turn(2, "cleanup_branch")
|
|
branch_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Branch-only question"},
|
|
{"role": "assistant", "content": "Branch-only answer"},
|
|
]
|
|
await session.add_items(branch_items)
|
|
|
|
await session.delete_branch("cleanup_branch", force=True)
|
|
|
|
with session._locked_connection() as conn:
|
|
rows = conn.execute(
|
|
f"""
|
|
SELECT message_data
|
|
FROM {session.messages_table}
|
|
WHERE session_id = ?
|
|
ORDER BY id
|
|
""",
|
|
(session.session_id,),
|
|
).fetchall()
|
|
|
|
contents = [json.loads(message_data)["content"] for (message_data,) in rows]
|
|
assert contents == [
|
|
"First question",
|
|
"First answer",
|
|
"Second question",
|
|
"Second answer",
|
|
]
|
|
assert await session.get_items(branch_id="main") == main_items
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_delete_branch_keeps_messages_still_referenced_by_another_branch():
|
|
"""Deleting one branch should keep messages inherited by a surviving branch."""
|
|
session = AdvancedSQLiteSession(
|
|
session_id="branch_delete_shared_descendant_test",
|
|
create_tables=True,
|
|
)
|
|
|
|
main_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Main first question"},
|
|
{"role": "assistant", "content": "Main first answer"},
|
|
{"role": "user", "content": "Main second question"},
|
|
{"role": "assistant", "content": "Main second answer"},
|
|
]
|
|
branch_a_shared_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Branch A shared question"},
|
|
{"role": "assistant", "content": "Branch A shared answer"},
|
|
]
|
|
branch_a_only_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Branch A only question"},
|
|
{"role": "assistant", "content": "Branch A only answer"},
|
|
]
|
|
|
|
try:
|
|
await session.add_items(main_items)
|
|
await session.create_branch_from_turn(2, "branch_a")
|
|
await session.add_items(branch_a_shared_items + branch_a_only_items)
|
|
|
|
await session.create_branch_from_turn(3, "branch_b")
|
|
await session.delete_branch("branch_a")
|
|
|
|
with session._locked_connection() as conn:
|
|
rows = conn.execute(
|
|
f"""
|
|
SELECT message_data
|
|
FROM {session.messages_table}
|
|
WHERE session_id = ?
|
|
ORDER BY id
|
|
""",
|
|
(session.session_id,),
|
|
).fetchall()
|
|
|
|
contents = [json.loads(message_data)["content"] for (message_data,) in rows]
|
|
assert "Branch A shared question" in contents
|
|
assert "Branch A shared answer" in contents
|
|
assert "Branch A only question" not in contents
|
|
assert "Branch A only answer" not in contents
|
|
assert await session.get_items(branch_id="branch_b") == [
|
|
*main_items[:2],
|
|
*branch_a_shared_items,
|
|
]
|
|
finally:
|
|
session.close()
|
|
|
|
|
|
async def test_orphan_cleanup_uses_set_based_delete_for_many_messages():
|
|
"""Orphan cleanup should not build one DELETE parameter per orphaned row."""
|
|
|
|
class RecordingCursor:
|
|
def __init__(self, cursor: Any, connection: "RecordingConnection") -> None:
|
|
self._cursor = cursor
|
|
self._connection = connection
|
|
|
|
@property
|
|
def rowcount(self) -> int:
|
|
return cast(int, self._cursor.rowcount)
|
|
|
|
def execute(self, sql: str, parameters: Any = None) -> Any:
|
|
normalized_sql = " ".join(sql.split()).upper()
|
|
if normalized_sql.startswith("DELETE"):
|
|
self._connection.delete_parameter_counts.append(len(parameters or ()))
|
|
if parameters is None:
|
|
return self._cursor.execute(sql)
|
|
return self._cursor.execute(sql, parameters)
|
|
|
|
def fetchall(self) -> Any:
|
|
return self._cursor.fetchall()
|
|
|
|
def close(self) -> None:
|
|
self._cursor.close()
|
|
|
|
class RecordingConnection:
|
|
def __init__(self, conn: Any) -> None:
|
|
self._conn = conn
|
|
self.delete_parameter_counts: list[int] = []
|
|
|
|
def cursor(self) -> RecordingCursor:
|
|
return RecordingCursor(self._conn.cursor(), self)
|
|
|
|
session = AdvancedSQLiteSession(
|
|
session_id="branch_delete_many_orphans_cleanup",
|
|
create_tables=True,
|
|
)
|
|
orphan_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": f"orphan {i}"} for i in range(1200)
|
|
]
|
|
|
|
try:
|
|
with session._locked_connection() as conn:
|
|
session._insert_items(conn, orphan_items)
|
|
conn.commit()
|
|
|
|
recording_conn = RecordingConnection(conn)
|
|
deleted_count = session._cleanup_orphaned_messages_sync(cast(Any, recording_conn))
|
|
conn.commit()
|
|
|
|
remaining_count = conn.execute(
|
|
f"SELECT COUNT(*) FROM {session.messages_table} WHERE session_id = ?",
|
|
(session.session_id,),
|
|
).fetchone()[0]
|
|
|
|
assert deleted_count == len(orphan_items)
|
|
assert remaining_count == 0
|
|
assert recording_conn.delete_parameter_counts == [2]
|
|
finally:
|
|
session.close()
|
|
|
|
|
|
async def test_get_conversation_turns():
|
|
"""Test get_conversation_turns functionality."""
|
|
session_id = "conversation_turns_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Add multiple turns
|
|
turn_1_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Hello there"},
|
|
{"role": "assistant", "content": "Hi!"},
|
|
]
|
|
await session.add_items(turn_1_items)
|
|
|
|
turn_2_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "How are you doing today?"},
|
|
{"role": "assistant", "content": "I'm doing well, thanks!"},
|
|
]
|
|
await session.add_items(turn_2_items)
|
|
|
|
# Get conversation turns
|
|
turns = await session.get_conversation_turns()
|
|
assert len(turns) == 2
|
|
|
|
# Verify turn structure
|
|
assert turns[0]["turn"] == 1
|
|
assert turns[0]["content"] == "Hello there"
|
|
assert turns[0]["full_content"] == "Hello there"
|
|
assert turns[0]["can_branch"] is True
|
|
assert "timestamp" in turns[0]
|
|
|
|
assert turns[1]["turn"] == 2
|
|
assert turns[1]["content"] == "How are you doing today?"
|
|
assert turns[1]["full_content"] == "How are you doing today?"
|
|
assert turns[1]["can_branch"] is True
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_find_turns_by_content():
|
|
"""Test find_turns_by_content functionality."""
|
|
session_id = "find_turns_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Add multiple turns with different content
|
|
turn_1_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Tell me about cats"},
|
|
{"role": "assistant", "content": "Cats are great pets"},
|
|
]
|
|
await session.add_items(turn_1_items)
|
|
|
|
turn_2_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "What about dogs?"},
|
|
{"role": "assistant", "content": "Dogs are also great pets"},
|
|
]
|
|
await session.add_items(turn_2_items)
|
|
|
|
turn_3_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Tell me about cats again"},
|
|
{"role": "assistant", "content": "Cats are wonderful companions"},
|
|
]
|
|
await session.add_items(turn_3_items)
|
|
|
|
# Search for turns containing "cats"
|
|
cat_turns = await session.find_turns_by_content("cats")
|
|
assert len(cat_turns) == 2
|
|
assert cat_turns[0]["turn"] == 1
|
|
assert cat_turns[1]["turn"] == 3
|
|
|
|
# Search for turns containing "dogs"
|
|
dog_turns = await session.find_turns_by_content("dogs")
|
|
assert len(dog_turns) == 1
|
|
assert dog_turns[0]["turn"] == 2
|
|
|
|
# Search for non-existent content
|
|
no_turns = await session.find_turns_by_content("elephants")
|
|
assert len(no_turns) == 0
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_get_conversation_turns_with_list_content():
|
|
"""List (multimodal) content is previewed as a string instead of crashing or leaking a list."""
|
|
session_id = "conversation_turns_list_content_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# A short list content must be previewed as a string, not returned as the raw list.
|
|
short_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": [{"type": "input_text", "text": "hello"}]},
|
|
]
|
|
await session.add_items(short_items)
|
|
|
|
# A long list content must not raise when the preview is built.
|
|
long_items: list[TResponseInputItem] = [
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "input_text", "text": str(i)} for i in range(101)],
|
|
},
|
|
]
|
|
await session.add_items(long_items)
|
|
|
|
turns = await session.get_conversation_turns()
|
|
assert len(turns) == 2
|
|
|
|
# 'content' is the documented truncated preview string, while 'full_content' keeps the list.
|
|
assert isinstance(turns[0]["content"], str)
|
|
assert isinstance(turns[0]["full_content"], list)
|
|
|
|
assert isinstance(turns[1]["content"], str)
|
|
assert turns[1]["content"].endswith("...")
|
|
assert isinstance(turns[1]["full_content"], list)
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_find_turns_by_content_with_list_content():
|
|
"""find_turns_by_content returns a string preview for list (multimodal) content."""
|
|
session_id = "find_turns_list_content_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": [{"type": "input_text", "text": "hello world"}]},
|
|
]
|
|
await session.add_items(items)
|
|
|
|
matches = await session.find_turns_by_content("hello")
|
|
assert len(matches) == 1
|
|
assert isinstance(matches[0]["content"], str)
|
|
assert isinstance(matches[0]["full_content"], list)
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_create_branch_from_content():
|
|
"""Test create_branch_from_content functionality."""
|
|
session_id = "branch_from_content_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Add multiple turns
|
|
turn_1_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "First question about math"},
|
|
{"role": "assistant", "content": "Math answer"},
|
|
]
|
|
await session.add_items(turn_1_items)
|
|
|
|
turn_2_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Second question about science"},
|
|
{"role": "assistant", "content": "Science answer"},
|
|
]
|
|
await session.add_items(turn_2_items)
|
|
|
|
turn_3_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Another math question"},
|
|
{"role": "assistant", "content": "Another math answer"},
|
|
]
|
|
await session.add_items(turn_3_items)
|
|
|
|
# Create branch from first occurrence of "math"
|
|
branch_name = await session.create_branch_from_content("math", "math_branch")
|
|
assert branch_name == "math_branch"
|
|
|
|
# Verify we're on the new branch
|
|
assert session._current_branch_id == "math_branch"
|
|
|
|
# Verify branch contains only items up to the first math turn (copies messages before turn 1)
|
|
branch_items = await session.get_items()
|
|
assert len(branch_items) == 0 # No messages before turn 1
|
|
|
|
# Test error case - search term not found
|
|
with pytest.raises(ValueError, match="No user turns found containing 'nonexistent'"):
|
|
await session.create_branch_from_content("nonexistent", "error_branch")
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_branch_specific_operations():
|
|
"""Test operations that work with specific branches."""
|
|
session_id = "branch_specific_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Add items to main branch
|
|
turn_1_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Main branch question"},
|
|
{"role": "assistant", "content": "Main branch answer"},
|
|
]
|
|
await session.add_items(turn_1_items)
|
|
|
|
# Add usage data for main branch
|
|
usage_main = Usage(requests=1, input_tokens=50, output_tokens=30, total_tokens=80)
|
|
run_result_main = create_mock_run_result(usage_main)
|
|
await session.store_run_usage(run_result_main)
|
|
|
|
# Create a branch from turn 1 (copies messages before turn 1, so empty)
|
|
await session.create_branch_from_turn(1, "test_branch")
|
|
|
|
# Add items to the new branch
|
|
turn_2_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Branch question"},
|
|
{"role": "assistant", "content": "Branch answer"},
|
|
]
|
|
await session.add_items(turn_2_items)
|
|
|
|
# Add usage data for branch
|
|
usage_branch = Usage(requests=1, input_tokens=40, output_tokens=20, total_tokens=60)
|
|
run_result_branch = create_mock_run_result(usage_branch)
|
|
await session.store_run_usage(run_result_branch)
|
|
|
|
# Test get_items with branch_id parameter
|
|
main_items = await session.get_items(branch_id="main")
|
|
assert len(main_items) == 2
|
|
assert main_items[0].get("content") == "Main branch question"
|
|
|
|
current_items = await session.get_items() # Should get from current branch
|
|
assert len(current_items) == 2 # Only the items added to the branch (copied branch is empty)
|
|
|
|
# Test get_conversation_turns with branch_id
|
|
main_turns = await session.get_conversation_turns(branch_id="main")
|
|
assert len(main_turns) == 1
|
|
assert main_turns[0]["content"] == "Main branch question"
|
|
|
|
current_turns = await session.get_conversation_turns() # Should get from current branch
|
|
assert len(current_turns) == 1 # Only one turn in the current branch
|
|
|
|
# Test get_session_usage with branch_id
|
|
main_usage = await session.get_session_usage(branch_id="main")
|
|
assert main_usage is not None
|
|
assert main_usage["total_turns"] == 1
|
|
|
|
all_usage = await session.get_session_usage() # Should get from all branches
|
|
assert all_usage is not None
|
|
assert all_usage["total_turns"] == 2 # Main branch has 1, current branch has 1
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_branch_error_handling():
|
|
"""Test error handling in branching operations."""
|
|
session_id = "branch_error_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Test creating branch from non-existent turn
|
|
with pytest.raises(ValueError, match="Turn 5 does not contain a user message"):
|
|
await session.create_branch_from_turn(5, "error_branch")
|
|
|
|
# Test switching to non-existent branch
|
|
with pytest.raises(ValueError, match="Branch 'nonexistent' does not exist"):
|
|
await session.switch_to_branch("nonexistent")
|
|
|
|
# Test deleting non-existent branch
|
|
with pytest.raises(ValueError, match="Branch 'nonexistent' does not exist"):
|
|
await session.delete_branch("nonexistent")
|
|
|
|
# Test deleting main branch
|
|
with pytest.raises(ValueError, match="Cannot delete the 'main' branch"):
|
|
await session.delete_branch("main")
|
|
|
|
# Test deleting empty branch ID
|
|
with pytest.raises(ValueError, match="Branch ID cannot be empty"):
|
|
await session.delete_branch("")
|
|
|
|
# Test deleting empty branch ID (whitespace only)
|
|
with pytest.raises(ValueError, match="Branch ID cannot be empty"):
|
|
await session.delete_branch(" ")
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_branch_deletion_with_force():
|
|
"""Test branch deletion with force parameter."""
|
|
session_id = "force_delete_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Add items to main branch
|
|
await session.add_items([{"role": "user", "content": "Main question"}])
|
|
await session.add_items([{"role": "user", "content": "Second question"}])
|
|
|
|
# Create and switch to a branch from turn 2
|
|
await session.create_branch_from_turn(2, "temp_branch")
|
|
assert session._current_branch_id == "temp_branch"
|
|
|
|
# Add some content to the branch so it exists
|
|
await session.add_items([{"role": "user", "content": "Branch question"}])
|
|
|
|
# Verify branch exists
|
|
branches = await session.list_branches()
|
|
branch_ids = [b["branch_id"] for b in branches]
|
|
assert "temp_branch" in branch_ids
|
|
|
|
# Try to delete current branch without force (should fail)
|
|
with pytest.raises(ValueError, match="Cannot delete current branch"):
|
|
await session.delete_branch("temp_branch")
|
|
|
|
# Delete current branch with force (should succeed and switch to main)
|
|
await session.delete_branch("temp_branch", force=True)
|
|
|
|
# Verify we're back on main branch
|
|
assert session._current_branch_id == "main"
|
|
|
|
# Verify branch is deleted
|
|
branches_after = await session.list_branches()
|
|
assert len(branches_after) == 1
|
|
assert branches_after[0]["branch_id"] == "main"
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_get_items_with_parameters():
|
|
"""Test get_items with new parameters (include_inactive, branch_id)."""
|
|
session_id = "get_items_params_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Add items to main branch
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "First question"},
|
|
{"role": "assistant", "content": "First answer"},
|
|
{"role": "user", "content": "Second question"},
|
|
{"role": "assistant", "content": "Second answer"},
|
|
]
|
|
await session.add_items(items)
|
|
|
|
# Test get_items with limit (gets most recent N items)
|
|
limited_items = await session.get_items(limit=2)
|
|
assert len(limited_items) == 2
|
|
assert limited_items[0].get("content") == "Second question" # Most recent first
|
|
assert limited_items[1].get("content") == "Second answer"
|
|
|
|
# Test get_items with branch_id
|
|
main_items = await session.get_items(branch_id="main")
|
|
assert len(main_items) == 4
|
|
|
|
# Test get_items (no longer has include_inactive parameter)
|
|
all_items = await session.get_items()
|
|
assert len(all_items) == 4
|
|
|
|
# Create a branch from turn 2 and test branch-specific get_items
|
|
await session.create_branch_from_turn(2, "test_branch")
|
|
|
|
# Add items to branch
|
|
branch_items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Branch question"},
|
|
{"role": "assistant", "content": "Branch answer"},
|
|
]
|
|
await session.add_items(branch_items)
|
|
|
|
# Test getting items from specific branch (should include copied items + new items)
|
|
branch_items_result = await session.get_items(branch_id="test_branch")
|
|
assert len(branch_items_result) == 4 # 2 copied from main (before turn 2) + 2 new items
|
|
|
|
# Test getting items from main branch while on different branch
|
|
main_items_from_branch = await session.get_items(branch_id="main")
|
|
assert len(main_items_from_branch) == 4
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_usage_tracking_storage(agent: Agent, usage_data: Usage):
|
|
"""Test usage data storage and retrieval."""
|
|
session_id = "usage_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Simulate adding items for turn 1 to increment turn counter
|
|
await session.add_items([{"role": "user", "content": "First turn"}])
|
|
run_result_1 = create_mock_run_result(usage_data)
|
|
await session.store_run_usage(run_result_1)
|
|
|
|
# Create different usage data for turn 2
|
|
usage_data_2 = Usage(
|
|
requests=2,
|
|
input_tokens=75,
|
|
output_tokens=45,
|
|
total_tokens=120,
|
|
input_tokens_details=InputTokensDetails.model_validate(
|
|
{"cache_write_tokens": 0, "cached_tokens": 20}
|
|
),
|
|
output_tokens_details=OutputTokensDetails(reasoning_tokens=15),
|
|
)
|
|
|
|
# Simulate adding items for turn 2 to increment turn counter
|
|
await session.add_items([{"role": "user", "content": "Second turn"}])
|
|
run_result_2 = create_mock_run_result(usage_data_2)
|
|
await session.store_run_usage(run_result_2)
|
|
|
|
# Test session-level usage aggregation
|
|
session_usage = await session.get_session_usage()
|
|
assert session_usage is not None
|
|
assert session_usage["requests"] == 3 # 1 + 2
|
|
assert session_usage["total_tokens"] == 200 # 80 + 120
|
|
assert session_usage["input_tokens"] == 125 # 50 + 75
|
|
assert session_usage["output_tokens"] == 75 # 30 + 45
|
|
assert session_usage["total_turns"] == 2
|
|
|
|
# Test turn-level usage retrieval
|
|
turn_1_usage = await session.get_turn_usage(1)
|
|
assert isinstance(turn_1_usage, dict)
|
|
assert turn_1_usage["requests"] == 1
|
|
assert turn_1_usage["total_tokens"] == 80
|
|
assert turn_1_usage["input_tokens_details"]["cached_tokens"] == 10
|
|
assert turn_1_usage["output_tokens_details"]["reasoning_tokens"] == 5
|
|
|
|
turn_2_usage = await session.get_turn_usage(2)
|
|
assert isinstance(turn_2_usage, dict)
|
|
assert turn_2_usage["requests"] == 2
|
|
assert turn_2_usage["total_tokens"] == 120
|
|
assert turn_2_usage["input_tokens_details"]["cached_tokens"] == 20
|
|
assert turn_2_usage["output_tokens_details"]["reasoning_tokens"] == 15
|
|
|
|
# Test getting all turn usage
|
|
all_turn_usage = await session.get_turn_usage()
|
|
assert isinstance(all_turn_usage, list)
|
|
assert len(all_turn_usage) == 2
|
|
assert all_turn_usage[0]["user_turn_number"] == 1
|
|
assert all_turn_usage[1]["user_turn_number"] == 2
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_runner_integration_with_usage_tracking(agent: Agent):
|
|
"""Test integration with Runner and automatic usage tracking pattern."""
|
|
session_id = "integration_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
async def store_session_usage(result: Any, session: AdvancedSQLiteSession):
|
|
"""Helper function to store usage after runner completes."""
|
|
try:
|
|
await session.store_run_usage(result)
|
|
except Exception:
|
|
# Ignore errors in test helper
|
|
pass
|
|
|
|
# Set up fake model responses
|
|
assert isinstance(agent.model, FakeModel)
|
|
fake_model = agent.model
|
|
fake_model.set_next_output([get_text_message("San Francisco")])
|
|
|
|
# First turn
|
|
result1 = await Runner.run(
|
|
agent,
|
|
"What city is the Golden Gate Bridge in?",
|
|
session=session,
|
|
)
|
|
assert result1.final_output == "San Francisco"
|
|
await store_session_usage(result1, session)
|
|
|
|
# Second turn
|
|
fake_model.set_next_output([get_text_message("California")])
|
|
result2 = await Runner.run(agent, "What state is it in?", session=session)
|
|
assert result2.final_output == "California"
|
|
await store_session_usage(result2, session)
|
|
|
|
# Verify conversation structure
|
|
conversation_turns = await session.get_conversation_by_turns()
|
|
assert len(conversation_turns) == 2
|
|
|
|
# Verify usage was tracked
|
|
session_usage = await session.get_session_usage()
|
|
assert session_usage is not None
|
|
assert session_usage["total_turns"] == 2
|
|
# FakeModel doesn't generate realistic usage data, so we just check structure exists
|
|
assert "requests" in session_usage
|
|
assert "total_tokens" in session_usage
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_sequence_ordering():
|
|
"""Test that sequence ordering works correctly even with same timestamps."""
|
|
session_id = "sequence_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Add multiple items quickly to test sequence ordering
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Message 1"},
|
|
{"role": "assistant", "content": "Response 1"},
|
|
{"role": "user", "content": "Message 2"},
|
|
{"role": "assistant", "content": "Response 2"},
|
|
]
|
|
await session.add_items(items)
|
|
|
|
# Get items and verify order is preserved
|
|
retrieved = await session.get_items()
|
|
assert len(retrieved) == 4
|
|
assert retrieved[0].get("content") == "Message 1"
|
|
assert retrieved[1].get("content") == "Response 1"
|
|
assert retrieved[2].get("content") == "Message 2"
|
|
assert retrieved[3].get("content") == "Response 2"
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_conversation_structure_with_multiple_turns():
|
|
"""Test conversation structure tracking with multiple user turns."""
|
|
session_id = "multi_turn_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Turn 1
|
|
turn_1: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi!"},
|
|
]
|
|
await session.add_items(turn_1)
|
|
|
|
# Turn 2
|
|
turn_2: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "How are you?"},
|
|
{"type": "function_call", "name": "mood_check", "arguments": "{}"}, # type: ignore
|
|
{"type": "function_call_output", "output": "I'm good"}, # type: ignore
|
|
{"role": "assistant", "content": "I'm doing well!"},
|
|
]
|
|
await session.add_items(turn_2)
|
|
|
|
# Turn 3
|
|
turn_3: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Goodbye"},
|
|
{"role": "assistant", "content": "See you later!"},
|
|
]
|
|
await session.add_items(turn_3)
|
|
|
|
# Verify conversation structure
|
|
conversation_turns = await session.get_conversation_by_turns()
|
|
assert len(conversation_turns) == 3
|
|
|
|
# Turn 1 should have 2 items
|
|
assert len(conversation_turns[1]) == 2
|
|
assert conversation_turns[1][0]["type"] == "user"
|
|
assert conversation_turns[1][1]["type"] == "assistant"
|
|
|
|
# Turn 2 should have 4 items including tool calls
|
|
assert len(conversation_turns[2]) == 4
|
|
turn_2_types = [item["type"] for item in conversation_turns[2]]
|
|
assert "user" in turn_2_types
|
|
assert "function_call" in turn_2_types
|
|
assert "function_call_output" in turn_2_types
|
|
assert "assistant" in turn_2_types
|
|
|
|
# Turn 3 should have 2 items
|
|
assert len(conversation_turns[3]) == 2
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_empty_session_operations():
|
|
"""Test operations on empty sessions."""
|
|
session_id = "empty_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Test getting items from empty session
|
|
items = await session.get_items()
|
|
assert len(items) == 0
|
|
|
|
# Test getting conversation from empty session
|
|
conversation = await session.get_conversation_by_turns()
|
|
assert len(conversation) == 0
|
|
|
|
# Test getting tool usage from empty session
|
|
tool_usage = await session.get_tool_usage()
|
|
assert len(tool_usage) == 0
|
|
|
|
# Test getting session usage from empty session
|
|
session_usage = await session.get_session_usage()
|
|
assert session_usage is None
|
|
|
|
# Test getting turns from empty session
|
|
turns = await session.get_conversation_turns()
|
|
assert len(turns) == 0
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_json_serialization_edge_cases(usage_data: Usage):
|
|
"""Test edge cases in JSON serialization of usage data."""
|
|
session_id = "json_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Test with normal usage data (need to add user message first to create turn)
|
|
await session.add_items([{"role": "user", "content": "First test"}])
|
|
run_result_1 = create_mock_run_result(usage_data)
|
|
await session.store_run_usage(run_result_1)
|
|
|
|
# Test with None usage data
|
|
run_result_none = create_mock_run_result(None)
|
|
await session.store_run_usage(run_result_none)
|
|
|
|
# Test with usage data missing details
|
|
minimal_usage = Usage(
|
|
requests=1,
|
|
input_tokens=10,
|
|
output_tokens=5,
|
|
total_tokens=15,
|
|
)
|
|
await session.add_items([{"role": "user", "content": "Second test"}])
|
|
run_result_2 = create_mock_run_result(minimal_usage)
|
|
await session.store_run_usage(run_result_2)
|
|
|
|
# Verify we can retrieve the data
|
|
turn_1_usage = await session.get_turn_usage(1)
|
|
assert isinstance(turn_1_usage, dict)
|
|
assert turn_1_usage["requests"] == 1
|
|
assert turn_1_usage["input_tokens_details"]["cached_tokens"] == 10
|
|
|
|
turn_2_usage = await session.get_turn_usage(2)
|
|
assert isinstance(turn_2_usage, dict)
|
|
assert turn_2_usage["requests"] == 1
|
|
# Should have default values for minimal data (Usage class provides defaults)
|
|
assert turn_2_usage["input_tokens_details"]["cached_tokens"] == 0
|
|
assert turn_2_usage["output_tokens_details"]["reasoning_tokens"] == 0
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_session_isolation():
|
|
"""Test that different session IDs maintain separate data."""
|
|
session1 = AdvancedSQLiteSession(session_id="session_1", create_tables=True)
|
|
session2 = AdvancedSQLiteSession(session_id="session_2", create_tables=True)
|
|
|
|
# Add data to session 1
|
|
await session1.add_items([{"role": "user", "content": "Session 1 message"}])
|
|
|
|
# Add data to session 2
|
|
await session2.add_items([{"role": "user", "content": "Session 2 message"}])
|
|
|
|
# Verify isolation
|
|
session1_items = await session1.get_items()
|
|
session2_items = await session2.get_items()
|
|
|
|
assert len(session1_items) == 1
|
|
assert len(session2_items) == 1
|
|
assert session1_items[0].get("content") == "Session 1 message"
|
|
assert session2_items[0].get("content") == "Session 2 message"
|
|
|
|
# Test conversation structure isolation
|
|
session1_turns = await session1.get_conversation_by_turns()
|
|
session2_turns = await session2.get_conversation_by_turns()
|
|
|
|
assert len(session1_turns) == 1
|
|
assert len(session2_turns) == 1
|
|
|
|
session1.close()
|
|
session2.close()
|
|
|
|
|
|
async def test_error_handling_in_usage_tracking(usage_data: Usage):
|
|
"""Test that usage tracking errors don't break the main flow."""
|
|
session_id = "error_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Test normal operation
|
|
run_result = create_mock_run_result(usage_data)
|
|
await session.store_run_usage(run_result)
|
|
|
|
# Close the session to simulate database errors
|
|
session.close()
|
|
|
|
# This should not raise an exception (error should be caught)
|
|
await session.store_run_usage(run_result)
|
|
|
|
|
|
async def test_advanced_tool_name_extraction():
|
|
"""Test advanced tool name extraction for different tool types."""
|
|
session_id = "advanced_tool_names_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Add items with various tool types and naming patterns
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Use various tools"},
|
|
# MCP tools with server labels
|
|
{"type": "mcp_call", "server_label": "filesystem", "name": "read_file", "arguments": "{}"}, # type: ignore
|
|
{
|
|
"type": "mcp_approval_request",
|
|
"server_label": "database",
|
|
"name": "execute_query",
|
|
"arguments": "{}",
|
|
}, # type: ignore
|
|
# Built-in tool types
|
|
{"type": "computer_call", "arguments": "{}"}, # type: ignore
|
|
{"type": "file_search_call", "arguments": "{}"}, # type: ignore
|
|
{"type": "web_search_call", "arguments": "{}"}, # type: ignore
|
|
{"type": "code_interpreter_call", "arguments": "{}"}, # type: ignore
|
|
# Regular function calls
|
|
{"type": "function_call", "name": "calculator", "arguments": "{}"}, # type: ignore
|
|
{"type": "custom_tool_call", "name": "custom_tool", "arguments": "{}"}, # type: ignore
|
|
]
|
|
await session.add_items(items)
|
|
|
|
# Get conversation structure and verify tool names
|
|
conversation_turns = await session.get_conversation_by_turns()
|
|
turn_items = conversation_turns[1]
|
|
|
|
tool_items = [item for item in turn_items if item["tool_name"]]
|
|
tool_names = [item["tool_name"] for item in tool_items]
|
|
|
|
# Verify MCP tools get server_label.name format
|
|
assert "filesystem.read_file" in tool_names
|
|
assert "database.execute_query" in tool_names
|
|
|
|
# Verify built-in tools use their type as name
|
|
assert "computer_call" in tool_names
|
|
assert "file_search_call" in tool_names
|
|
assert "web_search_call" in tool_names
|
|
assert "code_interpreter_call" in tool_names
|
|
|
|
# Verify regular function calls use their name
|
|
assert "calculator" in tool_names
|
|
assert "custom_tool" in tool_names
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_branch_usage_tracking():
|
|
"""Test usage tracking across different branches."""
|
|
session_id = "branch_usage_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Add items and usage to main branch
|
|
await session.add_items([{"role": "user", "content": "Main question"}])
|
|
usage_main = Usage(requests=1, input_tokens=50, output_tokens=30, total_tokens=80)
|
|
run_result_main = create_mock_run_result(usage_main)
|
|
await session.store_run_usage(run_result_main)
|
|
|
|
# Create a branch and add usage there
|
|
await session.create_branch_from_turn(1, "usage_branch")
|
|
await session.add_items([{"role": "user", "content": "Branch question"}])
|
|
usage_branch = Usage(requests=2, input_tokens=100, output_tokens=60, total_tokens=160)
|
|
run_result_branch = create_mock_run_result(usage_branch)
|
|
await session.store_run_usage(run_result_branch)
|
|
|
|
# Test branch-specific usage
|
|
main_usage = await session.get_session_usage(branch_id="main")
|
|
assert main_usage is not None
|
|
assert main_usage["requests"] == 1
|
|
assert main_usage["total_tokens"] == 80
|
|
assert main_usage["total_turns"] == 1
|
|
|
|
branch_usage = await session.get_session_usage(branch_id="usage_branch")
|
|
assert branch_usage is not None
|
|
assert branch_usage["requests"] == 2
|
|
assert branch_usage["total_tokens"] == 160
|
|
assert branch_usage["total_turns"] == 1
|
|
|
|
# Test total usage across all branches
|
|
total_usage = await session.get_session_usage()
|
|
assert total_usage is not None
|
|
assert total_usage["requests"] == 3 # 1 + 2
|
|
assert total_usage["total_tokens"] == 240 # 80 + 160
|
|
assert total_usage["total_turns"] == 2
|
|
|
|
# Test turn usage for specific branch
|
|
branch_turn_usage = await session.get_turn_usage(branch_id="usage_branch")
|
|
assert isinstance(branch_turn_usage, list)
|
|
assert len(branch_turn_usage) == 1
|
|
assert branch_turn_usage[0]["requests"] == 2
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_tool_name_extraction():
|
|
"""Test that tool names are correctly extracted from different item types."""
|
|
session_id = "tool_names_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Add items with different ways of specifying tool names
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Use tools please"}, # Need user message to create turn
|
|
{"type": "function_call", "name": "search_web", "arguments": "{}"}, # type: ignore
|
|
{"type": "function_call_output", "tool_name": "search_web", "output": "result"}, # type: ignore
|
|
{"type": "function_call", "name": "calculator", "arguments": "{}"}, # type: ignore
|
|
]
|
|
await session.add_items(items)
|
|
|
|
# Get conversation structure and verify tool names
|
|
conversation_turns = await session.get_conversation_by_turns()
|
|
turn_items = conversation_turns[1]
|
|
|
|
tool_items = [item for item in turn_items if item["tool_name"]]
|
|
tool_names = [item["tool_name"] for item in tool_items]
|
|
|
|
assert "search_web" in tool_names
|
|
assert "calculator" in tool_names
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_tool_execution_integration(agent: Agent):
|
|
"""Test integration with actual tool execution."""
|
|
session_id = "tool_integration_test"
|
|
session = AdvancedSQLiteSession(session_id=session_id, create_tables=True)
|
|
|
|
# Set up the fake model to trigger a tool call
|
|
fake_model = cast(FakeModel, agent.model)
|
|
fake_model.set_next_output(
|
|
[
|
|
{ # type: ignore
|
|
"type": "function_call",
|
|
"name": "test_tool",
|
|
"arguments": '{"query": "test query"}',
|
|
"call_id": "call_123",
|
|
}
|
|
]
|
|
)
|
|
|
|
# Then set the final response
|
|
fake_model.set_next_output([get_text_message("Tool executed successfully")])
|
|
|
|
# Run the agent
|
|
result = await Runner.run(
|
|
agent,
|
|
"Please use the test tool",
|
|
session=session,
|
|
)
|
|
|
|
# Verify the tool was executed
|
|
assert "Tool result for: test query" in str(result.new_items)
|
|
|
|
# Verify tool usage was tracked
|
|
tool_usage = await session.get_tool_usage()
|
|
assert len(tool_usage) > 0
|
|
|
|
session.close()
|
|
|
|
|
|
# ============================================================================
|
|
# SessionSettings Tests
|
|
# ============================================================================
|
|
|
|
|
|
async def test_session_settings_default():
|
|
"""Test that session_settings defaults to empty SessionSettings."""
|
|
from agents.memory import SessionSettings
|
|
|
|
session = AdvancedSQLiteSession(session_id="default_settings_test", create_tables=True)
|
|
|
|
# Should have default SessionSettings (inherited from SQLiteSession)
|
|
assert isinstance(session.session_settings, SessionSettings)
|
|
assert session.session_settings.limit is None
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_session_settings_constructor():
|
|
"""Test passing session_settings via constructor."""
|
|
from agents.memory import SessionSettings
|
|
|
|
session = AdvancedSQLiteSession(
|
|
session_id="constructor_settings_test",
|
|
create_tables=True,
|
|
session_settings=SessionSettings(limit=5),
|
|
)
|
|
|
|
assert session.session_settings is not None
|
|
assert session.session_settings.limit == 5
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_get_items_uses_session_settings_limit():
|
|
"""Test that get_items uses session_settings.limit as default."""
|
|
from agents.memory import SessionSettings
|
|
|
|
session = AdvancedSQLiteSession(
|
|
session_id="uses_settings_limit_test",
|
|
create_tables=True,
|
|
session_settings=SessionSettings(limit=3),
|
|
)
|
|
|
|
# Add 5 items
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": f"Message {i}"} for i in range(5)
|
|
]
|
|
await session.add_items(items)
|
|
|
|
# get_items() with no limit should use session_settings.limit=3
|
|
retrieved = await session.get_items()
|
|
assert len(retrieved) == 3
|
|
# Should get the last 3 items
|
|
assert retrieved[0].get("content") == "Message 2"
|
|
assert retrieved[1].get("content") == "Message 3"
|
|
assert retrieved[2].get("content") == "Message 4"
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_get_items_explicit_limit_overrides_session_settings():
|
|
"""Test that explicit limit parameter overrides session_settings."""
|
|
from agents.memory import SessionSettings
|
|
|
|
session = AdvancedSQLiteSession(
|
|
session_id="explicit_override_test",
|
|
create_tables=True,
|
|
session_settings=SessionSettings(limit=5),
|
|
)
|
|
|
|
# Add 10 items
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": f"Message {i}"} for i in range(10)
|
|
]
|
|
await session.add_items(items)
|
|
|
|
# Explicit limit=2 should override session_settings.limit=5
|
|
retrieved = await session.get_items(limit=2)
|
|
assert len(retrieved) == 2
|
|
assert retrieved[0].get("content") == "Message 8"
|
|
assert retrieved[1].get("content") == "Message 9"
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_session_settings_resolve():
|
|
"""Test SessionSettings.resolve() method."""
|
|
from agents.memory import SessionSettings
|
|
|
|
base = SessionSettings(limit=100)
|
|
override = SessionSettings(limit=50)
|
|
|
|
final = base.resolve(override)
|
|
|
|
assert final.limit == 50 # Override wins
|
|
assert base.limit == 100 # Original unchanged
|
|
|
|
# Resolving with None returns self
|
|
final_none = base.resolve(None)
|
|
assert final_none.limit == 100
|
|
|
|
|
|
async def test_runner_with_session_settings_override(agent: Agent):
|
|
"""Test that RunConfig can override session's default settings."""
|
|
from agents import RunConfig
|
|
from agents.memory import SessionSettings
|
|
|
|
# Session with default limit=100
|
|
session = AdvancedSQLiteSession(
|
|
session_id="runner_override_test",
|
|
create_tables=True,
|
|
session_settings=SessionSettings(limit=100),
|
|
)
|
|
|
|
# Add some history
|
|
items: list[TResponseInputItem] = [{"role": "user", "content": f"Turn {i}"} for i in range(10)]
|
|
await session.add_items(items)
|
|
|
|
# Use RunConfig to override limit to 2
|
|
assert isinstance(agent.model, FakeModel)
|
|
agent.model.set_next_output([get_text_message("Got it")])
|
|
|
|
await Runner.run(
|
|
agent,
|
|
"New question",
|
|
session=session,
|
|
run_config=RunConfig(
|
|
session_settings=SessionSettings(limit=2) # Override to 2
|
|
),
|
|
)
|
|
|
|
# Verify the agent received only the last 2 history items + new question
|
|
last_input = agent.model.last_turn_args["input"]
|
|
# Filter out the new "New question" input
|
|
history_items = [item for item in last_input if item.get("content") != "New question"]
|
|
# Should have 2 history items (last two from the 10 we added)
|
|
assert len(history_items) == 2
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_concurrent_add_items_preserves_message_structure_for_file_db():
|
|
"""Concurrent add_items calls should keep agent_messages and message_structure aligned."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "advanced_concurrent.db"
|
|
session = AdvancedSQLiteSession(
|
|
session_id="advanced_concurrent",
|
|
db_path=db_path,
|
|
create_tables=True,
|
|
)
|
|
|
|
async def add_batch(worker_id: int) -> list[str]:
|
|
contents = [f"worker-{worker_id}-message-{index}" for index in range(10)]
|
|
await session.add_items([{"role": "user", "content": content} for content in contents])
|
|
return contents
|
|
|
|
expected_batches = await asyncio.gather(*(add_batch(worker_id) for worker_id in range(8)))
|
|
expected_contents = {content for batch in expected_batches for content in batch}
|
|
|
|
retrieved_items = await session.get_items()
|
|
retrieved_contents = {
|
|
content
|
|
for item in retrieved_items
|
|
for content in [item.get("content")]
|
|
if isinstance(content, str)
|
|
}
|
|
|
|
assert retrieved_contents == expected_contents
|
|
assert len(retrieved_items) == len(expected_contents)
|
|
|
|
with session._locked_connection() as conn:
|
|
rows = conn.execute(
|
|
f"""
|
|
SELECT m.message_data
|
|
FROM {session.messages_table} m
|
|
JOIN message_structure s ON s.message_id = m.id
|
|
WHERE m.session_id = ?
|
|
ORDER BY s.sequence_number ASC
|
|
""",
|
|
(session.session_id,),
|
|
).fetchall()
|
|
|
|
structured_contents = {json.loads(message_data).get("content") for (message_data,) in rows}
|
|
|
|
assert structured_contents == expected_contents
|
|
assert len(rows) == len(expected_contents)
|
|
|
|
session.close()
|
|
|
|
|
|
async def test_output_tokens_details_persisted_when_input_details_missing():
|
|
"""Regression: output_tokens_details must persist even if input_tokens_details is None.
|
|
|
|
Previously the output serialization branch was nested inside the input branch,
|
|
silently dropping output_tokens_details whenever input_tokens_details was falsy
|
|
(e.g., when a provider populated only output details).
|
|
"""
|
|
session = AdvancedSQLiteSession(session_id="output_only_usage", create_tables=True)
|
|
usage = Usage(
|
|
requests=1,
|
|
input_tokens=10,
|
|
output_tokens=5,
|
|
total_tokens=15,
|
|
output_tokens_details=OutputTokensDetails(reasoning_tokens=42),
|
|
)
|
|
# Mimic providers that bypass validation and leave input_tokens_details unset.
|
|
object.__setattr__(usage, "input_tokens_details", None)
|
|
|
|
await session.add_items([{"role": "user", "content": "hi"}])
|
|
await session.store_run_usage(create_mock_run_result(usage))
|
|
|
|
turn_usage = await session.get_turn_usage(1)
|
|
assert isinstance(turn_usage, dict)
|
|
assert turn_usage["output_tokens_details"] == {"reasoning_tokens": 42}
|
|
assert turn_usage["input_tokens_details"] is None
|
|
session.close()
|