From 098918569b207b1819eacf4ed3658678cbd6c686 Mon Sep 17 00:00:00 2001
From: Forrest Yan <2674978072@qq.com>
Date: Sun, 9 Mar 2025 21:08:48 +0800
Subject: [PATCH] Reorganize workflow examples and improve code structure

---
 sample-app/main.go             | 41 +++++++++++++++++++++++++++++-----
 sample-app/workflow_example.go | 41 +++++-----------------------------
 2 files changed, 41 insertions(+), 41 deletions(-)

diff --git a/sample-app/main.go b/sample-app/main.go
index 249fc3c..c48f2d1 100644
--- a/sample-app/main.go
+++ b/sample-app/main.go
@@ -10,7 +10,7 @@ import (
 	tlp "github.com/traceloop/go-openllmetry/traceloop-sdk"
 )
 
-func workflow_example() {
+func main() {
 	ctx := context.Background()
 
 	traceloop, err := tlp.NewClient(ctx, tlp.Config{
@@ -24,6 +24,14 @@ func workflow_example() {
 		return
 	}
 
+	wf := traceloop.NewWorkflow(ctx, tlp.WorkflowAttributes{
+		Name: "history_generation",
+	})
+	defer wf.End()
+
+	factGenTask := wf.NewTask("current_date_fact_generation")
+	defer factGenTask.End()
+
 	request, err := traceloop.GetOpenAIChatCompletionRequest("example-prompt", map[string]interface{}{"date": time.Now().Format("01/02")})
 	if err != nil {
 		fmt.Printf("GetOpenAIChatCompletionRequest error: %v\n", err)
@@ -39,17 +47,13 @@ func workflow_example() {
 		})
 	}
 
-	llmSpan, err := traceloop.LogPrompt(
-		ctx,
+	llmSpan, err := factGenTask.LogPrompt(
 		tlp.Prompt{
 			Vendor:   "openai",
 			Mode:     "chat",
 			Model:    request.Model,
 			Messages: promptMsgs,
 		},
-		tlp.WorkflowAttributes{
-			Name: "example-workflow",
-		},
 	)
 	if err != nil {
 		fmt.Printf("LogPrompt error: %v\n", err)
@@ -84,5 +88,30 @@ func workflow_example() {
 		PromptTokens:     resp.Usage.PromptTokens,
 	})
 
+	someOtherTask := wf.NewTask("some_other_task")
+	defer someOtherTask.End()
+
+	otherPrompt, _ := someOtherTask.LogPrompt(tlp.Prompt{
+		Vendor: "openai",
+		Mode:   "chat",
+		Model:  request.Model,
+		Messages: []tlp.Message{
+			{
+				Index:   0,
+				Content: "some other prompt",
+				Role:    "user",
+			},
+		},
+	})
+
+	otherPrompt.LogCompletion(ctx, tlp.Completion{
+		Model:    resp.Model,
+		Messages: completionMsgs,
+	}, tlp.Usage{
+		TotalTokens:      resp.Usage.TotalTokens,
+		CompletionTokens: resp.Usage.CompletionTokens,
+		PromptTokens:     resp.Usage.PromptTokens,
+	})
+
 	fmt.Println(resp.Choices[0].Message.Content)
 }
diff --git a/sample-app/workflow_example.go b/sample-app/workflow_example.go
index c48f2d1..249fc3c 100644
--- a/sample-app/workflow_example.go
+++ b/sample-app/workflow_example.go
@@ -10,7 +10,7 @@ import (
 	tlp "github.com/traceloop/go-openllmetry/traceloop-sdk"
 )
 
-func main() {
+func workflow_example() {
 	ctx := context.Background()
 
 	traceloop, err := tlp.NewClient(ctx, tlp.Config{
@@ -24,14 +24,6 @@ func main() {
 		return
 	}
 
-	wf := traceloop.NewWorkflow(ctx, tlp.WorkflowAttributes{
-		Name: "history_generation",
-	})
-	defer wf.End()
-
-	factGenTask := wf.NewTask("current_date_fact_generation")
-	defer factGenTask.End()
-
 	request, err := traceloop.GetOpenAIChatCompletionRequest("example-prompt", map[string]interface{}{"date": time.Now().Format("01/02")})
 	if err != nil {
 		fmt.Printf("GetOpenAIChatCompletionRequest error: %v\n", err)
@@ -47,13 +39,17 @@ func main() {
 		})
 	}
 
-	llmSpan, err := factGenTask.LogPrompt(
+	llmSpan, err := traceloop.LogPrompt(
+		ctx,
 		tlp.Prompt{
 			Vendor:   "openai",
 			Mode:     "chat",
 			Model:    request.Model,
 			Messages: promptMsgs,
 		},
+		tlp.WorkflowAttributes{
+			Name: "example-workflow",
+		},
 	)
 	if err != nil {
 		fmt.Printf("LogPrompt error: %v\n", err)
@@ -88,30 +84,5 @@ func main() {
 		PromptTokens:     resp.Usage.PromptTokens,
 	})
 
-	someOtherTask := wf.NewTask("some_other_task")
-	defer someOtherTask.End()
-
-	otherPrompt, _ := someOtherTask.LogPrompt(tlp.Prompt{
-		Vendor: "openai",
-		Mode:   "chat",
-		Model:  request.Model,
-		Messages: []tlp.Message{
-			{
-				Index:   0,
-				Content: "some other prompt",
-				Role:    "user",
-			},
-		},
-	})
-
-	otherPrompt.LogCompletion(ctx, tlp.Completion{
-		Model:    resp.Model,
-		Messages: completionMsgs,
-	}, tlp.Usage{
-		TotalTokens:      resp.Usage.TotalTokens,
-		CompletionTokens: resp.Usage.CompletionTokens,
-		PromptTokens:     resp.Usage.PromptTokens,
-	})
-
 	fmt.Println(resp.Choices[0].Message.Content)
 }