adk-dynamic-prompt.py

# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Standard library imports
import os
import asyncio

# Load environment variables
from dotenv import load_dotenv
load_dotenv()

# ADK Imports
from google.adk.agents import LlmAgent
from google.adk.agents.callback_context import CallbackContext
from google.adk.runners import InMemoryRunner
from google.genai import types
from typing import Optional

# Define the model
GEMINI_2_FLASH = "gemini-2.0-flash"

# --- 1. User Context Functions ---

def extract_user_id_from_message(message: str) -> Optional[str]:
    """Extract user ID from message using common patterns."""
    import re

    patterns = [
        r'user\s+ID\s+is\s+(\d+)',      # "user ID is 456" - prioritize this
        r'ID\s+is\s+(\d+)',             # "ID is 456"
        r'ID:(\w+)',                    # "ID:123"
        r'id:(\w+)',                    # "id:abc"
        r'user_id=(\w+)',               # "user_id=456"
        r'user[_\s-]?id[:\s=]+(\w+)',   # "user-id: 789"
        r'\b(\d{3,})\b',                # Any 3+ digit number (fallback)
    ]

    for pattern in patterns:
        match = re.search(pattern, message, re.IGNORECASE)
        if match:
            return match.group(1)

    return None

def get_user_metadata(user_id: str) -> dict:
    """
    Simulate fetching user metadata from a database.
    In production, this would query your user database/API.
    """
    # Simulated user database
    user_db = {
        "123": {
            "name": "Alice Johnson",
            "role": "Senior Developer",
            "department": "Engineering",
            "location": "San Francisco",
            "timezone": "PST",
            "skill_level": "expert",
            "preferences": {
                "communication_style": "technical",
                "detail_level": "comprehensive"
            },
            "current_projects": ["E-commerce Platform", "API Optimization"],
            "recent_activity": "Working on database performance improvements"
        },
        "456": {
            "name": "Bob Smith",
            "role": "Product Manager",
            "department": "Product",
            "location": "New York",
            "timezone": "EST",
            "skill_level": "intermediate",
            "preferences": {
                "communication_style": "business-focused",
                "detail_level": "summary"
            },
            "current_projects": ["Mobile App Launch", "Q1 Planning"],
            "recent_activity": "Reviewing user feedback and metrics"
        },
        "789": {
            "name": "Charlie Brown",
            "role": "Junior Developer",
            "department": "Engineering",
            "location": "Austin",
            "timezone": "CST",
            "skill_level": "beginner",
            "preferences": {
                "communication_style": "educational",
                "detail_level": "step-by-step"
            },
            "current_projects": ["Learning Codebase", "First Feature"],
            "recent_activity": "Completing onboarding tasks"
        }
    }

    # Return specific user data or generic data
    if user_id in user_db:
        return user_db[user_id]
    else:
        return {
            "name": f"User {user_id}",
            "role": "Team Member",
            "department": "Unknown",
            "skill_level": "intermediate",
            "preferences": {"communication_style": "standard", "detail_level": "balanced"},
            "note": f"Generic profile for user {user_id}"
        }

def create_personalized_instruction(metadata: dict, user_id: str) -> str:
    """Create a personalized instruction based on user metadata."""

    name = metadata.get("name", f"User {user_id}")
    role = metadata.get("role", "Team Member")
    department = metadata.get("department", "Unknown")
    skill_level = metadata.get("skill_level", "intermediate")
    preferences = metadata.get("preferences", {})
    projects = metadata.get("current_projects", [])

    # Adapt tone based on skill level
    if skill_level == "expert":
        tone = "technical and detailed"
        approach = "dive deep into advanced concepts and provide comprehensive technical details"
    elif skill_level == "beginner":
        tone = "friendly and educational"
        approach = "provide clear step-by-step explanations with examples"
    else:
        tone = "professional and helpful"
        approach = "provide balanced explanations with practical examples"

    instruction = f"""
You are an AI assistant helping {name}, who is a {role} in the {department} department.

USER CONTEXT:
- Name: {name}
- Role: {role}
- Department: {department}
- Skill Level: {skill_level}
- Location: {metadata.get("location", "Not specified")}
- Current Projects: {", ".join(projects) if projects else "Not specified"}

COMMUNICATION PREFERENCES:
- Style: {preferences.get("communication_style", "standard")}
- Detail Level: {preferences.get("detail_level", "balanced")}

INSTRUCTIONS:
1. Always address the user as {name}
2. Tailor your responses to their {role} role and {skill_level} skill level
3. Use a {tone} tone
4. {approach}
5. Keep their current projects in mind when providing advice
6. Be helpful and provide actionable guidance

Remember: This user prefers {preferences.get("communication_style", "standard")} communication
with {preferences.get("detail_level", "balanced")} level of detail.
"""

    return instruction.strip()

