Skip to content

Commit 4f54878

Browse files
habemaseratch
habema
and
seratch
authored
Feat: Add AdvancedSQLiteSession with conversation branching & usage tracking (openai#1662)
Co-authored-by: Kazuhiro Sera <[email protected]>
1 parent e87552a commit 4f54878

File tree

4 files changed

+2560
-0
lines changed

4 files changed

+2560
-0
lines changed

‎examples/memory/advanced_sqlite_session_example.py‎

Lines changed: 278 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,278 @@
1+
"""
2+
Comprehensive example demonstrating AdvancedSQLiteSession functionality.
3+
4+
This example shows both basic session memory features and advanced conversation
5+
branching capabilities, including usage statistics, turn-based organization,
6+
and multi-timeline conversation management.
7+
"""
8+
9+
import asyncio
10+
11+
from agents import Agent, Runner, function_tool
12+
from agents.extensions.memory import AdvancedSQLiteSession
13+
14+
15+
@function_tool
16+
async def get_weather(city: str) -> str:
17+
if city.strip().lower() == "new york":
18+
return f"The weather in {city} is cloudy."
19+
return f"The weather in {city} is sunny."
20+
21+
22+
async def main():
23+
# Create an agent
24+
agent = Agent(
25+
name="Assistant",
26+
instructions="Reply very concisely.",
27+
tools=[get_weather],
28+
)
29+
30+
# Create an advanced session instance
31+
session = AdvancedSQLiteSession(
32+
session_id="conversation_comprehensive",
33+
create_tables=True,
34+
)
35+
36+
print("=== AdvancedSQLiteSession Comprehensive Example ===")
37+
print("This example demonstrates both basic and advanced session features.\n")
38+
39+
# === PART 1: Basic Session Functionality ===
40+
print("=== PART 1: Basic Session Memory ===")
41+
print("The agent will remember previous messages with structured tracking.\n")
42+
43+
# First turn
44+
print("First turn:")
45+
print("User: What city is the Golden Gate Bridge in?")
46+
result = await Runner.run(
47+
agent,
48+
"What city is the Golden Gate Bridge in?",
49+
session=session,
50+
)
51+
print(f"Assistant: {result.final_output}")
52+
print(f"Usage: {result.context_wrapper.usage.total_tokens} tokens")
53+
54+
# Store usage data automatically
55+
await session.store_run_usage(result)
56+
print()
57+
58+
# Second turn - continuing the conversation
59+
print("Second turn:")
60+
print("User: What's the weather in that city?")
61+
result = await Runner.run(
62+
agent,
63+
"What's the weather in that city?",
64+
session=session,
65+
)
66+
print(f"Assistant: {result.final_output}")
67+
print(f"Usage: {result.context_wrapper.usage.total_tokens} tokens")
68+
69+
# Store usage data automatically
70+
await session.store_run_usage(result)
71+
print()
72+
73+
# Third turn
74+
print("Third turn:")
75+
print("User: What's the population of that city?")
76+
result = await Runner.run(
77+
agent,
78+
"What's the population of that city?",
79+
session=session,
80+
)
81+
print(f"Assistant: {result.final_output}")
82+
print(f"Usage: {result.context_wrapper.usage.total_tokens} tokens")
83+
84+
# Store usage data automatically
85+
await session.store_run_usage(result)
86+
print()
87+
88+
# === PART 2: Usage Tracking and Analytics ===
89+
print("=== PART 2: Usage Tracking and Analytics ===")
90+
session_usage = await session.get_session_usage()
91+
if session_usage:
92+
print("Session Usage (aggregated from turns):")
93+
print(f" Total requests: {session_usage['requests']}")
94+
print(f" Total tokens: {session_usage['total_tokens']}")
95+
print(f" Input tokens: {session_usage['input_tokens']}")
96+
print(f" Output tokens: {session_usage['output_tokens']}")
97+
print(f" Total turns: {session_usage['total_turns']}")
98+
99+
# Show usage by turn
100+
turn_usage_list = await session.get_turn_usage()
101+
if turn_usage_list and isinstance(turn_usage_list, list):
102+
print("\nUsage by turn:")
103+
for turn_data in turn_usage_list:
104+
turn_num = turn_data["user_turn_number"]
105+
tokens = turn_data["total_tokens"]
106+
print(f" Turn {turn_num}: {tokens} tokens")
107+
else:
108+
print("No usage data found.")
109+
110+
print("\n=== Structured Query Demo ===")
111+
conversation_turns = await session.get_conversation_by_turns()
112+
print("Conversation by turns:")
113+
for turn_num, items in conversation_turns.items():
114+
print(f" Turn {turn_num}: {len(items)} items")
115+
for item in items:
116+
if item["tool_name"]:
117+
print(f" - {item['type']} (tool: {item['tool_name']})")
118+
else:
119+
print(f" - {item['type']}")
120+
121+
# Show tool usage
122+
tool_usage = await session.get_tool_usage()
123+
if tool_usage:
124+
print("\nTool usage:")
125+
for tool_name, count, turn in tool_usage:
126+
print(f" {tool_name}: used {count} times in turn {turn}")
127+
else:
128+
print("\nNo tool usage found.")
129+
130+
print("\n=== Original Conversation Complete ===")
131+
132+
# Show current conversation
133+
print("Current conversation:")
134+
current_items = await session.get_items()
135+
for i, item in enumerate(current_items, 1):
136+
role = str(item.get("role", item.get("type", "unknown")))
137+
if item.get("type") == "function_call":
138+
content = f"{item.get('name', 'unknown')}({item.get('arguments', '{}')})"
139+
elif item.get("type") == "function_call_output":
140+
content = str(item.get("output", ""))
141+
else:
142+
content = str(item.get("content", item.get("output", "")))
143+
print(f" {i}. {role}: {content}")
144+
145+
print(f"\nTotal items: {len(current_items)}")
146+
147+
# === PART 3: Conversation Branching ===
148+
print("\n=== PART 3: Conversation Branching ===")
149+
print("Let's explore a different path from turn 2...")
150+
151+
# Show available turns for branching
152+
print("\nAvailable turns for branching:")
153+
turns = await session.get_conversation_turns()
154+
for turn in turns:
155+
print(f" Turn {turn['turn']}: {turn['content']}")
156+
157+
# Create a branch from turn 2
158+
print("\nCreating new branch from turn 2...")
159+
branch_id = await session.create_branch_from_turn(2)
160+
print(f"Created branch: {branch_id}")
161+
162+
# Show what's in the new branch (should have conversation up to turn 2)
163+
branch_items = await session.get_items()
164+
print(f"Items copied to new branch: {len(branch_items)}")
165+
print("New branch contains:")
166+
for i, item in enumerate(branch_items, 1):
167+
role = str(item.get("role", item.get("type", "unknown")))
168+
if item.get("type") == "function_call":
169+
content = f"{item.get('name', 'unknown')}({item.get('arguments', '{}')})"
170+
elif item.get("type") == "function_call_output":
171+
content = str(item.get("output", ""))
172+
else:
173+
content = str(item.get("content", item.get("output", "")))
174+
print(f" {i}. {role}: {content}")
175+
176+
# Continue conversation in new branch
177+
print("\nContinuing conversation in new branch...")
178+
print("Turn 2 (new branch): User asks about New York instead")
179+
result = await Runner.run(
180+
agent,
181+
"Actually, what's the weather in New York instead?",
182+
session=session,
183+
)
184+
print(f"Assistant: {result.final_output}")
185+
await session.store_run_usage(result)
186+
187+
# Continue the new branch
188+
print("Turn 3 (new branch): User asks about NYC attractions")
189+
result = await Runner.run(
190+
agent,
191+
"What are some famous attractions in New York?",
192+
session=session,
193+
)
194+
print(f"Assistant: {result.final_output}")
195+
await session.store_run_usage(result)
196+
197+
# Show the new conversation
198+
print("\n=== New Conversation Branch ===")
199+
new_conversation = await session.get_items()
200+
print("New conversation with branch:")
201+
for i, item in enumerate(new_conversation, 1):
202+
role = str(item.get("role", item.get("type", "unknown")))
203+
if item.get("type") == "function_call":
204+
content = f"{item.get('name', 'unknown')}({item.get('arguments', '{}')})"
205+
elif item.get("type") == "function_call_output":
206+
content = str(item.get("output", ""))
207+
else:
208+
content = str(item.get("content", item.get("output", "")))
209+
print(f" {i}. {role}: {content}")
210+
211+
print(f"\nTotal items in new branch: {len(new_conversation)}")
212+
213+
# === PART 4: Branch Management ===
214+
print("\n=== PART 4: Branch Management ===")
215+
# Show all branches
216+
branches = await session.list_branches()
217+
print("All branches in this session:")
218+
for branch in branches:
219+
current = " (current)" if branch["is_current"] else ""
220+
print(
221+
f" {branch['branch_id']}: {branch['user_turns']} user turns, {branch['message_count']} total messages{current}"
222+
)
223+
224+
# Show conversation turns in current branch
225+
print("\nConversation turns in current branch:")
226+
current_turns = await session.get_conversation_turns()
227+
for turn in current_turns:
228+
print(f" Turn {turn['turn']}: {turn['content']}")
229+
230+
print("\n=== Branch Switching Demo ===")
231+
print("We can switch back to the main branch...")
232+
233+
# Switch back to main branch
234+
await session.switch_to_branch("main")
235+
print("Switched to main branch")
236+
237+
# Show what's in main branch
238+
main_items = await session.get_items()
239+
print(f"Items in main branch: {len(main_items)}")
240+
241+
# Switch back to new branch
242+
await session.switch_to_branch(branch_id)
243+
branch_items = await session.get_items()
244+
print(f"Items in new branch: {len(branch_items)}")
245+
246+
print("\n=== Final Summary ===")
247+
await session.switch_to_branch("main")
248+
main_final = len(await session.get_items())
249+
await session.switch_to_branch(branch_id)
250+
branch_final = len(await session.get_items())
251+
252+
print(f"Main branch items: {main_final}")
253+
print(f"New branch items: {branch_final}")
254+
255+
# Show that branches are completely independent
256+
print("\nBranches are completely independent:")
257+
print("- Main branch has full original conversation")
258+
print("- New branch has turn 1 + new conversation path")
259+
print("- No interference between branches!")
260+
261+
print("\n=== Comprehensive Example Complete ===")
262+
print("This demonstrates the full AdvancedSQLiteSession capabilities!")
263+
print("Key features:")
264+
print("- Structured conversation tracking with usage analytics")
265+
print("- Turn-based organization and querying")
266+
print("- Create branches from any user message")
267+
print("- Branches inherit conversation history up to the branch point")
268+
print("- Complete branch isolation - no interference between branches")
269+
print("- Easy branch switching and management")
270+
print("- No complex soft deletion - clean branch-based architecture")
271+
print("- Perfect for building AI systems with conversation editing capabilities!")
272+
273+
# Cleanup
274+
session.close()
275+
276+
277+
if __name__ == "__main__":
278+
asyncio.run(main())

‎src/agents/extensions/memory/__init__.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
__all__: list[str] = [
1414
"EncryptedSession",
1515
"SQLAlchemySession",
16+
"AdvancedSQLiteSession",
1617
]
1718

1819

@@ -39,4 +40,14 @@ def __getattr__(name: str) -> Any:
3940
"Install it with: pip install openai-agents[sqlalchemy]"
4041
) from e
4142

43+
if name == "AdvancedSQLiteSession":
44+
try:
45+
from .advanced_sqlite_session import AdvancedSQLiteSession # noqa: F401
46+
47+
return AdvancedSQLiteSession
48+
except ModuleNotFoundError as e:
49+
raise ImportError(
50+
f"Failed to import AdvancedSQLiteSession: {e}"
51+
) from e
52+
4253
raise AttributeError(f"module {__name__} has no attribute {name}")

0 commit comments

Comments
 (0)