Feature: Chat with Documents - Conversational Interface for Document Q&A

[manavgup]

Feature: Chat with Documents - Conversational Interface for Document Q&A

Overview

Implement a conversational interface that allows users to have multi-turn conversations with their documents, maintaining context across interactions and providing a more natural, interactive experience compared to the current single-turn Q&A system.

Business Value

• Enhanced User Experience: Natural conversational flow vs single queries
• Improved Understanding: Users can refine and clarify through follow-ups
• Higher Engagement: Interactive sessions keep users engaged longer
• Better Insights: Context retention enables deeper document exploration

---

Epic 1: Conversation Session Management

Enable stateful conversation sessions with context persistence and management.

Features

• F1.1: Session lifecycle management (create, maintain, terminate)
• F1.2: Conversation history storage and retrieval
• F1.3: Context window management (token limits)
• F1.4: Session timeout and cleanup

User Stories

US1.1: As a user, I want to start a new chat session with my documents

• Acceptance Criteria:
  • Can initiate a new chat session with a specific collection
  • Receive a unique session ID for tracking
  • Session persists across page refreshes
• Technical Changes:
  • Create ConversationSession model in rag_solution/models/
  • Add session endpoints to rag_solution/router/chat_router.py
  • Implement SessionService in rag_solution/services/
• Metrics:
  • Session creation time < 500ms
  • 99% session persistence reliability

US1.2: As a user, I want my conversation history preserved within a session

• Acceptance Criteria:
  • All Q&A pairs stored chronologically
  • Can scroll back to see previous exchanges
  • History cleared only on explicit action or timeout
• Technical Changes:
  • Add ConversationHistory table with foreign key to sessions
  • Implement history pagination in API
  • Add Redis for session caching (optional optimization)
• Metrics:
  • History retrieval < 200ms for last 50 messages
  • Zero data loss during active sessions

---

Epic 2: Context-Aware Response Generation

Implement context management for maintaining conversational coherence.

Features

• F2.1: Context accumulation across turns
• F2.2: Pronoun resolution and reference handling
• F2.3: Smart context pruning (relevance-based)
• F2.4: Integration with existing CoT capabilities

User Stories

US2.1: As a user, I want the system to understand context from previous questions

• Acceptance Criteria:
  • "Tell me more" understands previous topic
  • "What about X?" relates X to previous context
  • Pronouns (it, they, this) correctly resolved
• Technical Changes:
  • Create ContextManager class in rag_solution/services/
  • Implement sliding window context in ChatService
  • Add context preprocessing to prompt templates
• Metrics:
  • 90% accuracy in pronoun resolution (manual evaluation)
  • Context relevance score > 0.8

US2.2: As a user, I want to ask follow-up questions naturally

• Acceptance Criteria:
  • Follow-ups don't need to repeat context
  • System maintains topic continuity
  • Can handle topic shifts gracefully
• Technical Changes:
  • Implement FollowUpDetector in services
  • Add context injection to search queries
  • Modify SearchService to accept conversation context
• Metrics:
  • 85% user satisfaction with follow-up handling
  • Reduced query length by 40% for follow-ups

---

Epic 3: Conversational UI/UX

Build intuitive chat interface and API endpoints.

Features

• F3.1: Chat-style UI component
• F3.2: Real-time streaming responses
• F3.3: Typing indicators and status feedback
• F3.4: Quick action buttons (suggested questions)

User Stories

US3.1: As a user, I want a familiar chat interface like ChatGPT

• Acceptance Criteria:
  • Message bubbles with user/assistant distinction
  • Auto-scroll to latest message
  • Responsive design for mobile/desktop
• Technical Changes:
  • Create React chat components in webui/src/components/Chat/
  • Add WebSocket support for real-time updates
  • Implement ChatController in backend
• Metrics:
  • UI response time < 100ms
  • 95% positive UI feedback

US3.2: As a user, I want suggested follow-up questions

• Acceptance Criteria:
  • 3-5 relevant suggestions after each response
  • Click to send suggestion as next question
  • Suggestions based on current context
• Technical Changes:
  • Add QuestionSuggestionService
  • Implement suggestion generation using LLM
  • Add suggestion caching for performance
• Metrics:
  • 60% click-through rate on suggestions
  • Suggestion generation < 1s

