perf(discovery): Optimize token usage in Socratic question generation

[frankbria]

Summary

The Socratic discovery system includes full conversation history in each AI prompt. With MAX_DISCOVERY_QUESTIONS=20, this can grow to ~10,000 tokens just for history, plus prompt overhead.

Current Behavior

Each call to _generate_next_discovery_question() builds a prompt containing:

• Full conversation history (all Q&A pairs)
• Structured answers by topic
• Uncovered categories
• Socratic guidelines

Worst case scenario:

• 20 turns × ~500 tokens/turn = ~10,000 tokens for history
• Plus ~1,000 tokens for prompt template and metadata
• Total: ~11,000 input tokens per question generation

Cost impact:

• At $3/M input tokens (Claude 3.5 Sonnet): ~$0.033 per call
• Full 20-question discovery: ~$0.66 in input tokens alone

Proposed Optimizations

Option 1: Conversation Summarization

After N turns (e.g., 10), summarize earlier turns into a condensed summary:

if len(conversation_history) > 10:
    summary = self._summarize_turns(conversation_history[:10])
    recent_turns = conversation_history[10:]
    prompt += f"## Earlier Conversation Summary\n{summary}\n\n"
    prompt += "## Recent Conversation\n"
    for turn in recent_turns:
        ...

Option 2: Truncate to Recent K Turns

Keep only the most recent K turns (e.g., 5-7) plus category summaries:

MAX_HISTORY_TURNS = 7
recent_history = conversation_history[-MAX_HISTORY_TURNS:]

Option 3: Token Budget with tiktoken

Count tokens and truncate when approaching a limit:

import tiktoken

def _build_discovery_question_prompt(self, ...):
    enc = tiktoken.encoding_for_model("claude-3-5-sonnet-20241022")
    MAX_PROMPT_TOKENS = 4000
    
    # Add turns until budget reached
    for turn in reversed(conversation_history):
        turn_text = f"Q: {turn['question']}\nA: {turn['answer']}\n"
        if current_tokens + len(enc.encode(turn_text)) > MAX_PROMPT_TOKENS:
            break
        included_turns.insert(0, turn)

Option 4: Semantic Deduplication

Use embeddings to identify and remove redundant information across turns.

Recommendation

Start with Option 2 (truncate to recent K turns) as it's simplest and provides immediate benefit. Consider Option 1 (summarization) for better context preservation if quality degrades.

Acceptance Criteria

• Implement token optimization strategy
• Add token counting/logging to monitor usage
• Ensure question quality doesn't degrade with truncation
• Update documentation with token budget guidance
• Add tests for truncation/summarization logic

Priority

P3 - Nice to have optimization, not blocking for MVP

Related

• PR feat(discovery): implement AI-powered Socratic questioning system #257: Initial Socratic discovery implementation
• Issue feat(discovery): Improve category coverage detection with AI classification #258: Improve category coverage detection with AI classification

[traycerai]

Plan

Observations

The Socratic discovery system currently generates only the first question using AI (in start_discovery()), while subsequent questions come from a predefined framework of 10 questions. However, the chat() method loads full conversation history without truncation, which could grow unbounded. The issue proposes optimizing token usage for either: (1) a future implementation where each discovery question is AI-generated, or (2) the existing conversation history management. The codebase already has file:codeframe/lib/token_counter.py with tiktoken support, providing a foundation for token counting.

Approach

Implement Option 2 (Truncate to Recent K Turns) as the primary strategy, with Option 3 (Token Budget) as a safety mechanism. This hybrid approach provides immediate benefits with minimal complexity while ensuring token limits are never exceeded. The implementation will focus on optimizing conversation history loading in file:codeframe/agents/lead_agent.py and adding configuration for discovery-specific token budgets in file:codeframe/core/config.py. This approach is extensible—if AI-generated questions are implemented later, the same infrastructure will apply.

Implementation Steps

1. Add Configuration Constants

File: file:codeframe/core/config.py

Add discovery-specific configuration to the ContextManagementConfig class:

• max_discovery_questions: int = 20 - Maximum number of discovery questions
• discovery_history_window: int = 7 - Number of recent conversation turns to include
• discovery_max_prompt_tokens: int = 4000 - Token budget for discovery prompts
• enable_conversation_truncation: bool = True - Feature flag for truncation

2. Enhance Token Counter with Conversation Support

File: file:codeframe/lib/token_counter.py

Add new methods to the TokenCounter class:

• count_conversation_tokens(messages: List[Dict[str, str]]) -> int - Count tokens in a conversation history list
• truncate_conversation_to_budget(messages: List[Dict[str, str]], max_tokens: int, preserve_system: bool = True) -> List[Dict[str, str]] - Truncate conversation to fit within token budget, preserving system messages and most recent turns
• get_conversation_summary_stats(messages: List[Dict[str, str]]) -> Dict[str, Any] - Return statistics about conversation token distribution

3. Implement Conversation History Truncation

File: file:codeframe/agents/lead_agent.py

Modify the get_conversation_history() method (around line 160):

• Add optional parameters: max_turns: Optional[int] = None, max_tokens: Optional[int] = None
• Load full conversation from database as currently done
• If max_turns is specified, keep only the most recent N turns (user + assistant pairs)
• If max_tokens is specified, use TokenCounter.truncate_conversation_to_budget() to enforce token limit
• Log truncation events with metrics (original count, truncated count, tokens saved)
• Return truncated conversation history

4. Update Discovery Question Generation

File: file:codeframe/agents/lead_agent.py

Modify _build_discovery_start_prompt() method (around line 460):

• Accept additional parameter: conversation_context: Optional[List[Dict[str, str]]] = None
• If conversation context is provided, include a summarized version in the prompt
• Use token counter to ensure prompt stays within budget
• Add structured sections: "Previous Topics Discussed", "Remaining Categories to Cover"

Create new method _build_discovery_question_prompt():

• Build prompts for AI-generated discovery questions (for future use)
• Include truncated conversation history using get_conversation_history(max_turns=7, max_tokens=3000)
• Include structured answers by topic from answer_capture.get_structured_data()
• Include list of uncovered categories from discovery_framework
• Ensure total prompt stays under configured token budget

5. Add Token Usage Logging

File: file:codeframe/agents/lead_agent.py

Enhance existing token logging in discovery methods:

• In start_discovery() (line 424-429): Add prompt token breakdown (history tokens, template tokens, metadata tokens)
• In process_discovery_answer(): Log token usage if AI question generation is used
• Create helper method _log_discovery_token_usage(prompt_tokens: int, response_tokens: int, history_tokens: int, metadata_tokens: int) for consistent logging format

6. Implement Conversation Summarization (Optional Enhancement)

File: file:codeframe/lib/conversation_summarizer.py (new file)

Create a ConversationSummarizer class for future use:

• summarize_turns(turns: List[Dict[str, str]], max_summary_tokens: int = 500) -> str - Summarize older conversation turns
• Use simple extraction strategy: key topics, decisions made, information gathered
• Integrate with TokenCounter to ensure summaries fit budget
• This enables Option 1 (Conversation Summarization) for future iterations

7. Update Tests

File: file:tests/lib/test_token_counter.py

Add test cases:

• test_count_conversation_tokens() - Verify conversation token counting
• test_truncate_conversation_to_budget() - Test truncation logic with various budgets
• test_truncate_preserves_recent_messages() - Ensure most recent messages are kept
• test_truncate_preserves_system_messages() - Verify system messages are retained

File: file:tests/discovery/test_discovery_questions.py

Add test cases:

• test_discovery_prompt_within_token_budget() - Verify prompts stay under configured limit
• test_conversation_history_truncation_in_discovery() - Test truncation during discovery
• test_token_usage_logging() - Verify token metrics are logged correctly

File: file:tests/agents/test_lead_agent.py

Add test cases:

