chore: Update SDK dependency, use provider helpers, and fix broken examples (#19)

Author: jsonbailey

examples/bedrock_example.py

@@ -2,7 +2,7 @@
 import ldclient
 from ldclient import Context
 from ldclient.config import Config
-from ldai.client import LDAIClient, AIConfig, ModelConfig, ProviderConfig, LDMessage
+from ldai.client import LDAIClient
 import boto3
 
 client = boto3.client("bedrock-runtime", region_name="us-east-1")
@@ -49,12 +49,13 @@ def main():
     #       provider=ProviderConfig(name='bedrock'),
     #       messages=[LDMessage(role='system', content='You are a helpful assistant.')],
     #   )
-    #   config_value, tracker = aiclient.config(ai_config_key, context, default, {'myUserVariable': "Testing Variable"})
-    config_value, tracker = aiclient.config(
+    #   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"}
     )
+    tracker = config_value.tracker
 
     if not config_value.enabled:
         print("AI Config is disabled")

examples/gemini_example.py

@@ -2,7 +2,7 @@
 import ldclient
 from ldclient import Context
 from ldclient.config import Config
-from ldai.client import LDAIClient, AIConfig, ModelConfig, ProviderConfig, LDMessage
+from ldai.client import LDAIClient, LDMessage
 from ldai.tracker import TokenUsage
 from google import genai
 from google.genai import types
@@ -119,12 +119,13 @@ def main():
     #       provider=ProviderConfig(name='google'),
     #       messages=[LDMessage(role='system', content='You are a helpful assistant.')],
     #   )
-    #   config_value, tracker = aiclient.config(ai_config_key, context, default, {'myUserVariable': "Testing Variable"})
-    config_value, tracker = aiclient.config(
+    #   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"}
     )
+    tracker = config_value.tracker
 
     if not config_value.enabled:
         print("AI Config is disabled")

examples/langchain_example.py

@@ -1,9 +1,10 @@
 import os
+import asyncio
 import ldclient
 from ldclient import Context
 from ldclient.config import Config
-from ldai.client import LDAIClient, AIConfig, ModelConfig, ProviderConfig, LDMessage
-from ldai.tracker import TokenUsage
+from ldai.client import LDAIClient
+from ldai_langchain import LangChainProvider
 from langchain.chat_models import init_chat_model
 
 # Set sdk_key to your LaunchDarkly SDK key.
@@ -21,43 +22,7 @@ def map_provider_to_langchain(provider_name):
     lower_provider = provider_name.lower()
     return provider_mapping.get(lower_provider, lower_provider)
 
-def track_langchain_metrics(tracker, func):
-    """
-    Track LangChain-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=usage_data.get("input_tokens", 0),
-                output=usage_data.get("output_tokens", 0),
-                total=usage_data.get("total_tokens", 0) # LangChain also has values for input_token_details { cache_creation, cache_read }
-            )
-            tracker.track_tokens(token_usage)
-    except Exception:
-        tracker.track_error()
-        raise
-
-    return result
-
-def main():
+async def async_main():
     if not sdk_key:
         print("*** Please set the LAUNCHDARKLY_SDK_KEY env first")
         exit()
@@ -92,12 +57,13 @@ def main():
     #       provider=ProviderConfig(name='openai'),
     #       messages=[LDMessage(role='system', content='You are a helpful assistant.')],
     #   )
-    #   config_value, tracker = aiclient.config(ai_config_key, context, default, {'myUserVariable': "Testing Variable"})
-    config_value, tracker = aiclient.config(
+    #   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"}
     )
+    tracker = config_value.tracker
 
     if not config_value.enabled:
         print("AI Config is disabled")
@@ -120,8 +86,11 @@ def main():
         print("User Input:\n", USER_INPUT)
         messages.append({'role': 'user', 'content': USER_INPUT})
 
-        # Track the LangChain completion with LaunchDarkly metrics
-        completion = track_langchain_metrics(tracker, lambda: llm.invoke(messages))
+        # Track the LangChain completion with LaunchDarkly metrics using the LD LangChain provider's extractor
+        completion = await tracker.track_metrics_of(
+            lambda: llm.ainvoke(messages),
+            LangChainProvider.get_ai_metrics_from_response,
+        )
         ai_response = completion.content
 
         # Add the AI response to the conversation history.
@@ -139,5 +108,9 @@ def main():
     ldclient.get().close()
 
 
+def main():
+    asyncio.run(async_main())
+
+
 if __name__ == "__main__":
     main()

examples/langgraph_agent_example.py

@@ -3,8 +3,9 @@
 from pprint import pprint
 from ldclient import Context
 from ldclient.config import Config
-from ldai.client import LDAIClient, LDAIAgentConfig, LDAIAgentDefaults
+from ldai.client import LDAIClient
 from ldai.tracker import TokenUsage
+from ldai_langchain import LangChainProvider
 from langchain.chat_models import init_chat_model
 from langgraph.prebuilt import create_react_agent
 
@@ -29,28 +30,25 @@ def track_langgraph_metrics(tracker, func):
     try:
         result = tracker.track_duration_of(func)
         tracker.track_success()
-        
-        # For LangGraph agents, usage_metadata is included on all messages that used AI
+
         total_input_tokens = 0
         total_output_tokens = 0
         total_tokens = 0
-
         if "messages" in result:
-            for message in result['messages']:
-                # Check for usage_metadata directly on the message
-                if hasattr(message, "usage_metadata") and message.usage_metadata:
-                    usage_data = message.usage_metadata
-                    total_input_tokens += usage_data.get("input_tokens", 0)
-                    total_output_tokens += usage_data.get("output_tokens", 0)
-                    total_tokens += usage_data.get("total_tokens", 0)
-
+            for message in result["messages"]:
+                metrics = LangChainProvider.get_ai_metrics_from_response(message)
+                if metrics.usage:
+                    total_input_tokens += metrics.usage.input
+                    total_output_tokens += metrics.usage.output
+                    total_tokens += metrics.usage.total
         if total_tokens > 0:
-            token_usage = TokenUsage(
-                input=total_input_tokens,
-                output=total_output_tokens,
-                total=total_tokens
+            tracker.track_tokens(
+                TokenUsage(
+                    input=total_input_tokens,
+                    output=total_output_tokens,
+                    total=total_tokens,
+                )
             )
-            tracker.track_tokens(token_usage)
     except Exception:
         tracker.track_error()
         raise
@@ -88,17 +86,12 @@ def main():
     # Pass a default for improved resiliency when the agent config is unavailable
     # or LaunchDarkly is unreachable; omit for a disabled default.
     # Example (enabled default):
-    #   default = LDAIAgentDefaults(
+    #   default = AIAgentConfigDefault(
     #       enabled=True,
     #       instructions='You are a helpful assistant.',
     #   )
     #   agent_config = aiclient.agent_config(agent_config_key, context, default=default)
-    agent_config = aiclient.agent_config(
-        LDAIAgentConfig(
-            key=agent_config_key,
-        ),
-        context
-    )
+    agent_config = aiclient.agent_config(agent_config_key, context)
 
     if not agent_config.enabled:
         print("AI Agent Config is disabled")

examples/langgraph_multi_agent_example.py

@@ -2,8 +2,9 @@
 import ldclient
 from ldclient import Context
 from ldclient.config import Config
-from ldai.client import LDAIClient, LDAIAgentConfig, LDAIAgentDefaults
+from ldai.client import LDAIClient
 from ldai.tracker import TokenUsage
+from ldai_langchain import LangChainProvider
 from langchain.chat_models import init_chat_model
 from langgraph.prebuilt import create_react_agent
 from langgraph.graph import StateGraph, END
@@ -39,30 +40,26 @@ def track_langgraph_metrics(tracker, func, prev_message_count=0):
     try:
         result = tracker.track_duration_of(func)
         tracker.track_success()
-        
-        # For LangGraph agents, usage_metadata is included on all messages that used AI
+
         total_input_tokens = 0
         total_output_tokens = 0
         total_tokens = 0
-
         if "messages" in result:
-            # Only look at messages that were added during this function call
-            new_messages = result['messages'][prev_message_count:]
+            new_messages = result["messages"][prev_message_count:]
             for message in new_messages:
-                # Check for usage_metadata directly on the message
-                if hasattr(message, "usage_metadata") and message.usage_metadata:
-                    usage_data = message.usage_metadata
-                    total_input_tokens += usage_data.get("input_tokens", 0)
-                    total_output_tokens += usage_data.get("output_tokens", 0)
-                    total_tokens += usage_data.get("total_tokens", 0)
-
+                metrics = LangChainProvider.get_ai_metrics_from_response(message)
+                if metrics.usage:
+                    total_input_tokens += metrics.usage.input
+                    total_output_tokens += metrics.usage.output
+                    total_tokens += metrics.usage.total
         if total_tokens > 0:
-            token_usage = TokenUsage(
-                input=total_input_tokens,
-                output=total_output_tokens,
-                total=total_tokens
+            tracker.track_tokens(
+                TokenUsage(
+                    input=total_input_tokens,
+                    output=total_output_tokens,
+                    total=total_tokens,
+                )
             )
-            tracker.track_tokens(token_usage)
     except Exception:
         tracker.track_error()
         raise
@@ -76,17 +73,12 @@ def create_agent_with_config(aiclient, config_key, context):
     # Pass a default for improved resiliency when the agent config is unavailable
     # or LaunchDarkly is unreachable; omit for a disabled default.
     # Example (enabled default):
-    #   default = LDAIAgentDefaults(
+    #   default = AIAgentConfigDefault(
     #       enabled=True,
     #       instructions='You are a helpful assistant.',
     #   )
-    #   agent_config = aiclient.agent_config(LDAIAgentConfig(key=config_key, default=default), context)
-    agent_config = aiclient.agent_config(
-        LDAIAgentConfig(
-            key=config_key,
-        ),
-        context
-    )
+    #   agent_config = aiclient.agent_config(config_key, context, default=default)
+    agent_config = aiclient.agent_config(config_key, context)
 
     if not agent_config.enabled:
         return None, None, True

examples/openai_example.py

@@ -2,7 +2,7 @@
 import ldclient
 from ldclient import Context
 from ldclient.config import Config
-from ldai.client import LDAIClient, AIConfig, ModelConfig, ProviderConfig, LDMessage
+from ldai.client import LDAIClient
 from openai import OpenAI
 
 openai_client = OpenAI()
@@ -50,12 +50,13 @@ def main():
     #       provider=ProviderConfig(name='openai'),
     #       messages=[LDMessage(role='system', content='You are a helpful assistant.')],
     #   )
-    #   config_value, tracker = aiclient.config(ai_config_key, context, default, {'myUserVariable': "Testing Variable"})
-    config_value, tracker = aiclient.config(
+    #   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"}
     )
+    tracker = config_value.tracker
 
     if not config_value.enabled:
         print("AI Config is disabled")

pyproject.toml

@@ -20,9 +20,9 @@ direct-judge-example = 'examples.direct_judge_example:main'
 
 [tool.poetry.dependencies]
 python = "^3.10"
-launchdarkly-server-sdk-ai = "^0.16.0"
-launchdarkly-server-sdk-ai-langchain = "^0.3.0"
-launchdarkly-server-sdk-ai-openai = "^0.1.0"
+launchdarkly-server-sdk-ai = "^0.16.1"
+launchdarkly-server-sdk-ai-langchain = "^0.3.2"
+launchdarkly-server-sdk-ai-openai = "^0.2.1"
 launchdarkly-observability = { version = ">=0.1.0", optional = true }
 
 boto3 = { version = ">=0.2.0", optional = true }