Skip to content

Commit b3ed91f

Browse files
habemahabema
committed
feat: StructuredSQLiteSession
1 parent 0a7bb1b commit b3ed91f

File tree

5 files changed

+1465
-0
lines changed

5 files changed

+1465
-0
lines changed

‎examples/memory/advanced_structured_session_example.py‎

Lines changed: 191 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,191 @@
1+
"""
2+
Advanced example demonstrating conversation branching with StructuredSQLiteSession.
3+
4+
This example shows how to use soft deletion for conversation editing/branching,
5+
allowing you to "undo" parts of a conversation and continue from any point.
6+
"""
7+
8+
import asyncio
9+
10+
from agents import Agent, Runner, function_tool
11+
from agents.extensions.memory import StructuredSQLiteSession
12+
13+
14+
@function_tool
15+
async def get_weather(city: str) -> str:
16+
if city.strip().lower() == "new york":
17+
return f"The weather in {city} is cloudy."
18+
return f"The weather in {city} is sunny."
19+
20+
21+
async def main():
22+
# Create an agent
23+
agent = Agent(
24+
name="Assistant",
25+
instructions="Reply very concisely.",
26+
tools=[get_weather],
27+
)
28+
29+
# Create a structured session instance
30+
session = StructuredSQLiteSession("conversation_advanced")
31+
32+
print("=== Advanced Structured Session: Conversation Branching ===")
33+
print("This example demonstrates conversation editing and branching.\n")
34+
35+
# Build initial conversation
36+
print("Building initial conversation...")
37+
38+
# Turn 1
39+
print("Turn 1: User asks about Golden Gate Bridge")
40+
result = await Runner.run(
41+
agent,
42+
"What city is the Golden Gate Bridge in?",
43+
session=session,
44+
)
45+
print(f"Assistant: {result.final_output}")
46+
await session.store_run_usage(result)
47+
48+
# Turn 2
49+
print("Turn 2: User asks about weather")
50+
result = await Runner.run(
51+
agent,
52+
"What's the weather in that city?",
53+
session=session,
54+
)
55+
print(f"Assistant: {result.final_output}")
56+
await session.store_run_usage(result)
57+
58+
# Turn 3
59+
print("Turn 3: User asks about population")
60+
result = await Runner.run(
61+
agent,
62+
"What's the population of that city?",
63+
session=session,
64+
)
65+
print(f"Assistant: {result.final_output}")
66+
await session.store_run_usage(result)
67+
68+
print("\n=== Original Conversation Complete ===")
69+
70+
# Show current conversation
71+
print("Current conversation:")
72+
current_items = await session.get_items()
73+
for i, item in enumerate(current_items, 1):
74+
role = str(item.get("role", item.get("type", "unknown")))
75+
if item.get("type") == "function_call":
76+
content = f"{item.get('name', 'unknown')}({item.get('arguments', '{}')})"
77+
elif item.get("type") == "function_call_output":
78+
content = str(item.get("output", ""))
79+
else:
80+
content = str(item.get("content", item.get("output", "")))
81+
print(f" {i}. {role}: {content}")
82+
83+
print(f"\nTotal items: {len(current_items)}")
84+
85+
# Demonstrate conversation branching
86+
print("\n=== Conversation Branching Demo ===")
87+
print("Let's say we want to edit the conversation from turn 2 onwards...")
88+
89+
# Soft delete from turn 2 to create a branch point
90+
print("\nSoft deleting from turn 2 onwards to create branch point...")
91+
deleted = await session.soft_delete_from_turn(2)
92+
print(f"Deleted: {deleted}")
93+
94+
# Show only active items (turn 1 only)
95+
active_items = await session.get_items()
96+
print(f"Active items after deletion: {len(active_items)}")
97+
print("Active conversation (turn 1 only):")
98+
for i, item in enumerate(active_items, 1):
99+
role = str(item.get("role", item.get("type", "unknown")))
100+
if item.get("type") == "function_call":
101+
content = f"{item.get('name', 'unknown')}({item.get('arguments', '{}')})"
102+
elif item.get("type") == "function_call_output":
103+
content = str(item.get("output", ""))
104+
else:
105+
content = str(item.get("content", item.get("output", "")))
106+
print(f" {i}. {role}: {content}")
107+
108+
# Create a new branch
109+
print("\nCreating new conversation branch...")
110+
print("Turn 2 (new branch): User asks about New York instead")
111+
result = await Runner.run(
112+
agent,
113+
"Actually, what's the weather in New York instead?",
114+
session=session,
115+
)
116+
print(f"Assistant: {result.final_output}")
117+
await session.store_run_usage(result)
118+
119+
# Continue the new branch
120+
print("Turn 3 (new branch): User asks about NYC attractions")
121+
result = await Runner.run(
122+
agent,
123+
"What are some famous attractions in New York?",
124+
session=session,
125+
)
126+
print(f"Assistant: {result.final_output}")
127+
await session.store_run_usage(result)
128+
129+
# Show the new conversation
130+
print("\n=== New Conversation Branch ===")
131+
new_conversation = await session.get_items()
132+
print("New conversation with branch:")
133+
for i, item in enumerate(new_conversation, 1):
134+
role = str(item.get("role", item.get("type", "unknown")))
135+
if item.get("type") == "function_call":
136+
content = f"{item.get('name', 'unknown')}({item.get('arguments', '{}')})"
137+
elif item.get("type") == "function_call_output":
138+
content = str(item.get("output", ""))
139+
else:
140+
content = str(item.get("content", item.get("output", "")))
141+
print(f" {i}. {role}: {content}")
142+
143+
# Show that we can still access the original branch
144+
all_items = await session.get_items(include_inactive=True)
145+
print(f"\nTotal items including inactive (original + new branch): {len(all_items)}")
146+
147+
print("\n=== Conversation Structure Analysis ===")
148+
# Show conversation turns (active only)
149+
conversation_turns = await session.get_conversation_by_turns()
150+
print("Active conversation turns:")
151+
for turn_num, items in conversation_turns.items():
152+
print(f" Turn {turn_num}: {len(items)} items")
153+
for item in items: # type: ignore
154+
if item["tool_name"]: # type: ignore
155+
print(
156+
f" - {item['type']} (tool: {item['tool_name']}) [active: {item['active']}]" # type: ignore
157+
)
158+
else:
159+
print(f" - {item['type']} [active: {item['active']}]") # type: ignore
160+
161+
# Show all conversation turns (including inactive)
162+
all_conversation_turns = await session.get_conversation_by_turns(include_inactive=True)
163+
print("\nAll conversation turns (including inactive):")
164+
for turn_num, items in all_conversation_turns.items():
165+
print(f" Turn {turn_num}: {len(items)} items")
166+
for item in items: # type: ignore
167+
status = "ACTIVE" if item["active"] else "INACTIVE" # type: ignore
168+
if item["tool_name"]: # type: ignore
169+
print(f" - {item['type']} (tool: {item['tool_name']}) [{status}]") # type: ignore
170+
else:
171+
print(f" - {item['type']} [{status}]")
172+
173+
print("\n=== Reactivation Demo ===")
174+
print("We can also reactivate the original conversation...")
175+
176+
# Reactivate the original conversation
177+
reactivated = await session.reactivate_from_turn(2)
178+
print(f"Reactivated: {reactivated}")
179+
180+
# Show all active items now
181+
all_active = await session.get_items()
182+
print(f"All active items after reactivation: {len(all_active)}")
183+
184+
print("\n=== Advanced Example Complete ===")
185+
print("This demonstrates how soft deletion enables conversation editing/branching!")
186+
print("You can 'undo' parts of a conversation and continue from any point.")
187+
print("Perfect for building conversational AI systems with editing capabilities.")
188+
189+
190+
if __name__ == "__main__":
191+
asyncio.run(main())

