Time-Travel Replay (v1.0.8)

from pathlib import Path
import sys
ROOT = Path.cwd().resolve().parent if Path.cwd().name == 'notebooks' else Path.cwd().resolve()
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from shipit_agent.models import AgentEvent
from shipit_agent.tracing import InMemoryTraceStore
from shipit_agent.replay import TraceReplayer, diff_traces

store = InMemoryTraceStore()
trace_id = 'demo-run-1'
events = [
    AgentEvent('run_started', 'starting', {'user_prompt': 'Recommend a movie.'}),
    AgentEvent('tool_called', 'web_search', {'tool': 'web_search', 'query': 'best thriller 2024'}),
    AgentEvent('tool_completed', 'ok', {'text': 'Top hit: Dune Part Two — sci-fi, not thriller'}),
    AgentEvent('tool_called', 'web_search', {'tool': 'web_search', 'query': 'best 2024 movies'}),
    AgentEvent('tool_completed', 'ok', {'text': 'Top hit: Dune Part Two'}),
    AgentEvent('run_completed', 'done', {'output': 'Dune Part Two — strongly recommended.'}),
]
for ev in events:
    store.append_event(trace_id, ev)

print(f'saved {len(events)} events to trace {trace_id!r}')

replayer = TraceReplayer.from_store(store, trace_id)
print(f'trace_id: {replayer.trace_id}')
print(f'events:   {len(replayer)}')
print()
print('every event by type:')
for i, ev in enumerate(replayer.events):
    print(f'  {i}: {ev.type:<18} {ev.message}')

print('user prompts:')
for idx, text in replayer.find_user_messages():
    print(f'  event {idx}: {text}')

print()
print('tool calls:', replayer.event_indices_by_type('tool_called'))
print('tool completions:', replayer.event_indices_by_type('tool_completed'))

for ev_idx in (0, 2, 4, 5):
    msgs = replayer.messages_at(ev_idx)
    print(f'after event {ev_idx} ({replayer.events[ev_idx].type}): {len(msgs)} message(s)')
    for m in msgs:
        print(f'  [{m.role}] {m.content[:60]}')

fork = replayer.fork(
    at_event=1,
    edit_user_message='Recommend a thriller specifically — not sci-fi or fantasy.',
    extra_metadata={'reason': 'wrong-genre-fix'},
)

print('fork point:')
print(f'  source: {fork.fork.source_trace_id}')
print(f'  at event: {fork.fork.at_event}')
print(f'  edits: {fork.fork.edits}')
print()
print('resume prompt:', fork.user_prompt)
print(f'history seeded: {len(fork.messages)} message(s)')

# In a real app you'd resume on a real Agent. Here we use a tiny stub
# so the notebook runs without API keys.
from shipit_agent.models import AgentResult, Message

class StubAgent:
    def __init__(self):
        self.history = []
    def run(self, prompt, **kwargs):
        return AgentResult(
            output=f'Resumed with prompt: {prompt!r} → recommendation: "Heat (1995)"',
            messages=list(self.history) + [Message(role='user', content=prompt)],
            events=[],
            tool_results=[],
            metadata={},
            parsed=None,
            rag_sources=[],
        )

result = fork.continue_from(agent=StubAgent())
print(result.output)
print()
print('result.fork:', result.fork)

# Pretend we have a second run in the store with a different middle
trace2_id = 'demo-run-2'
events2 = [
    AgentEvent('run_started', 'starting', {'user_prompt': 'Recommend a movie.'}),
    AgentEvent('tool_called', 'web_search', {'tool': 'web_search', 'query': 'thriller 1990s'}),
    AgentEvent('tool_completed', 'ok', {'text': 'Heat (1995)'}),
    AgentEvent('run_completed', 'done', {'output': 'Heat (1995) — go watch it.'}),
]
for ev in events2:
    store.append_event(trace2_id, ev)

replayer1 = TraceReplayer.from_store(store, 'demo-run-1')
replayer2 = TraceReplayer.from_store(store, 'demo-run-2')

diff = diff_traces(replayer1, replayer2)
for line in diff.to_lines():
    print(line)