• test_get_conversation_history_with_max_turns() - Test turn-based truncation
• test_get_conversation_history_with_max_tokens() - Test token-based truncation
• test_discovery_with_long_conversation_history() - Simulate 20-turn discovery and verify token optimization

8. Update Documentation

File: file:docs/discovery-token-optimization.md (new file)

Document the token optimization strategy:

• Configuration options and recommended values
• Token budget calculations and examples
• Monitoring token usage in discovery
• When to adjust discovery_history_window or discovery_max_prompt_tokens
• Future enhancements (summarization, semantic deduplication)

File: file:README.md or relevant user documentation

Add section on discovery token management:

• Explain the 7-turn default window
• How to configure for longer/shorter conversations
• Expected token usage and costs per discovery session

Architecture Diagram

sequenceDiagram
    participant User
    participant LeadAgent
    participant TokenCounter
    participant Database
    participant AnthropicAPI

    User->>LeadAgent: start_discovery() or process_discovery_answer()
    LeadAgent->>Database: Load conversation history
    Database-->>LeadAgent: Full conversation (N turns)
    
    LeadAgent->>TokenCounter: count_conversation_tokens(messages)
    TokenCounter-->>LeadAgent: Total tokens
    
    alt Tokens > Budget
        LeadAgent->>TokenCounter: truncate_conversation_to_budget(messages, max_tokens)
        TokenCounter-->>LeadAgent: Truncated conversation
        LeadAgent->>LeadAgent: Log truncation metrics
    end
    
    LeadAgent->>LeadAgent: Build discovery prompt with truncated history
    LeadAgent->>TokenCounter: count_tokens(prompt)
    TokenCounter-->>LeadAgent: Prompt token count
    
    LeadAgent->>AnthropicAPI: send_message(prompt)
    AnthropicAPI-->>LeadAgent: Response + usage metrics
    
    LeadAgent->>LeadAgent: Log token usage breakdown
    LeadAgent-->>User: Next discovery question

Loading

Configuration Example

# In .codeframe/config.json
{
  "context_management": {
    "enabled": true,
    "max_discovery_questions": 20,
    "discovery_history_window": 7,
    "discovery_max_prompt_tokens": 4000,
    "enable_conversation_truncation": true
  }
}

Token Budget Breakdown

Component	Tokens	Notes
System prompt & guidelines	~500	Socratic instructions, categories
Structured answers by topic	~1000	Extracted data from previous answers
Uncovered categories	~200	List of remaining topics
Conversation history (7 turns)	~1750	7 × ~250 tokens/turn
Safety buffer	~550	For variations in response length
Total Budget	4000	Configurable via discovery_max_prompt_tokens

Monitoring & Metrics

Add logging for:

• discovery.prompt.tokens.history - Tokens used by conversation history
• discovery.prompt.tokens.metadata - Tokens used by structured data
• discovery.prompt.tokens.total - Total prompt tokens
• discovery.truncation.turns_removed - Number of turns truncated
• discovery.truncation.tokens_saved - Tokens saved by truncation

Import In IDE

🤖 Prompt for AI Agents

## Observations

The Socratic discovery system currently generates only the first question using AI (in `start_discovery()`), while subsequent questions come from a predefined framework of 10 questions. However, the `chat()` method loads full conversation history without truncation, which could grow unbounded. The issue proposes optimizing token usage for either: (1) a future implementation where each discovery question is AI-generated, or (2) the existing conversation history management. The codebase already has `file:codeframe/lib/token_counter.py` with tiktoken support, providing a foundation for token counting.

## Approach

Implement **Option 2 (Truncate to Recent K Turns)** as the primary strategy, with **Option 3 (Token Budget)** as a safety mechanism. This hybrid approach provides immediate benefits with minimal complexity while ensuring token limits are never exceeded. The implementation will focus on optimizing conversation history loading in `file:codeframe/agents/lead_agent.py` and adding configuration for discovery-specific token budgets in `file:codeframe/core/config.py`. This approach is extensible—if AI-generated questions are implemented later, the same infrastructure will apply.

## Implementation Steps

