From e39e84266fd7166dec6e0f694c6e4a7b5dce0516 Mon Sep 17 00:00:00 2001
From: Joschka Braun <joschka.braun@gmail.com>
Date: Mon, 12 Aug 2024 11:28:03 -0400
Subject: [PATCH] feat: sample logs

---
 src/types.ts                   | 3 +++
 src/utils/core/Trace.ts        | 3 ++-
 src/utils/core/TraceManager.ts | 4 ++++
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/src/types.ts b/src/types.ts
index 24e049d..4eb7641 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -55,6 +55,7 @@ export type Completion = {
   log_omit_inputs?: boolean;
   log_omit_outputs?: boolean;
   log_omit?: boolean;
+  log_sample_rate?: number;
   experiment_uuid?: string | null;
   project_uuid?: string;
 };
@@ -164,6 +165,7 @@ export type TraceLog = EvaluatedLog & {
   deployment_id?: string;
   output_for_eval_metrics?: string;
   apply_eval_frac?: number;
+  log_sample_rate?: number;
   cache_hit?: boolean;
   evaluation_metric_names?: string[];
   feedback_score?: number;
@@ -197,6 +199,7 @@ export type TraceOptions = {
   applyEvalFrac?: number;
   deploymentId?: string;
   target?: string;
+  logSampleRate?: number;
 };
 
 export type UpdateLog = {
diff --git a/src/utils/core/Trace.ts b/src/utils/core/Trace.ts
index 9b2925c..69cf847 100644
--- a/src/utils/core/Trace.ts
+++ b/src/utils/core/Trace.ts
@@ -37,7 +37,7 @@ export class Trace {
     executionOrder: number = 0,
   ) {
     this.evalFuncs = options?.evalFuncs || [];
-    const { endUserIdentifier, sessionId, deploymentId, ...opts } = options || {};
+    const { endUserIdentifier, sessionId, deploymentId, logSampleRate, ...opts } = options || {};
 
     this.id = genTraceId();
     this.name = name;
@@ -56,6 +56,7 @@ export class Trace {
       root_trace_id: this.rootId,
       children: this.children,
       end_user_identifier: endUserIdentifier,
+      log_sample_rate: logSampleRate,
       session_id: sessionId,
       deployment_id: deploymentId,
       ...opts,
diff --git a/src/utils/core/TraceManager.ts b/src/utils/core/TraceManager.ts
index 891b5ea..5412984 100644
--- a/src/utils/core/TraceManager.ts
+++ b/src/utils/core/TraceManager.ts
@@ -53,6 +53,10 @@ export class TraceManager {
       options.target = parentTarget;
     }
 
+    if (parentTrace && parentTrace.getLog().log_sample_rate !== options.logSampleRate) {
+      options.logSampleRate = parentTrace.getLog().log_sample_rate;
+    }
+
     const trace = new Trace(name, options, parentId, rootId, depth, executionOrder);
 
     if (parentTrace) {