---

Epic 4: Advanced Conversation Features

Enhance conversational capabilities with advanced features.

Features

• F4.1: Multi-collection conversations
• F4.2: Conversation summarization
• F4.3: Export conversation history
• F4.4: Conversation branching (explore different paths)

User Stories

US4.1: As a user, I want to save and share conversations

• Acceptance Criteria:
  • Export as PDF/Markdown/JSON
  • Shareable links with read-only access
  • Bookmark important exchanges
• Technical Changes:
  • Add export formatters in rag_solution/utils/
  • Implement sharing service with access controls
  • Add bookmarking to conversation model
• Metrics:
  • Export generation < 3s
  • 99.9% export fidelity

US4.2: As a user, I want to compare information across documents

• Acceptance Criteria:
  • Reference multiple documents in one conversation
  • System indicates source for each fact
  • Can switch document context mid-conversation
• Technical Changes:
  • Extend SearchService for multi-collection queries
  • Add source tracking to response generation
  • Implement collection context switching
• Metrics:
  • Multi-doc query latency < 2x single-doc
  • 95% source attribution accuracy

---

Implementation Plan

Phase 1 (Sprint 1-2): Foundation

• Design database schema for conversations
• Implement basic session management
• Create minimal chat API endpoints
• Build proof-of-concept UI

Phase 2 (Sprint 3-4): Core Functionality

• Implement context management
• Add conversation history
• Integrate with existing search pipeline
• Enhance UI with real-time features

Phase 3 (Sprint 5-6): Enhancement

• Add follow-up suggestions
• Implement conversation export
• Add multi-collection support
• Performance optimization

Phase 4 (Sprint 7): Polish & Scale

• Add comprehensive testing
• Implement caching strategies
• Add monitoring and analytics
• Documentation and examples

---

Success Metrics

Technical Metrics

• Latency: P95 response time < 2s
• Throughput: Support 100 concurrent conversations
• Reliability: 99.9% uptime for chat service
• Scalability: Linear scaling with session count

User Metrics

• Engagement: 3x increase in average session duration
• Satisfaction: 85% positive feedback on chat experience
• Adoption: 60% of users prefer chat over single Q&A
• Retention: 40% higher return rate for chat users

Business Metrics

• Efficiency: 50% reduction in repeated queries
• Insights: 2x more questions per user session
• Value: 30% increase in document utilization
• ROI: Positive ROI within 6 months

---

Technical Architecture

┌─────────────┐     ┌─────────────┐     ┌─────────────┐
│   Chat UI   │────▶│  Chat API   │────▶│   Session   │
│  (React)    │     │  (FastAPI)  │     │   Service   │
└─────────────┘     └─────────────┘     └─────────────┘
                            │                    │
                            ▼                    ▼
                    ┌─────────────┐     ┌─────────────┐
                    │   Context   │     │  Database   │
                    │   Manager   │     │ (Sessions)  │
                    └─────────────┘     └─────────────┘
                            │
                            ▼
                    ┌─────────────┐
                    │   Search    │──── Existing RAG
                    │   Service   │     Pipeline
                    └─────────────┘

---

Risks & Mitigations

Risk	Impact	Mitigation
Context window limits	High	Implement smart pruning algorithm
Response latency	Medium	Add caching and async processing
Session state management	High	Use Redis for fast session access
UI complexity	Medium	Iterative design with user feedback
LLM costs	High	Implement usage quotas and optimization

---

Dependencies

• Existing Search Service with CoT
• LLM Provider Service
• User Authentication System
• Document Collection Service

---

Related Issues

• Improve Pipeline Association Architecture for Better UX and Flexibility #222 - Simplified Pipeline Resolution
• Chain of Thought Implementation (completed)

---

Acceptance Criteria for Epic

• All user stories implemented and tested
• Performance metrics met
• Documentation complete
• Security review passed
• User acceptance testing completed

---

Labels

• enhancement
• epic
• backend
• frontend
• high-priority

Assignees

@manavgup

Milestone

Q1 2025 - Conversational AI Features

[manavgup]

CLI Support Epic for Issue #229: Chat with Documents

Overview

This epic adds comprehensive Command Line Interface (CLI) support to the Chat with Documents feature, enabling users to interact with conversational sessions directly from the terminal. This complements the existing web-based interface and provides power users with efficient command-line access to all chat functionality.

