works

Author: nina-kollman

sample-app/generate_joke_workflow_example.go

@@ -264,7 +264,7 @@ func runJokeWorkflow() {
 			"user_id": "user_12345",
 			"chat_id": "chat_1234",
 		},
-	})
+	}, nil)
 	defer wf.End()
 
 	// Execute workflow steps

sample-app/recipe_agent_example.go

@@ -7,11 +7,18 @@ import (
 
 	"github.com/sashabaranov/go-openai"
 	sdk "github.com/traceloop/go-openllmetry/traceloop-sdk"
+	"github.com/traceloop/go-openllmetry/traceloop-sdk/model"
 )
 
 var associationProperties = map[string]string{
 	"user_id": "user_67890",
-	"ab_testing_variant": "variant_a",
+}
+
+var abTest = &model.ABTest{
+	VarientsKeys: map[string]bool{
+		"variant_a": false,
+		"variant_b": true,
+	},
 }
 
 // ingredientValidatorTool validates that requested ingredients are available/safe
@@ -236,7 +243,10 @@ func runRecipeAgent() {
 	client := openai.NewClient(os.Getenv("OPENAI_API_KEY"))
 
 	// Create standalone agent with association properties
-	agent := traceloop.NewAgent(ctx, "recipe_generator", associationProperties)
+	agent := traceloop.NewAgent(ctx, "recipe_generator", sdk.AgentAttributes{
+		Name: "recipe_generator",
+		AssociationProperties: associationProperties,
+	}, abTest)
 	defer agent.End()
 
 	// User request

sample-app/workflow_example.go

@@ -26,7 +26,7 @@ func workflowMain() {
 
 	wf := traceloop.NewWorkflow(ctx, tlp.WorkflowAttributes{
 		Name: "history_generation",
-	})
+	}, nil)
 	defer wf.End()
 
 	factGenTask := wf.NewTask("current_date_fact_generation")

traceloop-sdk/agent.go

@@ -5,16 +5,17 @@ import (
 	"fmt"
 
 	semconvai "github.com/traceloop/go-openllmetry/semconv-ai"
+	"github.com/traceloop/go-openllmetry/traceloop-sdk/model"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
-	"github.com/traceloop/go-openllmetry/traceloop-sdk/model"
 )
 
 type Agent struct {
 	sdk        *Traceloop
 	workflow   *Workflow
 	ctx        context.Context
 	Attributes AgentAttributes `json:"agent_attributes"`
+	ABTest     *model.ABTest   `json:"ab_test"`
 }
 
 func (agent *Agent) End() {
@@ -70,5 +71,3 @@ func (agent *Agent) NewTool(name string, toolType string, toolFunction ToolFunct
 		Function: toolFunction,
 	}
 }
-
-

traceloop-sdk/model/types.go

@@ -8,3 +8,8 @@ const (
 	SpanKindTask SpanKind = "task"
 	SpanKindWorkflow SpanKind = "workflow"
 )
+
+// The varient that is active will be added to the trace.
+type ABTest struct {
+	VarientsKeys map[string]bool `json:"varients_keys"`
+}

traceloop-sdk/sdk.go

@@ -143,7 +143,7 @@ func (instance *Traceloop) getTracer() apitrace.Tracer {
 }
 
 // NewAgent creates a standalone agent (without a workflow)
-func (instance *Traceloop) NewAgent(ctx context.Context, name string, associationProperties map[string]string) *Agent {
+func (instance *Traceloop) NewAgent(ctx context.Context, name string, agentAttrs AgentAttributes, abTest *model.ABTest) *Agent {
 	aCtx, span := instance.getTracer().Start(ctx, fmt.Sprintf("%s.agent", name), apitrace.WithNewRoot())
 
 	attrs := []attribute.KeyValue{
@@ -152,8 +152,17 @@ func (instance *Traceloop) NewAgent(ctx context.Context, name string, associatio
 		semconvai.LLMAgentName.String(name),
 	}
 
+	if abTest != nil {
+		for key, activeVarient := range abTest.VarientsKeys {
+			if activeVarient {
+				agentAttrs.AssociationProperties["ab_testing_variant"] = key
+				break
+			}
+		}
+	}
+
 	// Add association properties if provided
-	for key, value := range associationProperties {
+	for key, value := range agentAttrs.AssociationProperties {
 		attrs = append(attrs, attribute.String("traceloop.association.properties."+key, value))
 	}
 
@@ -163,10 +172,8 @@ func (instance *Traceloop) NewAgent(ctx context.Context, name string, associatio
 		sdk:      instance,
 		workflow: nil,
 		ctx:      aCtx,
-		Attributes: AgentAttributes{
-			Name:                  name,
-			AssociationProperties: associationProperties,
-		},
+		Attributes: agentAttrs,
+		ABTest: abTest,
 	}
 }
 

traceloop-sdk/workflow.go

@@ -14,9 +14,10 @@ type Workflow struct {
 	sdk        *Traceloop
 	ctx        context.Context
 	Attributes WorkflowAttributes `json:"workflow_attributes"`
+	ABTest 	   *model.ABTest `json:"ab_test"`
 }
 
-func (instance *Traceloop) NewWorkflow(ctx context.Context, attrs WorkflowAttributes) *Workflow {
+func (instance *Traceloop) NewWorkflow(ctx context.Context, attrs WorkflowAttributes, abTest *model.ABTest) *Workflow {
 	wCtx, span := instance.getTracer().Start(ctx, fmt.Sprintf("%s.workflow", attrs.Name), trace.WithNewRoot())
 
 	span.SetAttributes(
@@ -25,6 +26,15 @@ func (instance *Traceloop) NewWorkflow(ctx context.Context, attrs WorkflowAttrib
 		semconvai.TraceloopEntityName.String(attrs.Name),
 	)
 
+	if abTest != nil {
+		for key, activeVarient := range abTest.VarientsKeys {
+			if activeVarient {
+				span.SetAttributes(attribute.String("traceloop.association.properties.ab_testing_variant", key))
+				break
+			}
+		}	
+	}
+
 	return &Workflow{
 		sdk:        instance,
 		ctx:        wCtx,
@@ -69,6 +79,13 @@ func (workflow *Workflow) NewAgent(name string, associationProperties map[string
 		semconvai.TraceloopEntityName.String(name),
 	}
 
+	if workflow.ABTest != nil {
+		for key, activeVarient := range workflow.ABTest.VarientsKeys {
+			if activeVarient {
+				associationProperties["ab_testing_variant"] = key
+			}
+		}
+	}
 	// Add agent-specific association properties to the span
 	for key, value := range associationProperties {
 		attrs = append(attrs, attribute.String("traceloop.association.properties."+key, value))