From f52bf8f1a0e66069d2ee1436e3575c23fb5fc803 Mon Sep 17 00:00:00 2001 From: Daniel OBrien Date: Tue, 15 Oct 2024 12:53:57 -0400 Subject: [PATCH 1/7] updates to Python tracking --- ldai/client.py | 6 +-- ldai/tracking_utils.py | 24 ++---------- ldai/types.py | 84 +++++++++++++++++++++++++++++++++++------- 3 files changed, 76 insertions(+), 38 deletions(-) diff --git a/ldai/client.py b/ldai/client.py index 77fbc24..3391b29 100644 --- a/ldai/client.py +++ b/ldai/client.py @@ -1,11 +1,11 @@ import json from typing import Any, Dict, Optional from ldclient import Context -#from config import LDAIConfig, LDAIConfigTracker from ldclient.client import LDClient import chevron from ldai.tracker import LDAIConfigTracker +from ldai.types import AIConfig class LDAIClient: """The LaunchDarkly AI SDK client object.""" @@ -13,7 +13,7 @@ class LDAIClient: def __init__(self, client: LDClient): self.client = client - def model_config(self, key: str, context: Context, default_value: str, variables: Optional[Dict[str, Any]] = None) -> Any: + def model_config(self, key: str, context: Context, default_value: str, variables: Optional[Dict[str, Any]] = None) -> AIConfig: """Get the value of a model configuration asynchronously. Args: @@ -40,8 +40,6 @@ def model_config(self, key: str, context: Context, default_value: str, variables for entry in variation['prompt'] ] - #return detail.value, - return { 'config': variation, 'tracker': LDAIConfigTracker(self.client, variation['_ldMeta']['variationId'], key, context) diff --git a/ldai/tracking_utils.py b/ldai/tracking_utils.py index a44bf4a..8383b01 100644 --- a/ldai/tracking_utils.py +++ b/ldai/tracking_utils.py @@ -1,23 +1,5 @@ from typing import Union -from ldai.types import BedrockTokenUsage, TokenMetrics, TokenUsage, UnderscoreTokenUsage +from ldai.types import BedrockTokenUsage, TokenMetrics, OpenAITokenUsage, UnderscoreTokenUsage -def usage_to_token_metrics(usage: Union[TokenUsage, UnderscoreTokenUsage, BedrockTokenUsage]) -> TokenMetrics: - def get_attr(obj, attr, default=0): - if isinstance(obj, dict): - return obj.get(attr, default) - return getattr(obj, attr, default) - - if 'inputTokens' in usage and 'outputTokens' in usage: - # Bedrock usage - return { - 'total': get_attr(usage, 'totalTokens'), - 'input': get_attr(usage, 'inputTokens'), - 'output': get_attr(usage, 'outputTokens'), - } - - # OpenAI usage (both camelCase and snake_case) - return { - 'total': get_attr(usage, 'total_tokens', get_attr(usage, 'totalTokens', 0)), - 'input': get_attr(usage, 'prompt_tokens', get_attr(usage, 'promptTokens', 0)), - 'output': get_attr(usage, 'completion_tokens', get_attr(usage, 'completionTokens', 0)), - } \ No newline at end of file +def usage_to_token_metrics(usage: Union[OpenAITokenUsage, UnderscoreTokenUsage, BedrockTokenUsage]) -> TokenMetrics: + return usage.to_metrics() \ No newline at end of file diff --git a/ldai/types.py b/ldai/types.py index e5a962c..c9cf29c 100644 --- a/ldai/types.py +++ b/ldai/types.py @@ -1,5 +1,32 @@ from enum import Enum -from typing import TypedDict +from typing import Callable, TypedDict +from dataclasses import dataclass + +@dataclass +class TokenMetrics(TypedDict): + total: int + input: int + output: int # type: ignore + +class AIConfigData(TypedDict): + config: dict + prompt: any + _ldMeta: dict + +class AITracker(TypedDict): + track_duration: Callable[..., None] + track_tokens: Callable[..., None] + track_error: Callable[..., None] + track_generation: Callable[..., None] + track_feedback: Callable[..., None] + +class AIConfig(): + def __init__(self, config: AIConfigData, tracker: AITracker): + self._config = config + self._tracker = tracker + + config: AIConfigData + tracker: AITracker class FeedbackKind(Enum): Positive = "positive" @@ -10,17 +37,48 @@ class TokenUsage(TypedDict): prompt_tokens: int completion_tokens: int -class UnderscoreTokenUsage(TypedDict): - total_tokens: int - prompt_tokens: int - completion_tokens: int + def to_metrics(self): + return { + 'total': self['total_tokens'], + 'input': self['prompt_tokens'], + 'output': self['completion_tokens'], + } -class BedrockTokenUsage(TypedDict): - totalTokens: int - inputTokens: int - outputTokens: int +class OpenAITokenUsage: + def __init__(self, data: any): + self.total_tokens = data.total_tokens + self.prompt_tokens = data.prompt_tokens + self.completion_tokens = data.completion_tokens -class TokenMetrics(TypedDict): - total: int - input: int - output: int # type: ignore \ No newline at end of file + def to_metrics(self) -> TokenMetrics: + return { + 'total': self.total_tokens, + 'input': self.prompt_tokens, + 'output': self.completion_tokens, + } + +class UnderscoreTokenUsage: + def __init__(self, data: dict): + self.total_tokens = data.get('total_tokens', 0) + self.prompt_tokens = data.get('prompt_tokens', 0) + self.completion_tokens = data.get('completion_tokens', 0) + + def to_metrics(self) -> TokenMetrics: + return { + 'total': self.total_tokens, + 'input': self.prompt_tokens, + 'output': self.completion_tokens, + } + +class BedrockTokenUsage: + def __init__(self, data: dict): + self.totalTokens = data.get('totalTokens', 0) + self.inputTokens = data.get('inputTokens', 0) + self.outputTokens = data.get('outputTokens', 0) + + def to_metrics(self) -> TokenMetrics: + return { + 'total': self.totalTokens, + 'input': self.inputTokens, + 'output': self.outputTokens, + } \ No newline at end of file From dc073a48752820e601449df77613e535130c3fe9 Mon Sep 17 00:00:00 2001 From: Daniel OBrien Date: Wed, 16 Oct 2024 10:12:56 -0400 Subject: [PATCH 2/7] PR Feedback Remove tracking utils return AI Config Remove errant print --- ldai/client.py | 7 +------ ldai/tracker.py | 3 +-- ldai/tracking_utils.py | 5 ----- ldai/types.py | 7 ++----- 4 files changed, 4 insertions(+), 18 deletions(-) delete mode 100644 ldai/tracking_utils.py diff --git a/ldai/client.py b/ldai/client.py index 3391b29..a1463f8 100644 --- a/ldai/client.py +++ b/ldai/client.py @@ -1,4 +1,3 @@ -import json from typing import Any, Dict, Optional from ldclient import Context from ldclient.client import LDClient @@ -31,7 +30,6 @@ def model_config(self, key: str, context: Context, default_value: str, variables if variables: all_variables.update(variables) - print(variation) variation['prompt'] = [ { **entry, @@ -40,10 +38,7 @@ def model_config(self, key: str, context: Context, default_value: str, variables for entry in variation['prompt'] ] - return { - 'config': variation, - 'tracker': LDAIConfigTracker(self.client, variation['_ldMeta']['variationId'], key, context) - } + return AIConfig(config=variation, tracker=LDAIConfigTracker(self.client, variation['_ldMeta']['variationId'], key, context)) def interpolate_template(self, template: str, variables: Dict[str, Any]) -> str: """Interpolate the template with the given variables. diff --git a/ldai/tracker.py b/ldai/tracker.py index b934e22..fb803ad 100644 --- a/ldai/tracker.py +++ b/ldai/tracker.py @@ -1,6 +1,5 @@ from typing import Dict, Union from ldclient import Context, LDClient -from ldai.tracking_utils import usage_to_token_metrics from ldai.types import BedrockTokenUsage, FeedbackKind, TokenUsage, UnderscoreTokenUsage class LDAIConfigTracker: @@ -20,7 +19,7 @@ def track_duration(self, duration: int) -> None: self.ld_client.track('$ld:ai:duration:total', self.context, self.get_track_data(), duration) def track_tokens(self, tokens: Union[TokenUsage, UnderscoreTokenUsage, BedrockTokenUsage]) -> None: - token_metrics = usage_to_token_metrics(tokens) + token_metrics = tokens.to_metrics() if token_metrics['total'] > 0: self.ld_client.track('$ld:ai:tokens:total', self.context, self.get_track_data(), token_metrics['total']) if token_metrics['input'] > 0: diff --git a/ldai/tracking_utils.py b/ldai/tracking_utils.py deleted file mode 100644 index 8383b01..0000000 --- a/ldai/tracking_utils.py +++ /dev/null @@ -1,5 +0,0 @@ -from typing import Union -from ldai.types import BedrockTokenUsage, TokenMetrics, OpenAITokenUsage, UnderscoreTokenUsage - -def usage_to_token_metrics(usage: Union[OpenAITokenUsage, UnderscoreTokenUsage, BedrockTokenUsage]) -> TokenMetrics: - return usage.to_metrics() \ No newline at end of file diff --git a/ldai/types.py b/ldai/types.py index c9cf29c..a3c69af 100644 --- a/ldai/types.py +++ b/ldai/types.py @@ -22,11 +22,8 @@ class AITracker(TypedDict): class AIConfig(): def __init__(self, config: AIConfigData, tracker: AITracker): - self._config = config - self._tracker = tracker - - config: AIConfigData - tracker: AITracker + self.config = config + self.tracker = tracker class FeedbackKind(Enum): Positive = "positive" From 6c0742dbce310746f094d99fb295140d0f0bdc62 Mon Sep 17 00:00:00 2001 From: Daniel OBrien Date: Wed, 16 Oct 2024 10:14:38 -0400 Subject: [PATCH 3/7] set version back to 0.0.1 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 729c2d5..e51a9d1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "launchdarkly-server-sdk-ai" -version = "0.0.14" +version = "0.0.1" description = "LaunchDarkly SDK for AI" authors = ["LaunchDarkly "] license = "Apache-2.0" From fe50860adb89173d08efd298db0eea9818f7a8f3 Mon Sep 17 00:00:00 2001 From: Daniel OBrien Date: Fri, 18 Oct 2024 08:57:19 -0400 Subject: [PATCH 4/7] add tracking decorator --- ldai/tracker.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/ldai/tracker.py b/ldai/tracker.py index fb803ad..4a19e31 100644 --- a/ldai/tracker.py +++ b/ldai/tracker.py @@ -1,3 +1,4 @@ +import time from typing import Dict, Union from ldclient import Context, LDClient from ldai.types import BedrockTokenUsage, FeedbackKind, TokenUsage, UnderscoreTokenUsage @@ -18,6 +19,16 @@ def get_track_data(self): def track_duration(self, duration: int) -> None: self.ld_client.track('$ld:ai:duration:total', self.context, self.get_track_data(), duration) + def track_duration_of(self, func): + def wrapper(*args, **kwargs): + start_time = time.time() + result = func(*args, **kwargs) + end_time = time.time() + duration = int((end_time - start_time) * 1000) # duration in milliseconds + self.track_duration(duration) + return result + return wrapper + def track_tokens(self, tokens: Union[TokenUsage, UnderscoreTokenUsage, BedrockTokenUsage]) -> None: token_metrics = tokens.to_metrics() if token_metrics['total'] > 0: From bff84efe585c294347ddc23d686fe5321070a371 Mon Sep 17 00:00:00 2001 From: Daniel OBrien Date: Fri, 18 Oct 2024 09:42:18 -0400 Subject: [PATCH 5/7] change duration tracker implementation type --- ldai/tracker.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/ldai/tracker.py b/ldai/tracker.py index 4a19e31..c1ae9f5 100644 --- a/ldai/tracker.py +++ b/ldai/tracker.py @@ -19,15 +19,13 @@ def get_track_data(self): def track_duration(self, duration: int) -> None: self.ld_client.track('$ld:ai:duration:total', self.context, self.get_track_data(), duration) - def track_duration_of(self, func): - def wrapper(*args, **kwargs): - start_time = time.time() - result = func(*args, **kwargs) - end_time = time.time() - duration = int((end_time - start_time) * 1000) # duration in milliseconds - self.track_duration(duration) - return result - return wrapper + def track_duration_of(self, func, *args, **kwargs): + start_time = time.time() + result = func(*args, **kwargs) + end_time = time.time() + duration = int((end_time - start_time) * 1000) # duration in milliseconds + self.track_duration(duration) + return result def track_tokens(self, tokens: Union[TokenUsage, UnderscoreTokenUsage, BedrockTokenUsage]) -> None: token_metrics = tokens.to_metrics() From 1b5e2a554597ddf63fd8d3ab489294b21472cc46 Mon Sep 17 00:00:00 2001 From: Daniel OBrien Date: Fri, 18 Oct 2024 10:04:07 -0400 Subject: [PATCH 6/7] change types to dataclasses --- ldai/types.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/ldai/types.py b/ldai/types.py index a3c69af..efa8300 100644 --- a/ldai/types.py +++ b/ldai/types.py @@ -1,19 +1,21 @@ from enum import Enum -from typing import Callable, TypedDict +from typing import Callable from dataclasses import dataclass @dataclass -class TokenMetrics(TypedDict): +class TokenMetrics(): total: int input: int output: int # type: ignore -class AIConfigData(TypedDict): +@dataclass + +class AIConfigData(): config: dict prompt: any _ldMeta: dict -class AITracker(TypedDict): +class AITracker(): track_duration: Callable[..., None] track_tokens: Callable[..., None] track_error: Callable[..., None] @@ -29,7 +31,9 @@ class FeedbackKind(Enum): Positive = "positive" Negative = "negative" -class TokenUsage(TypedDict): +@dataclass + +class TokenUsage(): total_tokens: int prompt_tokens: int completion_tokens: int From f68fa4d18b45e98882794da114ce71d76a36c5cf Mon Sep 17 00:00:00 2001 From: Daniel OBrien Date: Fri, 18 Oct 2024 10:04:28 -0400 Subject: [PATCH 7/7] add openai tracker wrapper reorganize method order --- ldai/tracker.py | 34 ++++++++++++++++++++-------------- 1 file changed, 20 insertions(+), 14 deletions(-) diff --git a/ldai/tracker.py b/ldai/tracker.py index c1ae9f5..9be9d9c 100644 --- a/ldai/tracker.py +++ b/ldai/tracker.py @@ -1,7 +1,7 @@ import time from typing import Dict, Union from ldclient import Context, LDClient -from ldai.types import BedrockTokenUsage, FeedbackKind, TokenUsage, UnderscoreTokenUsage +from ldai.types import BedrockTokenUsage, FeedbackKind, OpenAITokenUsage, TokenUsage, UnderscoreTokenUsage class LDAIConfigTracker: def __init__(self, ld_client: LDClient, variation_id: str, config_key: str, context: Context): @@ -27,23 +27,29 @@ def track_duration_of(self, func, *args, **kwargs): self.track_duration(duration) return result - def track_tokens(self, tokens: Union[TokenUsage, UnderscoreTokenUsage, BedrockTokenUsage]) -> None: - token_metrics = tokens.to_metrics() - if token_metrics['total'] > 0: - self.ld_client.track('$ld:ai:tokens:total', self.context, self.get_track_data(), token_metrics['total']) - if token_metrics['input'] > 0: - self.ld_client.track('$ld:ai:tokens:input', self.context, self.get_track_data(), token_metrics['input']) - if token_metrics['output'] > 0: - self.ld_client.track('$ld:ai:tokens:output', self.context, self.get_track_data(), token_metrics['output']) - def track_error(self, error: int) -> None: self.ld_client.track('$ld:ai:error', self.context, self.get_track_data(), error) - def track_generation(self, generation: int) -> None: - self.ld_client.track('$ld:ai:generation', self.context, self.get_track_data(), generation) - def track_feedback(self, feedback: Dict[str, FeedbackKind]) -> None: if feedback['kind'] == FeedbackKind.Positive: self.ld_client.track('$ld:ai:feedback:user:positive', self.context, self.get_track_data(), 1) elif feedback['kind'] == FeedbackKind.Negative: - self.ld_client.track('$ld:ai:feedback:user:negative', self.context, self.get_track_data(), 1) \ No newline at end of file + self.ld_client.track('$ld:ai:feedback:user:negative', self.context, self.get_track_data(), 1) + + def track_generation(self, generation: int) -> None: + self.ld_client.track('$ld:ai:generation', self.context, self.get_track_data(), generation) + + def track_openai(self, func, *args, **kwargs): + result = self.track_duration_of(func, *args, **kwargs) + if result.usage: + self.track_tokens(OpenAITokenUsage(result.usage)) + return result + + def track_tokens(self, tokens: Union[TokenUsage, UnderscoreTokenUsage, BedrockTokenUsage]) -> None: + token_metrics = tokens.to_metrics() + if token_metrics['total'] > 0: + self.ld_client.track('$ld:ai:tokens:total', self.context, self.get_track_data(), token_metrics['total']) + if token_metrics['input'] > 0: + self.ld_client.track('$ld:ai:tokens:input', self.context, self.get_track_data(), token_metrics['input']) + if token_metrics['output'] > 0: + self.ld_client.track('$ld:ai:tokens:output', self.context, self.get_track_data(), token_metrics['output']) \ No newline at end of file