# --- 2. Dynamic Instruction Provider ---
def dynamic_instruction_provider(ctx) -> str:
    """Provides dynamic instructions based on user input analysis."""
    base_instruction = "You are a helpful AI assistant that adapts to user context."

    # Get the current user message from context
    user_message = ""
    if hasattr(ctx, 'user_content') and ctx.user_content:
        try:
            if ctx.user_content.parts and len(ctx.user_content.parts) > 0:
                user_message = ctx.user_content.parts[0].text
        except Exception as e:
            print(f"[Dynamic Instruction] Error getting user message: {e}")

    print(f"[Dynamic Instruction] Analyzing message: '{user_message}'")

    # Extract user ID from message
    user_id = extract_user_id_from_message(user_message)

    if user_id:
        print(f"[Dynamic Instruction] Extracted User ID: {user_id}")

        # Fetch user metadata
        user_metadata = get_user_metadata(user_id)
        print(f"[Dynamic Instruction] User Metadata: {user_metadata}")

        # Create personalized instruction
        enhanced_instruction = create_personalized_instruction(user_metadata, user_id)
        print(f"[Dynamic Instruction] Generated personalized instruction for: {user_metadata.get('name', f'User {user_id}')}")

        return enhanced_instruction
    else:
        print(f"[Dynamic Instruction] No user ID found, using base instruction")

    return base_instruction

# --- 3. Setup Agent without Callback ---
user_context_agent = LlmAgent(
    name="UserContextAgent",
    model=GEMINI_2_FLASH,
    instruction=dynamic_instruction_provider,  # All logic is now in the instruction provider
    description="An agent that extracts user IDs and personalizes responses"
)

# --- 4. Demo Function ---
async def main():
    app_name = "user_context_demo"
    user_id = "test_user"
    session_id = "user_session"

    # Create runner and session service
    runner = InMemoryRunner(agent=user_context_agent, app_name=app_name)
    session_service = runner.session_service

    # Create session
    await session_service.create_session(
        app_name=app_name,
        user_id=user_id,
        session_id=session_id
    )

    print("🎯 USER CONTEXT INJECTION DEMO")
    print("=" * 50)
    print("Testing automatic user ID extraction and metadata injection")

    # Test different user scenarios
    test_scenarios = [
        {
            "message": "Hello, I'm user ID:123 and I need help with my project.",
            "expected_user": "Alice Johnson (Senior Developer)"
        },
        {
            "message": "Hi! My user ID is 456. Can you help me with planning?",
            "expected_user": "Bob Smith (Product Manager)"
        },
        {
            "message": "I'm user_id=789 and I'm new here. Need guidance.",
            "expected_user": "Charlie Brown (Junior Developer)"
        },
        {
            "message": "This message has no ID - should use default behavior.",
            "expected_user": "No specific user context"
        }
    ]

    for i, scenario in enumerate(test_scenarios, 1):
        print(f"\n--- Test {i}: {scenario['expected_user']} ---")
        print(f"📥 Input: '{scenario['message']}'")

        async for event in runner.run_async(
            user_id=user_id,
            session_id=session_id,
            new_message=types.Content(role="user", parts=[types.Part(text=scenario['message'])])
        ):
            if hasattr(event, 'is_final_response') and event.is_final_response() and event.content:
                response = event.content.parts[0].text.strip()
                print(f"🤖 AI Response: {response}")
                break

    print(f"\n✅ Demo completed! User context injection working successfully.")
    print(f"💡 The agent automatically adapts its responses based on extracted user metadata.")

if __name__ == "__main__":
    asyncio.run(main())