Epic Objectives

Primary Goals

• Terminal-Based Chat Interface: Enable users to create, manage, and interact with conversation sessions via CLI commands
• Seamless Integration: Leverage existing CLI infrastructure (backend/rag_solution/cli) and entry point (rag-cli)
• Feature Parity: Provide CLI access to all Chat with Documents functionality available through the API
• Enhanced Productivity: Support scripting, automation, and batch operations for power users

Success Criteria

• All chat operations available via CLI commands
• Consistent user experience with existing CLI patterns
• Comprehensive test coverage for CLI functionality
• Complete documentation and examples
• Performance equivalent to API operations

Scope and Features

Core CLI Commands

Session Management

# Create a new conversation session
rag-cli chat create-session <user_id> <collection_id> <session_name> [options]

# List user's conversation sessions
rag-cli chat list-sessions <user_id> [--status <status>] [--limit <n>] [--offset <n>]

# Show session details
rag-cli chat show-session <session_id> <user_id>

# Update session settings
rag-cli chat update-session <session_id> <user_id> [--session-name <name>] [--context-window-size <size>] [--max-messages <count>]

# Delete session
rag-cli chat delete-session <session_id> <user_id> [--force]

# Archive session
rag-cli chat archive-session <session_id> <user_id>

# Restore archived session
rag-cli chat restore-session <session_id> <user_id>

Message Management

# Add message to session
rag-cli chat add-message <session_id> <content> [--role <role>] [--message-type <type>] [--metadata <json>]

# List session messages
rag-cli chat list-messages <session_id> <user_id> [--limit <n>] [--offset <n>]

Analytics and Export

# Get session statistics
rag-cli chat session-stats <session_id> <user_id>

# Export session data
rag-cli chat export-session <session_id> <user_id> [--format <json|markdown|pdf>] [--output-file <path>]

# Search sessions
rag-cli chat search-sessions <user_id> <query> [--limit <n>]

# Cleanup expired sessions
rag-cli chat cleanup

Advanced Features

Interactive Chat Mode

# Start interactive chat session
rag-cli chat interactive <session_id> <user_id>

• Real-time conversation interface
• Context-aware suggestions
• Multi-turn conversation support
• Auto-completion for commands

Batch Operations

# Batch message processing
rag-cli chat batch-messages <session_id> <messages_file>

# Bulk session operations
rag-cli chat bulk-archive <user_id> <session_ids_file>

Integration Features

# Pipe content to chat
echo "What is machine learning?" | rag-cli chat add-message <session_id>

# Export to various formats
rag-cli chat export-session <session_id> <user_id> --format markdown | pandoc -f markdown -t html

# Integration with search
rag-cli search query <collection_id> "AI concepts" | rag-cli chat add-message <session_id>

Technical Implementation

Architecture Integration

CLI Command Structure

• Location: backend/rag_solution/cli/commands/chat.py
• Base Class: Extends BaseCommand for consistency
• API Integration: Uses RAGAPIClient for all operations
• Error Handling: Comprehensive error handling and user feedback

Entry Point Integration

• Primary Entry: rag-cli (existing entry point)
• Command Group: rag-cli chat <subcommand>
• Help System: Integrated help and documentation
• Configuration: Uses existing CLI configuration system

Output Formats

• Table Format: Human-readable tabular output (default)
• JSON Format: Machine-readable JSON output (--output json)
• YAML Format: YAML output for configuration files
• CSV Format: Comma-separated values for data analysis

Data Flow

graph TD
    A[User CLI Command] --> B[rag-cli Entry Point]
    B --> C[ChatCommands Class]
    C --> D[RAGAPIClient]
    D --> E[Chat API Endpoints]
    E --> F[ConversationService]
    F --> G[Database]
    
    C --> H[Output Formatter]
    H --> I[Terminal Display]
    
    J[Configuration] --> C
    K[Authentication] --> D

Loading

Error Handling Strategy

User-Friendly Error Messages

• Clear, actionable error descriptions
• Suggested solutions for common issues
• Proper exit codes for scripting

Error Categories

• Authentication Errors: Login required, token expired
• Validation Errors: Invalid parameters, missing required fields
• API Errors: Service unavailable, rate limiting
• Network Errors: Connection timeouts, DNS resolution
• Permission Errors: Access denied, insufficient privileges

