gemini_example.py

import os
import logging
from dotenv import load_dotenv
import ldclient
from ldclient import Context
from ldclient.config import Config
from ldai import LDAIClient, LDMessage
from ldai.tracker import TokenUsage
from ldobserve import ObservabilityConfig, ObservabilityPlugin
from google import genai
from google.genai import types
from typing import List, Optional, Tuple

load_dotenv()

logging.basicConfig()
logging.getLogger('ldclient').setLevel(logging.WARNING)

# Set sdk_key to your LaunchDarkly SDK key.
sdk_key = os.getenv('LAUNCHDARKLY_SDK_KEY')

# Set config_key to the AI Config key you want to evaluate.
ai_config_key = os.getenv('LAUNCHDARKLY_COMPLETION_KEY', 'sample-completion')

# Set Google API key
google_api_key = os.getenv('GOOGLE_API_KEY')

def map_to_google_ai_messages(
    input_messages: List[LDMessage]
) -> Tuple[Optional[str], List[types.Content]]:
    messages: List[types.Content] = []

    system_messages: List[str] = []

    for message in input_messages:
        if message.role == 'system':
            system_messages.append(message.content)
            continue
        elif message.role == 'assistant':
            role = "model"
            parts = [types.Part(text=message.content)]
        elif message.role == 'user':
            role = "user"
            parts = [types.Part(text=message.content)]
        else:
            # Skip other message types
            continue

        messages.append(types.Content(role=role, parts=parts))

    # Concatenate system messages with spaces
    system_instruction = " ".join(system_messages) if system_messages else None
    return system_instruction, messages

def track_genai_metrics(tracker, func):
    """
    Track GenAi-specific operations.

    This function will track the duration of the operation, the token
    usage, and the success or error status.

    If the provided function throws, then this method will also throw.

    In the case the provided function throws, this function will record the
    duration and an error.

    A failed operation will not have any token usage data.

    :param tracker: The LaunchDarkly tracker instance.
    :param func: Function to track.
    :return: Result of the tracked function.
    """
    try:
        result = tracker.track_duration_of(func)
        tracker.track_success()
        if hasattr(result, "usage_metadata") and result.usage_metadata:
            # Extract token usage from LangChain response
            usage_data = result.usage_metadata
            token_usage = TokenUsage(
                input=getattr(usage_data, 'prompt_token_count', 0),
                output=getattr(usage_data, "candidates_token_count", 0),
                total=getattr(usage_data, 'total_token_count', 0)
            )
            tracker.track_tokens(token_usage)
    except Exception:
        tracker.track_error()
        raise

    return result

def main():
    if not sdk_key:
        print("*** Please set the LAUNCHDARKLY_SDK_KEY env first")
        exit()
    if not ai_config_key:
        print("*** Please set the LAUNCHDARKLY_COMPLETION_KEY env first")
        exit()
    if not google_api_key:
        print("*** Please set the GOOGLE_API_KEY env first")
        exit()

    ldclient.set_config(Config(sdk_key, plugins=[
        ObservabilityPlugin(ObservabilityConfig(
            service_name='hello-python-ai-gemini',
        ))
    ]))
    aiclient = LDAIClient(ldclient.get())

    if not ldclient.get().is_initialized():
        print("*** SDK failed to initialize. Please check your internet connection and SDK credential for any typo.")
        exit()

    print("*** SDK successfully initialized")

    # Set up the evaluation context. This context should appear on your
    # LaunchDarkly contexts dashboard soon after you run the demo.
    context = (
        Context
        .builder('example-user-key')
        .kind('user')
        .name('Sandy')
        .build()
    )

    # Pass a default for improved resiliency when the AI config is unavailable
    # or LaunchDarkly is unreachable; omit for a disabled default.
    # Example:
    #   default = AIConfig(
    #       enabled=True,
    #       model=ModelConfig(name='gemini-pro'),
    #       provider=ProviderConfig(name='google'),
    #       messages=[LDMessage(role='system', content='You are a helpful assistant.')],
    #   )
    #   config_value = aiclient.completion_config(ai_config_key, context, default, {'myUserVariable': "Testing Variable"})
    config_value = aiclient.completion_config(
        ai_config_key,
        context,
        variables={'myUserVariable': "Testing Variable"}
    )

    if not config_value.enabled:
        print(f"AI config '{ai_config_key}' is disabled. Verify the config key exists in your LaunchDarkly project and is not targeting a disabled variation.")
        return

    tracker = config_value.create_tracker()

    # Configure Google Generative AI
    client = genai.Client(
        api_key=google_api_key,
    )

    # Convert LaunchDarkly messages to Google AI format using the helper function
    system_instruction, messages = map_to_google_ai_messages(config_value.messages or [])

    SAMPLE_QUESTION = "What can you help me with?"
    user_message = types.Content(role="user", parts=[types.Part(text=SAMPLE_QUESTION)])
    messages.append(user_message)

    print(f'\nSending sample question to {config_value.model.name}: "{SAMPLE_QUESTION}"')
    print("Waiting for response...")

    completion = track_genai_metrics(tracker, lambda: client.models.generate_content(
        model=config_value.model.name,
        contents=messages,
        config=types.GenerateContentConfig(
            system_instruction=system_instruction,
        )
    ))
    ai_response = completion.text

    ai_message = types.Content(role="model", parts=[types.Part(text=ai_response)])
    messages.append(ai_message)

    print(f"\nModel response:\n{ai_response}")

    summary = tracker.get_summary()
    print("\nDone! The AI config was evaluated and the following metrics were tracked:")
    print(f"  Duration:      {summary.duration_ms}ms")
    print(f"  Success:       {summary.success}")
    if summary.tokens:
        print(f"  Input tokens:  {summary.tokens.input}")
        print(f"  Output tokens: {summary.tokens.output}")
        print(f"  Total tokens:  {summary.tokens.total}")
    if summary.tool_calls:
        print(f"  Tool calls:    {', '.join(summary.tool_calls)}")

    # Flush pending events and close the client.
    ldclient.get().flush()
    ldclient.get().close()

if __name__ == "__main__":
    main()