### 1. Add Configuration Constants

**File**: `file:codeframe/core/config.py`

Add discovery-specific configuration to the `ContextManagementConfig` class:
- `max_discovery_questions: int = 20` - Maximum number of discovery questions
- `discovery_history_window: int = 7` - Number of recent conversation turns to include
- `discovery_max_prompt_tokens: int = 4000` - Token budget for discovery prompts
- `enable_conversation_truncation: bool = True` - Feature flag for truncation

### 2. Enhance Token Counter with Conversation Support

**File**: `file:codeframe/lib/token_counter.py`

Add new methods to the `TokenCounter` class:
- `count_conversation_tokens(messages: List[Dict[str, str]]) -> int` - Count tokens in a conversation history list
- `truncate_conversation_to_budget(messages: List[Dict[str, str]], max_tokens: int, preserve_system: bool = True) -> List[Dict[str, str]]` - Truncate conversation to fit within token budget, preserving system messages and most recent turns
- `get_conversation_summary_stats(messages: List[Dict[str, str]]) -> Dict[str, Any]` - Return statistics about conversation token distribution

### 3. Implement Conversation History Truncation

**File**: `file:codeframe/agents/lead_agent.py`

Modify the `get_conversation_history()` method (around line 160):
- Add optional parameters: `max_turns: Optional[int] = None`, `max_tokens: Optional[int] = None`
- Load full conversation from database as currently done
- If `max_turns` is specified, keep only the most recent N turns (user + assistant pairs)
- If `max_tokens` is specified, use `TokenCounter.truncate_conversation_to_budget()` to enforce token limit
- Log truncation events with metrics (original count, truncated count, tokens saved)
- Return truncated conversation history

### 4. Update Discovery Question Generation

**File**: `file:codeframe/agents/lead_agent.py`

Modify `_build_discovery_start_prompt()` method (around line 460):
- Accept additional parameter: `conversation_context: Optional[List[Dict[str, str]]] = None`
- If conversation context is provided, include a summarized version in the prompt
- Use token counter to ensure prompt stays within budget
- Add structured sections: "Previous Topics Discussed", "Remaining Categories to Cover"

Create new method `_build_discovery_question_prompt()`:
- Build prompts for AI-generated discovery questions (for future use)
- Include truncated conversation history using `get_conversation_history(max_turns=7, max_tokens=3000)`
- Include structured answers by topic from `answer_capture.get_structured_data()`
- Include list of uncovered categories from `discovery_framework`
- Ensure total prompt stays under configured token budget

### 5. Add Token Usage Logging

**File**: `file:codeframe/agents/lead_agent.py`

Enhance existing token logging in discovery methods:
- In `start_discovery()` (line 424-429): Add prompt token breakdown (history tokens, template tokens, metadata tokens)
- In `process_discovery_answer()`: Log token usage if AI question generation is used
- Create helper method `_log_discovery_token_usage(prompt_tokens: int, response_tokens: int, history_tokens: int, metadata_tokens: int)` for consistent logging format

### 6. Implement Conversation Summarization (Optional Enhancement)

**File**: `file:codeframe/lib/conversation_summarizer.py` (new file)

Create a `ConversationSummarizer` class for future use:
- `summarize_turns(turns: List[Dict[str, str]], max_summary_tokens: int = 500) -> str` - Summarize older conversation turns
- Use simple extraction strategy: key topics, decisions made, information gathered
- Integrate with `TokenCounter` to ensure summaries fit budget
- This enables Option 1 (Conversation Summarization) for future iterations

### 7. Update Tests

**File**: `file:tests/lib/test_token_counter.py`

Add test cases:
- `test_count_conversation_tokens()` - Verify conversation token counting
- `test_truncate_conversation_to_budget()` - Test truncation logic with various budgets
- `test_truncate_preserves_recent_messages()` - Ensure most recent messages are kept
- `test_truncate_preserves_system_messages()` - Verify system messages are retained

**File**: `file:tests/discovery/test_discovery_questions.py`