Testing Strategy

Test Coverage

• Unit Tests: Individual command testing with mocks
• Integration Tests: End-to-end CLI workflow testing
• E2E Tests: Complete user journey testing
• Performance Tests: Load testing for batch operations

Test Categories

Unit Tests (test_chat_cli_tdd.py)

• Command method testing
• Parameter validation
• Error handling
• Output formatting

Integration Tests

• API client integration
• Authentication flow
• Configuration management
• Error propagation

E2E Tests

• Complete user workflows
• Multi-command scenarios
• File I/O operations
• Cross-platform compatibility

Documentation Requirements

User Documentation

• CLI Reference Guide: Complete command reference
• Tutorials: Step-by-step usage examples
• Best Practices: Recommended workflows and patterns
• Troubleshooting: Common issues and solutions

Developer Documentation

• API Integration: How CLI commands map to API endpoints
• Extension Guide: Adding new chat commands
• Testing Guide: Running and writing tests
• Architecture: CLI system design and patterns

Implementation Phases

Phase 1: Core Commands (Week 1-2)

• Session management commands (create, list, show, update, delete)
• Basic message commands (add, list)
• Error handling and validation
• Unit tests for core functionality

Phase 2: Advanced Features (Week 3-4)

• Archive/restore functionality
• Statistics and analytics commands
• Export functionality
• Search and cleanup commands
• Integration tests

Phase 3: Interactive Mode (Week 5-6)

• Interactive chat interface
• Real-time conversation support
• Auto-completion and suggestions
• E2E tests

Phase 4: Polish and Documentation (Week 7-8)

• Performance optimization
• Comprehensive documentation
• User tutorials and examples
• Final testing and validation

Dependencies

Internal Dependencies

