api/trace_implementation.md

Freeplay Traces Implementation Plan

Overview

This document outlines the implementation plan for integrating Freeplay Traces into the existing Huge League Soccer application. Traces provide enhanced observability by grouping LLM interactions within a session, allowing us to track input questions, LLM responses, and associated metadata.

Current Architecture

Existing Freeplay Integration Points

1. Session Management (api/server_gradio.py)

   • AppState.ensure_sessions() - Creates Freeplay session on first user interaction
   • Session ID preserved across state changes

2. Prompt Management (api/workflows/base.py)

   • call_model() - Retrieves persona-specific prompts via get_prompt_by_persona()
   • Uses Freeplay prompt templates (casual_fan_prompt, super_fan_prompt)

3. Session Recording (api/workflows/base.py)

   • should_continue() - Records final LLM responses to Freeplay session
   • Uses record_session() method

Proposed Traces Implementation

AppState Enhancement

The AppState class will be enhanced to store the actual Freeplay session object, allowing direct access to session.create_trace() as shown in the original documentation.

Modified Methods

# api/utils/freeplay_helpers.py - Modified method
def record_session(self, state, end: Optional[float] = time.time(),
                  formatted_prompt: Optional[FormattedPrompt] = None,
                  prompt_vars: Optional[dict] = None,
                  trace_info: Optional[TraceInfo] = None):  # NEW PARAMETER
    # ... existing code ...
    payload = RecordPayload(
        # ... existing fields ...
        trace_info=trace_info  # NEW FIELD
    )

Implementation Flow

1. AppState Enhancement

Location: api/server_gradio.py - AppState class

class AppState(BaseModel):
    # ... existing fields ...
    freeplay_session_id: str = ""
    freeplay_session: Optional[Any] = None  # NEW: Store the actual session object
    
    def ensure_sessions(self):
        if not self.zep_session_id:
            self.zep_session_id = ZepClient() \
                .get_or_create_user(self.email, self.first_name, self.last_name) \
                .create_session() \
                .session_id
        if not self.freeplay_session_id:
            freeplay_client = FreeplayClient()
            self.freeplay_session = freeplay_client.create_session()  # Store session object
            self.freeplay_session_id = self.freeplay_session.session_id

2. User Input & Trace Creation

Location: api/server_gradio.py - submit_helper()

def submit_helper(state, handler, user_query):
    state.ensure_sessions()  # Freeplay session already exists
    
    # Create trace directly on the session object
    trace_info = state.freeplay_session.create_trace(
        input=user_query, # optional metadata not included
    )
    
    # Pass trace_info to workflow
    workflow_bundle, workflow_state = build_workflow_with_state(
        handler=handler,
        zep_session_id=state.zep_session_id,
        freeplay_session_id=state.freeplay_session_id,
        email=state.email,
        first_name=state.first_name,
        last_name=state.last_name,
        persona=state.persona,
        messages=state.history,
        trace_info=trace_info,  # NEW: Pass trace_info to workflow
    )

3. Workflow State Enhancement

Location: api/workflows/base.py - AgentState

class AgentState(TypedDict):
    # ... existing fields ...
    trace_info: Optional[TraceInfo] = None  # NEW

4. Workflow Builder Enhancement

Location: api/workflows/base.py - build_workflow_with_state()

def build_workflow_with_state(handler: AsyncCallbackHandler,
                              zep_session_id: Optional[str] = None,
                              freeplay_session_id: Optional[str] = None,
                              email: Optional[str] = None,
                              first_name: Optional[str] = None,
                              last_name: Optional[str] = None,
                              persona: Optional[str] = None,
                              messages: Optional[List[BaseMessage]] = None,
                              trace_info: Optional[TraceInfo] = None) -> Tuple[WorkflowBundle, AgentState]:  # NEW PARAMETER
    """
    Utility to build workflow and initial state in one step.
    """
    bundle = build_workflow(handler)
    state = {
        "zep_session_id": zep_session_id,
        "freeplay_session_id": freeplay_session_id,
        "email": email,
        "first_name": first_name,
        "last_name": last_name,
        "persona": persona,
        "messages": messages or [],
        "start_time": time.time(),
        "zep_memory": None,
        "trace_info": trace_info,  # NEW: Populate trace_info in AgentState
    }
    return bundle, state

5. LLM Processing with Trace Association

Location: api/workflows/base.py - call_model()

async def call_model(state: AgentState,
                     handler: AsyncCallbackHandler,
                     zep_client: ZepClient,
                     freeplay_client: FreeplayClient) -> dict:
    
    # ... existing prompt and LLM logic ...
    
    # Record trace output with LLM response
    if state.get("trace_info"):
        final_response = response.content if hasattr(response, 'content') else str(response)
        state["trace_info"].record_output(
            project_id=FREEPLAY_PROJECT_ID,
            output=final_response
        )
    
    return {'messages': [response], 'zep_memory': memory}

6. Session Recording with Trace

Location: api/workflows/base.py - should_continue()

async def should_continue(state: AgentState,
                          handler: AsyncCallbackHandler,
                          zep_client: ZepClient,
                          freeplay_client: FreeplayClient) -> str:
    # ... existing logic ...
    
    if 'tool_calls' not in last_message.additional_kwargs:
        # Record session with trace association
        freeplay_client.record_session(state, trace_info=state.get("trace_info"))
        return 'end'
    return 'continue'

