Context Enrichers

Context enrichers add additional data to the ToolContext before tools are executed, enabling tools to access user preferences, session state, and other contextual information.

ContextEnricher Interface

from vanna.core.enricher import ContextEnricher
from vanna.core.tool import ToolContext

class ContextEnricher(ABC):
    @abstractmethod
    async def enrich_context(self, context: ToolContext) -> ToolContext:
        """Add data to the tool execution context"""
        return context

Registering Enrichers

from vanna import Agent

agent = Agent(
    llm_service=llm,
    context_enrichers=[
        UserPreferencesEnricher(db),
        TimezoneEnricher(),
        SessionStateEnricher()
    ]
)

Example Enrichers

Example 1: User Preferences

from vanna.core.enricher import ContextEnricher

class UserPreferencesEnricher(ContextEnricher):
    def __init__(self, database):
        self.db = database
    
    async def enrich_context(self, context: ToolContext) -> ToolContext:
        # Fetch user preferences from database
        prefs = await self.db.get_user_preferences(context.user.id)
        
        # Add to context metadata
        context.metadata['preferences'] = {
            'timezone': prefs.get('timezone', 'UTC'),
            'date_format': prefs.get('date_format', 'YYYY-MM-DD'),
            'currency': prefs.get('currency', 'USD'),
            'language': prefs.get('language', 'en')
        }
        
        return context

Tools can then access preferences:

async def execute(self, context: ToolContext, args: Args) -> ToolResult:
    user_timezone = context.metadata.get('preferences', {}).get('timezone', 'UTC')
    # Use timezone in tool logic

Example 2: Temporal Context

from datetime import datetime
import pytz

class TimezoneEnricher(ContextEnricher):
    async def enrich_context(self, context: ToolContext) -> ToolContext:
        # Get user's timezone from preferences or default
        tz_name = context.metadata.get('preferences', {}).get('timezone', 'UTC')
        tz = pytz.timezone(tz_name)
        
        # Add temporal context
        now = datetime.now(tz)
        context.metadata['temporal'] = {
            'current_time': now.isoformat(),
            'current_date': now.date().isoformat(),
            'day_of_week': now.strftime('%A'),
            'timezone': tz_name,
            'is_business_hours': 9 <= now.hour < 17
        }
        
        return context

Example 3: Session State

class SessionStateEnricher(ContextEnricher):
    def __init__(self, session_store):
        self.sessions = session_store
    
    async def enrich_context(self, context: ToolContext) -> ToolContext:
        # Fetch session state
        session = await self.sessions.get(context.conversation_id)
        
        if session:
            context.metadata['session'] = {
                'filters': session.get('active_filters', {}),
                'selected_database': session.get('database'),
                'recent_queries': session.get('query_history', [])[-5:],
                'dashboard_id': session.get('current_dashboard')
            }
        
        return context

Example 4: Database Schema Cache

class SchemaCacheEnricher(ContextEnricher):
    def __init__(self, sql_runner):
        self.sql_runner = sql_runner
        self.schema_cache = {}
    
    async def enrich_context(self, context: ToolContext) -> ToolContext:
        # Cache and provide schema information
        if 'schema' not in self.schema_cache:
            schema = await self.sql_runner.get_schema_info()
            self.schema_cache['schema'] = schema
        
        context.metadata['database_schema'] = self.schema_cache['schema']
        
        return context

Example 5: Feature Flags

class FeatureFlagEnricher(ContextEnricher):
    def __init__(self, feature_flag_service):
        self.flags = feature_flag_service
    
    async def enrich_context(self, context: ToolContext) -> ToolContext:
        # Evaluate feature flags for user
        flags = await self.flags.get_flags_for_user(context.user.id)
        
        context.metadata['features'] = {
            'experimental_charts': flags.get('exp_charts', False),
            'ai_suggestions': flags.get('ai_suggest', True),
            'beta_tools': flags.get('beta_tools', False)
        }
        
        return context

Example 6: Organization Context

class OrganizationEnricher(ContextEnricher):
    def __init__(self, org_service):
        self.org_service = org_service
    
    async def enrich_context(self, context: ToolContext) -> ToolContext:
        # Add organization-level context
        org_id = context.user.metadata.get('organization_id')
        
        if org_id:
            org = await self.org_service.get_organization(org_id)
            
            context.metadata['organization'] = {
                'id': org.id,
                'name': org.name,
                'plan': org.plan_tier,
                'limits': org.usage_limits,
                'settings': org.settings
            }
        
        return context

Accessing Enriched Context in Tools

Tools receive the enriched context:

class MyTool(Tool):
    async def execute(self, context: ToolContext, args: Args) -> ToolResult:
        # Access enriched data
        timezone = context.metadata.get('temporal', {}).get('timezone')
        user_prefs = context.metadata.get('preferences', {})
        session_filters = context.metadata.get('session', {}).get('filters', {})
        
        # Use in tool logic
        result = self.do_something(args, timezone, user_prefs)
        
        return ToolResult(success=True, result_for_llm=result)

Enricher Execution Order

Enrichers run in the order they’re registered:

agent = Agent(
    context_enrichers=[
        PreferencesEnricher(),  # Runs first
        TimezoneEnricher(),     # Can use preferences  
        SessionEnricher()       # Can use timezone
    ]
)

Conditional Enrichment

Only enrich when needed:

class ConditionalEnricher(ContextEnricher):
    async def enrich_context(self, context: ToolContext) -> ToolContext:
        # Only enrich for specific user groups
        if 'premium' in context.user.group_memberships:
            context.metadata['premium_features'] = await self.get_premium_features()
        
        return context

Caching in Enrichers

Cache expensive lookups:

class CachedEnricher(ContextEnricher):
    def __init__(self):
        self.cache = {}
        self.cache_ttl = 300  # 5 minutes
    
    async def enrich_context(self, context: ToolContext) -> ToolContext:
        cache_key = context.user.id
        now = time.time()
        
        # Check cache
        if cache_key in self.cache:
            cached_data, timestamp = self.cache[cache_key]
            if now - timestamp < self.cache_ttl:
                context.metadata['user_data'] = cached_data
                return context
        
        # Fetch and cache
        data = await self.fetch_user_data(context.user.id)
        self.cache[cache_key] = (data, now)
        context.metadata['user_data'] = data
        
        return context

Best Practices

1. Keep enrichers focused - One responsibility per enricher
2. Use caching - Avoid expensive lookups on every request
3. Handle missing data - Provide sensible defaults
4. Document metadata keys - So tools know what’s available
5. Consider performance - Enrichers run on every tool call
6. Use async properly - All enrichers are async
7. Test independently - Unit test each enricher

See Also

• Custom Tools
• Lifecycle Hooks
• Conversation Filters