‎examples/memory/basic_structured_session_example.py‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
"""
2+
Basic example demonstrating structured session memory functionality.
3+
4+
This example shows how to use StructuredSQLiteSession for conversation tracking
5+
with usage statistics and turn-based organization.
6+
"""
7+
8+
import asyncio
9+
10+
from agents import Agent, Runner, function_tool
11+
from agents.extensions.memory import StructuredSQLiteSession
12+
13+
14+
@function_tool
15+
async def get_weather(city: str) -> str:
16+
if city.strip().lower() == "new york":
17+
return f"The weather in {city} is cloudy."
18+
return f"The weather in {city} is sunny."
19+
20+
21+
async def main():
22+
# Create an agent
23+
agent = Agent(
24+
name="Assistant",
25+
instructions="Reply very concisely.",
26+
tools=[get_weather],
27+
)
28+
29+
# Create a structured session instance
30+
session = StructuredSQLiteSession("conversation_basic")
31+
32+
print("=== Basic Structured Session Example ===")
33+
print("The agent will remember previous messages with structured tracking.\n")
34+
35+
# First turn
36+
print("First turn:")
37+
print("User: What city is the Golden Gate Bridge in?")
38+
result = await Runner.run(
39+
agent,
40+
"What city is the Golden Gate Bridge in?",
41+
session=session,
42+
)
43+
print(f"Assistant: {result.final_output}")
44+
print(f"Usage: {result.context_wrapper.usage.total_tokens} tokens")
45+
46+
# Store usage data automatically
47+
await session.store_run_usage(result)
48+
print()
49+
50+
# Second turn - continuing the conversation
51+
print("Second turn:")
52+
print("User: What's the weather in that city?")
53+
result = await Runner.run(
54+
agent,
55+
"What's the weather in that city?",
56+
session=session,
57+
)
58+
print(f"Assistant: {result.final_output}")
59+
print(f"Usage: {result.context_wrapper.usage.total_tokens} tokens")
60+
61+
# Store usage data automatically
62+
await session.store_run_usage(result)
63+
print()
64+
65+
print("=== Usage Tracking Demo ===")
66+
session_usage = await session.get_session_usage()
67+
if session_usage:
68+
print("Session Usage (aggregated from turns):")
69+
print(f" Total requests: {session_usage['requests']}")
70+
print(f" Total tokens: {session_usage['total_tokens']}")
71+
print(f" Input tokens: {session_usage['input_tokens']}")
72+
print(f" Output tokens: {session_usage['output_tokens']}")
73+
print(f" Total turns: {session_usage['total_turns']}")
74+
75+
# Show usage by turn
76+
turn_usage_list = await session.get_turn_usage()
77+
if turn_usage_list and isinstance(turn_usage_list, list):
78+
print("\nUsage by turn:")
79+
for turn_data in turn_usage_list:
80+
turn_num = turn_data["user_turn_number"]
81+
tokens = turn_data["total_tokens"]
82+
print(f" Turn {turn_num}: {tokens} tokens")
83+
else:
84+
print("No usage data found.")
85+
86+
print("\n=== Structured Query Demo ===")
87+
conversation_turns = await session.get_conversation_by_turns()
88+
print("Conversation by turns:")
89+
for turn_num, items in conversation_turns.items():
90+
print(f" Turn {turn_num}: {len(items)} items")
91+
for item in items:
92+
if item["tool_name"]:
93+
print(f" - {item['type']} (tool: {item['tool_name']})")
94+
else:
95+
print(f" - {item['type']}")
96+
97+
# Show tool usage
98+
tool_usage = await session.get_tool_usage()
99+
if tool_usage:
100+
print("\nTool usage:")
101+
for tool_name, count, turn in tool_usage:
102+
print(f" {tool_name}: used {count} times in turn {turn}")
103+
else:
104+
print("\nNo tool usage found.")
105+
106+
print("\n=== Basic Example Complete ===")
107+
print("Session provides structured tracking with usage analytics!")
108+
109+
110+
if __name__ == "__main__":
111+
asyncio.run(main())

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,9 @@
99
from __future__ import annotations
1010

1111
from .sqlalchemy_session import SQLAlchemySession # noqa: F401
12+
from .structured_sqlite_session import StructuredSQLiteSession # noqa: F401
1213

1314
__all__: list[str] = [
1415
"SQLAlchemySession",
16+
"StructuredSQLiteSession",
1517
]

0 commit comments

Comments
 (0)