Data Flow Diagram

┌─────────────────┐    ┌──────────────────┐    ┌─────────────────┐
│   User Input    │    │  submit_helper   │    │ AppState        │
│                 │    │                  │    │                 │
│ "Tell me about  │───▶│ 1. ensure_sessions│───▶│ freeplay_session│
│  Ryan Martinez" │    │ 2. create_trace() │    │ .create_trace() │
└─────────────────┘    │ 3. pass trace_info│    │                 │
                       └──────────────────┘    └─────────────────┘
                                │                        │
                                ▼                        ▼
                       ┌──────────────────┐    ┌─────────────────┐
                       │ build_workflow_  │    │   Trace Info    │
                       │ with_state()     │    │                 │
                       │                  │    │ - trace_id      │
                       │ - trace_info     │    │ - input         │
                       │ - user_query     │    │ - metadata      │
                       └──────────────────┘    └─────────────────┘
                                │
                                ▼
                       ┌──────────────────┐
                       │   AgentState     │
                       │                  │
                       │ - trace_info     │
                       │ - messages       │
                       │ - persona        │
                       └──────────────────┘
                                │
                                ▼
                       ┌──────────────────┐
                       │   Workflow       │
                       │   Execution      │
                       │                  │
                       │ call_model()     │
                       │ - get_prompt()   │
                       │ - LLM call       │
                       │ - record_output()│
                       └──────────────────┘
                                │
                                ▼
                       ┌──────────────────┐
                       │ should_continue()│
                       │                  │
                       │ record_session() │
                       │ with trace_info  │
                       └──────────────────┘
                                │
                                ▼
                       ┌──────────────────┐
                       │ Freeplay Session │
                       │ Recording        │
                       │                  │
                       │ - Session ID     │
                       │ - Trace Info     │
                       │ - Input/Output   │
                       └──────────────────┘

Trace Lifecycle

1. Creation Phase

• Trigger: User submits question
• Location: submit_helper()
• Action: Create trace with user input and metadata
• Output: TraceInfo object with trace_id

2. Association Phase

• Trigger: Workflow initialization
• Location: build_workflow_with_state()
• Action: Pass trace_info to workflow state
• Output: AgentState with trace context

3. Processing Phase

• Trigger: LLM interaction
• Location: call_model()
• Action: Record LLM output to trace
• Output: Trace with input/output pair

4. Recording Phase

• Trigger: Workflow completion
• Location: should_continue()
• Action: Associate trace with session recording
• Output: Complete session with trace linkage

Benefits

Enhanced Observability

• Input/Output Tracking: Complete visibility of user questions and LLM responses
• Metadata Association: Persona, user info, and interaction context
• Performance Metrics: Response times, tool usage, and evaluation results

Debugging & Analysis

• Trace Isolation: Individual interactions can be analyzed separately
• Session Context: Traces maintain session-level context
• Evaluation Support: Built-in support for feedback and evaluation metrics

Backward Compatibility

• Optional Implementation: Traces are opt-in, existing functionality unchanged
• Graceful Degradation: System works without traces if disabled
• Incremental Rollout: Can be enabled per user or session

Implementation Phases

Phase 1: AppState Enhancement

1. Add freeplay_session property to AppState class
2. Modify ensure_sessions() to store session object
3. Test session object storage and retrieval

Phase 2: Workflow Integration

1. Update AgentState to include trace_info field
2. Modify build_workflow_with_state() to accept trace_info parameter
3. Update call_model() to record trace outputs
4. Update should_continue() to associate traces with sessions

Phase 3: Server Integration

1. Update submit_helper() to create traces using session object
2. Pass trace_info through workflow initialization
3. Test complete trace lifecycle

Phase 4: Session Recording Enhancement

1. Modify record_session() to accept trace_info parameter
2. Update RecordPayload to include trace information
3. Test session recording with trace association

Phase 5: Testing & Validation

1. Unit tests for trace creation and recording
2. Integration tests for workflow trace integration
3. End-to-end tests for complete trace lifecycle
4. Performance impact assessment

Risk Mitigation

High Risk Areas

1. Session State Management

   • Risk: Breaking existing session flow
   • Mitigation: Optional trace parameters, backward compatibility

2. Workflow State Changes

   • Risk: Breaking existing state structure
   • Mitigation: Optional fields with defaults, thorough testing

Medium Risk Areas

1. Performance Impact

   • Risk: Slowing down workflow execution
   • Mitigation: Lazy trace creation, optional features

2. Data Consistency

   • Risk: Trace/recording mismatches
   • Mitigation: Validation and error handling

Success Criteria

1. Functionality: Traces created and recorded for all user interactions
2. Performance: No measurable impact on workflow execution time
3. Compatibility: All existing functionality continues to work
4. Observability: Complete trace data available in Freeplay dashboard
5. Reliability: Graceful handling of trace creation/recording failures

Next Steps

1. User Approval: Review and approve implementation plan
2. Phase 1 Implementation: Begin with core infrastructure
3. Testing Strategy: Validate each phase before proceeding
4. Documentation: Update API documentation with trace features
5. Monitoring: Implement trace-specific logging and metrics