Adaptive System Prompts

from timbal import Agent
from timbal.state import get_run_context

def get_system_prompt() -> str:
    client_chain = get_run_context().current_span().client_chain

    prompts = {
        "whatsapp": "You are a WhatsApp business assistant. Keep responses concise and friendly. Use emojis appropriately.",
        "slack": "You are a Slack bot for internal team communication. Be professional but approachable. Use Slack formatting.",
    }
    
    return prompts.get(client_chain, "You are a helpful assistant.")

def pre_hook():
    span = get_run_context().current_span()
    prompt_text = span.input.get("prompt")
    
    if "whatsapp" in prompt_text:
        span.client_chain = "whatsapp"
    elif "slack" in prompt_text:
        span.client_chain = "slack"
    else:
        span.client_chain = "unknown"

agent = Agent(
    name="ContextAwareAgent",
    model="openai/gpt-4o-mini",
    system_prompt=get_system_prompt,
    pre_hook=pre_hook
)

# Different origins get different system prompts
result1 = await agent(prompt="whatsapp: Hi! I need help").collect()
result2 = await agent(prompt="slack: @bot Can you help?").collect()
result3 = await agent(prompt="Hi! I need help").collect()