• Chat with Documents API implementation (Issue Feature: Chat with Documents - Conversational Interface for Document Q&A #229)
• Existing CLI infrastructure (backend/rag_solution/cli)
• Authentication system
• Configuration management

External Dependencies

• Python argparse library
• requests library for API calls
• tabulate for table formatting
• rich for enhanced terminal output (optional)

Acceptance Criteria

Functional Requirements

• All chat operations accessible via CLI
• Consistent command syntax and behavior
• Proper error handling and user feedback
• Support for all output formats
• Integration with existing CLI infrastructure

Non-Functional Requirements

• Response time < 2 seconds for most operations
• Memory usage < 100MB for CLI process
• 100% test coverage for core functionality
• Cross-platform compatibility (Linux, macOS, Windows)
• Comprehensive documentation

Quality Requirements

• Code follows project style guidelines
• All tests pass consistently
• No memory leaks or resource issues
• Proper logging and debugging support
• Security best practices followed

Risks and Mitigation

Technical Risks

• API Changes: CLI tightly coupled to API
  • Mitigation: Use versioned API endpoints, implement graceful degradation
• Performance: CLI operations slower than expected
  • Mitigation: Implement caching, optimize API calls, add progress indicators
• Complexity: CLI becoming too complex for users
  • Mitigation: Focus on common use cases, provide clear help and examples

Integration Risks

• Breaking Changes: CLI changes affecting existing functionality
  • Mitigation: Comprehensive testing, backward compatibility
• Authentication: Complex auth flow for CLI users
  • Mitigation: Simplify auth process, provide clear setup instructions

Success Metrics

Usage Metrics

• Number of CLI commands executed per day
• User adoption rate for CLI vs web interface
• Most frequently used commands
• Error rates and resolution times

Performance Metrics

• Average command execution time
• Memory usage patterns
• API call efficiency
• User satisfaction scores

Future Enhancements

Potential Additions

• Plugin System: Allow third-party CLI extensions
• Configuration Profiles: Multiple configuration profiles for different environments
• Advanced Scripting: Support for complex automation workflows
• Real-time Updates: Live updates for active sessions
• Integration APIs: Direct integration with other tools and services

---

Implementation Notes

This epic should be implemented after the core Chat with Documents API is complete (Issue #229). The CLI implementation will depend on the API endpoints and data models defined in the main feature.

The existing CLI infrastructure in backend/rag_solution/cli provides a solid foundation, and the rag-cli entry point should be used as the primary interface for all chat-related commands.

All implementation should follow the existing patterns and conventions established in the current CLI codebase, ensuring consistency and maintainability.

[manavgup]

Seamless Integration Epic for Issue #229: Search + CoT + Conversation

Overview

This epic defines the seamless integration between three core services: Search Service, Chain of Thought (CoT) Service, and Conversation Service. The goal is to create a unified, context-aware system where conversation history enhances search capabilities and CoT reasoning, while maintaining the existing functionality of each service.

Epic Objectives

Primary Goals

• Unified Context Flow: Seamless context propagation from conversation → search → CoT
• Enhanced Reasoning: CoT that considers conversation history and context
• Preserved Functionality: All existing search and CoT capabilities remain intact
• Performance Optimization: Efficient context management without duplication
• User Experience: Natural conversational flow with enhanced reasoning

Success Criteria

• Conversation context enhances search quality by 30%+
• CoT reasoning considers conversation history
• Zero breaking changes to existing search/CoT APIs
• Context propagation adds <100ms latency
• All existing tests continue to pass

Architecture Integration

Service Layer Integration

ConversationService
    ↓
ContextManagerService (build context from history)
    ↓
SearchService (enhanced with conversation context)
    ↓
PipelineService (document retrieval)
    ↓
ChainOfThoughtService (reasoning with conversation context)
    ↓
AnswerSynthesizer (final answer with context)
    ↓
ConversationService (store response with metadata)
    ↓
QuestionSuggestionService (generate follow-ups)

Integration Points

1. Context Propagation

Conversation → Search

# Enhanced SearchInput with conversation context
search_input = SearchInput(
    question=enhanced_question,
    collection_id=session.collection_id,
    user_id=session.user_id,
    config_metadata={
        "conversation_context": context.context_window,
        "session_id": str(session.id),
        "message_history": [msg.content for msg in recent_messages],
        "conversation_entities": context.extracted_entities,
        "cot_enabled": True,
        "show_cot_steps": False  # Hide in conversation UI
    }
)

Search → CoT

# Enhanced CoT input with conversation context
cot_input = ChainOfThoughtInput(
    question=search_input.question,
    collection_id=search_input.collection_id,
    user_id=search_input.user_id,
    cot_config=search_input.config_metadata,
    context_metadata={
        "conversation_context": search_input.config_metadata.get("conversation_context"),
        "message_history": search_input.config_metadata.get("message_history", []),
        "conversation_entities": search_input.config_metadata.get("conversation_entities", []),
        "previous_answers": [msg.content for msg in recent_messages if msg.role == "assistant"]
    }
)

2. Context Enhancement Strategies

Question Enhancement

def enhance_question_with_conversation_context(
    self, 
    question: str, 
    conversation_context: str, 
    message_history: list[str]
) -> str:
    """Enhance question with conversation context for better search."""
    
    # Extract entities from conversation
    entities = self.extract_entities_from_context(conversation_context)
    
    # Build enhanced question
    if entities:
        entity_context = f" (in the context of {', '.join(entities)})"
        enhanced_question = f"{question}{entity_context}"
    else:
        enhanced_question = question
    
    # Add conversation context if question is ambiguous
    if self.is_ambiguous_question(question):
        recent_context = " ".join(message_history[-3:])  # Last 3 messages
        enhanced_question = f"{question} (referring to: {recent_context})"
    
    return enhanced_question

Service Integration Patterns

1. Lazy Initialization Pattern

class ConversationService:
    def __init__(self, db: Session, settings: Settings):
        self.db = db
        self.settings = settings
        # Lazy initialization for seamless integration
        self._search_service: SearchService | None = None
        self._context_manager_service: ContextManagerService | None = None
        self._question_suggestion_service: QuestionSuggestionService | None = None
    
    @property
    def search_service(self) -> SearchService:
        if self._search_service is None:
            self._search_service = SearchService(self.db, self.settings)
        return self._search_service

2. Context-Aware Search Enhancement

class SearchService:
    async def search(self, search_input: SearchInput) -> SearchOutput:
        """Enhanced search with conversation context."""
        
        # Check for conversation context
        conversation_context = search_input.config_metadata.get("conversation_context")
        if conversation_context:
            # Enhance question with conversation context
            search_input.question = self._enhance_question_with_conversation(
                search_input.question, 
                conversation_context,
                search_input.config_metadata.get("message_history", [])
            )
        
        # Continue with existing search logic
        # ... existing implementation

Testing Strategy

1. Integration Test Categories

Context Flow Tests

• Test context propagation from conversation → search → CoT
• Test context enhancement and pruning
• Test context caching and performance

Service Integration Tests

• Test ConversationService → SearchService integration
• Test SearchService → CoT integration with conversation context
• Test error handling across service boundaries

End-to-End Workflow Tests

• Test complete conversation flow with CoT
• Test multi-turn reasoning with context
• Test performance under load

Implementation Phases

Phase 1: Basic Integration (Week 1-2)

• Context propagation from conversation to search
• Basic context enhancement in SearchService
• CoT integration with conversation context
• Integration tests for context flow

Phase 2: Enhanced Integration (Week 3-4)

• Advanced context pruning and optimization
• Conversation-aware CoT reasoning
• Performance optimization and caching
• Comprehensive integration tests

Phase 3: Advanced Features (Week 5-6)

• Multi-turn reasoning chains
• Context-aware question suggestions
• Advanced context management strategies
• End-to-end workflow tests

Success Metrics

Functional Metrics

• Context enhancement improves search relevance by 30%+
• CoT reasoning considers conversation history
• Zero breaking changes to existing APIs
• All existing tests continue to pass

Performance Metrics

• Context propagation adds <100ms latency
• Context caching reduces build time by 50%+
• Memory usage stays within 10% increase
• Response time remains under 5 seconds

---

This epic should be implemented after the core Chat with Documents feature (Issue #229) is complete. The seamless integration will enhance the existing functionality without breaking changes, providing a natural conversational experience with enhanced reasoning capabilities.

[manavgup]

Token Tracking Implementation Plan - Sub-Issue

This comment outlines a comprehensive Token Tracking Implementation Plan for the Chat with Documents feature (Issue #229). This addresses the token limit handling requirements mentioned in Epic 1: Context window management (F1.3).

Overview

Instead of predicting token usage through complex tokenizer mapping, this implementation tracks actual token consumption from LLM API responses to provide:

• ✅ 100% Accurate Token Counts - Uses real data from API responses
• ✅ Real-time Warnings - Proactive alerts when approaching context limits
• ✅ Provider Agnostic - Works with OpenAI, Anthropic, IBM WatsonX, etc.
• ✅ Cost Tracking - Monitor actual API usage and costs
• ✅ Simple Implementation - No complex tokenizer management needed

Key Components

1. LLMUsage Data Structure - Standardized token usage tracking across services
2. Enhanced LLM Providers - All providers return usage data with responses
3. Token Warning System - Real-time notifications at 70%, 85%, and 95% thresholds
4. Service Integration - Updates to SearchService, ConversationService, and ChainOfThoughtService
5. Database Schema Updates - Add context_window limits to llm_models table

Implementation Approach

Core Token Tracking

@dataclass
class LLMUsage:
    """Actual token usage from LLM API response."""
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    model_name: str
    service_type: ServiceType  # search, conversation, chain_of_thought
    timestamp: datetime
    user_id: Optional[str] = None
    session_id: Optional[str] = None

Provider Integration

All LLM providers updated with new method:

async def generate_text_with_usage(
    self, prompt: str, service_type: ServiceType, **kwargs
) -> Tuple[str, LLMUsage]:
    # Generate response and extract actual usage from API response
    # OpenAI: response.usage.total_tokens
    # Anthropic: response.usage.input_tokens + output_tokens  
    # IBM: result.input_token_count + generated_token_count

Warning System

class TokenWarningService:
    async def check_usage_warning(self, usage: LLMUsage) -> Optional[TokenWarning]:
        # 70-85%: Info warnings
        # 85-95%: Warning severity
        # 95%+: Critical warnings with suggested actions

Service Updates

Search Service

• Tracks tokens for regular and Chain-of-Thought searches
• Includes token usage and warnings in SearchOutput metadata
• Aggregates multi-step CoT token consumption

Conversation Service

• Monitors session-level token accumulation
• Provides conversation-length warnings
• Returns token warnings with each message response
• New endpoint: /api/chat/sessions/{id}/token-stats

Chain of Thought Service

• Tracks tokens for each reasoning step (classification, decomposition, generation, synthesis)
• Aggregates total usage across all steps
• Includes per-step token breakdown in metadata

Database Changes

LLM Models Table Update

ALTER TABLE llm_models ADD COLUMN context_window INTEGER DEFAULT 4096;
ALTER TABLE llm_models ADD COLUMN max_output_tokens INTEGER DEFAULT 2048;

-- Update with known model limits:
-- GPT-3.5-turbo: 4K context, GPT-4-turbo: 128K context
-- Claude-3: 200K context, Granite: 8K context (conservative)

Optional Usage Tracking Table

CREATE TABLE token_usage (
    id UUID PRIMARY KEY,
    user_id UUID REFERENCES users(id),
    session_id UUID REFERENCES conversation_sessions(id),
    model_id VARCHAR(255),
    service_type VARCHAR(50), -- search, conversation, cot
    prompt_tokens INTEGER,
    completion_tokens INTEGER, 
    total_tokens INTEGER,
    created_at TIMESTAMP WITH TIME ZONE
);

API Response Updates

Chat Process Endpoint

POST /api/chat/sessions/{id}/process
{
  "message": { /* ConversationMessageOutput */ },
  "status": "success",
  "token_warning": {
    "type": "approaching_limit",
    "message": "Context window is 87% full. Approaching limit.",
    "percentage_used": 87.3,
    "severity": "warning", 
    "current_tokens": 3574,
    "limit_tokens": 4096,
    "suggested_action": "consider_new_session"
  }
}

Search Metadata

{
  "answer": "...",
  "metadata": {
    "token_usage": {
      "prompt_tokens": 1200,
      "completion_tokens": 300,
      "total_tokens": 1500,
      "model_name": "gpt-3.5-turbo"
    },
    "token_warning": { /* TokenWarning if applicable */ },
    "cot_token_breakdown": [
      {"step": "classification", "total_tokens": 150},
      {"step": "generation", "total_tokens": 800}
    ]
  }
}

4-Week Implementation Timeline

Week 1: Foundation

• Create LLMUsage and TokenWarning schemas
• Update LLM provider base class and implementations
• Create TokenWarningService
• Database migration for context windows

Week 2: Service Integration

• Update SearchService with token tracking
• Update ConversationService with warnings
• Update ChainOfThoughtService with multi-step tracking

Week 3: API and Testing

• Update API endpoints with token data
• Add token statistics endpoint
• Comprehensive test suite (unit + integration)

Week 4: Polish and Deploy

• E2E testing and performance optimization
• Documentation updates
• Production deployment

Benefits for Issue #229

This directly addresses Epic 1 requirements:

• F1.3 Context window management: Real-time token monitoring with accurate limits
• Session lifecycle management: Token warnings help determine when to start new sessions
• User Experience: Clear warnings and suggestions prevent unexpected truncation
• Conversation history: Token statistics show usage patterns over time

Integration with Existing Features

• Chain of Thought: Multi-step token aggregation and per-step tracking
• Search Service: Unified token tracking for both regular and CoT searches
• Conversation Sessions: Session-level usage monitoring and warnings
• LLM Providers: Provider-agnostic implementation using actual API data

Files to be Created/Modified

New Files:

• rag_solution/schemas/llm_usage_schema.py
• rag_solution/services/token_warning_service.py
• tests/unit/test_token_tracking.py
• tests/integration/test_token_integration.py
• Database migration files

Modified Files:

• rag_solution/generation/providers/base.py
• rag_solution/generation/providers/openai.py
• rag_solution/generation/providers/anthropic.py
• rag_solution/generation/providers/watsonx.py
• rag_solution/services/search_service.py
• rag_solution/services/conversation_service.py
• rag_solution/services/chain_of_thought_service.py
• rag_solution/router/chat_router.py
• Multiple schema files for token-related fields

This implementation provides a robust foundation for token management in the conversational interface while maintaining clean separation of concerns and provider independence.

Full Implementation Plan Available: /backend/TOKEN_TRACKING_IMPLEMENTATION_PLAN.md (1,700+ lines)

Would this be better implemented as a separate sub-issue, or integrated directly into the Epic 1 implementation?