move to att

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/main.go

@@ -124,4 +124,4 @@ func workflowExample() {
 	})
 
 	fmt.Println(resp.Choices[0].Message.Content)
-}
+}

sample-app/recipe_agent_example.go

@@ -53,7 +53,6 @@ func ingredientValidatorTool(ctx context.Context, agent *sdk.Agent, client *open
 
 	llmSpan := tool.LogPrompt(prompt)
 
-	// Make API call
 	resp, err := client.CreateChatCompletion(ctx, openai.ChatCompletionRequest{
 		Model: "gpt-3.5-turbo",
 		Messages: []openai.ChatCompletionMessage{
@@ -67,7 +66,6 @@ func ingredientValidatorTool(ctx context.Context, agent *sdk.Agent, client *open
 		return "", fmt.Errorf("CreateChatCompletion error: %w", err)
 	}
 
-	// Log completion
 	var completionMsgs []sdk.Message
 	for _, choice := range resp.Choices {
 		completionMsgs = append(completionMsgs, sdk.Message{
@@ -121,7 +119,6 @@ func nutritionCalculatorTool(ctx context.Context, agent *sdk.Agent, client *open
 
 	llmSpan := tool.LogPrompt(prompt)
 
-	// Make API call
 	resp, err := client.CreateChatCompletion(ctx, openai.ChatCompletionRequest{
 		Model: "gpt-3.5-turbo",
 		Messages: []openai.ChatCompletionMessage{
@@ -135,7 +132,6 @@ func nutritionCalculatorTool(ctx context.Context, agent *sdk.Agent, client *open
 		return "", fmt.Errorf("CreateChatCompletion error: %w", err)
 	}
 
-	// Log completion
 	var completionMsgs []sdk.Message
 	for _, choice := range resp.Choices {
 		completionMsgs = append(completionMsgs, sdk.Message{
@@ -191,7 +187,6 @@ func cookingTimeEstimatorTool(ctx context.Context, agent *sdk.Agent, client *ope
 
 	llmSpan := tool.LogPrompt(prompt)
 
-	// Make API call
 	resp, err := client.CreateChatCompletion(ctx, openai.ChatCompletionRequest{
 		Model: "gpt-3.5-turbo",
 		Messages: []openai.ChatCompletionMessage{
@@ -205,7 +200,6 @@ func cookingTimeEstimatorTool(ctx context.Context, agent *sdk.Agent, client *ope
 		return "", fmt.Errorf("CreateChatCompletion error: %w", err)
 	}
 
-	// Log completion
 	var completionMsgs []sdk.Message
 	for _, choice := range resp.Choices {
 		completionMsgs = append(completionMsgs, sdk.Message{
@@ -246,10 +240,10 @@ func runRecipeAgent() {
 	agent := traceloop.NewAgent(ctx, "recipe_generator", sdk.AgentAttributes{
 		Name: "recipe_generator",
 		AssociationProperties: associationProperties,
-	}, abTest)
+		ABTest: abTest,
+	})
 	defer agent.End()
 
-	// User request
 	userRequest := "Create a healthy pasta dish with vegetables"
 	fmt.Printf("User request: %s\n\n", userRequest)
 
@@ -284,7 +278,6 @@ func runRecipeAgent() {
 		return
 	}
 
-	// Log completion
 	var completionMsgs []sdk.Message
 	for _, choice := range resp.Choices {
 		completionMsgs = append(completionMsgs, sdk.Message{

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

@@ -15,7 +15,6 @@ type Agent struct {
 	workflow   *Workflow
 	ctx        context.Context
 	Attributes AgentAttributes `json:"agent_attributes"`
-	ABTest     *model.ABTest   `json:"ab_test"`
 }
 
 func (agent *Agent) End() {

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, agentAttrs AgentAttributes, abTest *model.ABTest) *Agent {
+func (instance *Traceloop) NewAgent(ctx context.Context, name string, agentAttrs AgentAttributes) *Agent {
 	aCtx, span := instance.getTracer().Start(ctx, fmt.Sprintf("%s.agent", name), apitrace.WithNewRoot())
 
 	attrs := []attribute.KeyValue{
@@ -152,8 +152,8 @@ func (instance *Traceloop) NewAgent(ctx context.Context, name string, agentAttrs
 		semconvai.LLMAgentName.String(name),
 	}
 
-	if abTest != nil {
-		for key, activeVarient := range abTest.VarientsKeys {
+	if agentAttrs.ABTest != nil {
+		for key, activeVarient := range agentAttrs.ABTest.VarientsKeys {
 			if activeVarient {
 				agentAttrs.AssociationProperties["ab_testing_variant"] = key
 				break
@@ -173,7 +173,6 @@ func (instance *Traceloop) NewAgent(ctx context.Context, name string, agentAttrs
 		workflow: nil,
 		ctx:      aCtx,
 		Attributes: agentAttrs,
-		ABTest: abTest,
 	}
 }
 

traceloop-sdk/tracing_types.go

@@ -1,5 +1,7 @@
 package traceloop
 
+import "github.com/traceloop/go-openllmetry/traceloop-sdk/model"
+
 type Message struct {
 	Index     int        `json:"index"`
 	Role      string     `json:"role"`
@@ -28,6 +30,7 @@ type Completion struct {
 type WorkflowAttributes struct {
 	Name                  string            `json:"workflow_name"`
 	AssociationProperties map[string]string `json:"association_properties"`
+	ABTest                *model.ABTest   `json:"ab_test"`
 }
 
 type ContextAttributes struct {
@@ -66,4 +69,5 @@ type ToolCallAttributes struct {
 type AgentAttributes struct {
 	Name                  string            `json:"agent_name"`
 	AssociationProperties map[string]string `json:"association_properties"`
+	ABTest                *model.ABTest   `json:"ab_test"`
 }

traceloop-sdk/workflow.go

@@ -14,10 +14,9 @@ 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, abTest *model.ABTest) *Workflow {
+func (instance *Traceloop) NewWorkflow(ctx context.Context, attrs WorkflowAttributes) *Workflow {
 	wCtx, span := instance.getTracer().Start(ctx, fmt.Sprintf("%s.workflow", attrs.Name), trace.WithNewRoot())
 
 	span.SetAttributes(
@@ -26,8 +25,8 @@ func (instance *Traceloop) NewWorkflow(ctx context.Context, attrs WorkflowAttrib
 		semconvai.TraceloopEntityName.String(attrs.Name),
 	)
 
-	if abTest != nil {
-		for key, activeVarient := range abTest.VarientsKeys {
+	if attrs.ABTest != nil {
+		for key, activeVarient := range attrs.ABTest.VarientsKeys {
 			if activeVarient {
 				span.SetAttributes(attribute.String("traceloop.association.properties.ab_testing_variant", key))
 				break
@@ -79,8 +78,8 @@ 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 workflow.Attributes.ABTest != nil {
+		for key, activeVarient := range workflow.Attributes.ABTest.VarientsKeys {
 			if activeVarient {
 				associationProperties["ab_testing_variant"] = key
 			}