Add test cases:
- `test_discovery_prompt_within_token_budget()` - Verify prompts stay under configured limit
- `test_conversation_history_truncation_in_discovery()` - Test truncation during discovery
- `test_token_usage_logging()` - Verify token metrics are logged correctly

**File**: `file:tests/agents/test_lead_agent.py`

Add test cases:
- `test_get_conversation_history_with_max_turns()` - Test turn-based truncation
- `test_get_conversation_history_with_max_tokens()` - Test token-based truncation
- `test_discovery_with_long_conversation_history()` - Simulate 20-turn discovery and verify token optimization

### 8. Update Documentation

**File**: `file:docs/discovery-token-optimization.md` (new file)

Document the token optimization strategy:
- Configuration options and recommended values
- Token budget calculations and examples
- Monitoring token usage in discovery
- When to adjust `discovery_history_window` or `discovery_max_prompt_tokens`
- Future enhancements (summarization, semantic deduplication)

**File**: `file:README.md` or relevant user documentation

Add section on discovery token management:
- Explain the 7-turn default window
- How to configure for longer/shorter conversations
- Expected token usage and costs per discovery session

## Architecture Diagram

```mermaid
sequenceDiagram
    participant User
    participant LeadAgent
    participant TokenCounter
    participant Database
    participant AnthropicAPI

    User->>LeadAgent: start_discovery() or process_discovery_answer()
    LeadAgent->>Database: Load conversation history
    Database-->>LeadAgent: Full conversation (N turns)
    
    LeadAgent->>TokenCounter: count_conversation_tokens(messages)
    TokenCounter-->>LeadAgent: Total tokens
    
    alt Tokens > Budget
        LeadAgent->>TokenCounter: truncate_conversation_to_budget(messages, max_tokens)
        TokenCounter-->>LeadAgent: Truncated conversation
        LeadAgent->>LeadAgent: Log truncation metrics
    end
    
    LeadAgent->>LeadAgent: Build discovery prompt with truncated history
    LeadAgent->>TokenCounter: count_tokens(prompt)
    TokenCounter-->>LeadAgent: Prompt token count
    
    LeadAgent->>AnthropicAPI: send_message(prompt)
    AnthropicAPI-->>LeadAgent: Response + usage metrics
    
    LeadAgent->>LeadAgent: Log token usage breakdown
    LeadAgent-->>User: Next discovery question
```

## Configuration Example

```python
# In .codeframe/config.json
{
  "context_management": {
    "enabled": true,
    "max_discovery_questions": 20,
    "discovery_history_window": 7,
    "discovery_max_prompt_tokens": 4000,
    "enable_conversation_truncation": true
  }
}
```

## Token Budget Breakdown

| Component | Tokens | Notes |
|-----------|--------|-------|
| System prompt & guidelines | ~500 | Socratic instructions, categories |
| Structured answers by topic | ~1000 | Extracted data from previous answers |
| Uncovered categories | ~200 | List of remaining topics |
| Conversation history (7 turns) | ~1750 | 7 × ~250 tokens/turn |
| Safety buffer | ~550 | For variations in response length |
| **Total Budget** | **4000** | Configurable via `discovery_max_prompt_tokens` |

## Monitoring & Metrics

Add logging for:
- `discovery.prompt.tokens.history` - Tokens used by conversation history
- `discovery.prompt.tokens.metadata` - Tokens used by structured data
- `discovery.prompt.tokens.total` - Total prompt tokens
- `discovery.truncation.turns_removed` - Number of turns truncated
- `discovery.truncation.tokens_saved` - Tokens saved by truncation

---

Execution Information

Branch: main
Commit: 64973be

---

💡 Tips

Supported Commands (Inside Comments)

• Use @traycerai generate to iterate on the previous version of the implementation plan.

Supported Commands (Inside Description)

• Add @traycerai ignore anywhere in the ticket description to prevent this ticket from being processed.
• Add @traycerai branch:<branch-name> anywhere in the ticket description to specify the target branch for the implementation plan.

Community

• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.