-
Notifications
You must be signed in to change notification settings - Fork 5
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: introduce workflows and tasks (#8)
- Loading branch information
1 parent
8dda467
commit 2d99d05
Showing
7 changed files
with
210 additions
and
13 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,117 @@ | ||
| package main | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
| "os" | ||
| "time" | ||
|
|
||
| "github.com/sashabaranov/go-openai" | ||
| tlp "github.com/traceloop/go-openllmetry/traceloop-sdk" | ||
| ) | ||
|
|
||
| func main() { | ||
| ctx := context.Background() | ||
|
|
||
| traceloop, err := tlp.NewClient(ctx, tlp.Config{ | ||
| BaseURL: "api-staging.traceloop.com", | ||
| APIKey: os.Getenv("TRACELOOP_API_KEY"), | ||
| }) | ||
| defer func() { traceloop.Shutdown(ctx) }() | ||
|
|
||
| if err != nil { | ||
| fmt.Printf("NewClient error: %v\n", err) | ||
| 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) | ||
| return | ||
| } | ||
|
|
||
| var promptMsgs []tlp.Message | ||
| for i, message := range request.Messages { | ||
| promptMsgs = append(promptMsgs, tlp.Message{ | ||
| Index: i, | ||
| Content: message.Content, | ||
| Role: message.Role, | ||
| }) | ||
| } | ||
|
|
||
| llmSpan, err := factGenTask.LogPrompt( | ||
| tlp.Prompt{ | ||
| Vendor: "openai", | ||
| Mode: "chat", | ||
| Model: request.Model, | ||
| Messages: promptMsgs, | ||
| }, | ||
| ) | ||
| if err != nil { | ||
| fmt.Printf("LogPrompt error: %v\n", err) | ||
| return | ||
| } | ||
|
|
||
| client := openai.NewClient(os.Getenv("OPENAI_API_KEY")) | ||
| resp, err := client.CreateChatCompletion( | ||
| context.Background(), | ||
| *request, | ||
| ) | ||
| if err != nil { | ||
| fmt.Printf("ChatCompletion error: %v\n", err) | ||
| return | ||
| } | ||
|
|
||
| var completionMsgs []tlp.Message | ||
| for _, choice := range resp.Choices { | ||
| completionMsgs = append(completionMsgs, tlp.Message{ | ||
| Index: choice.Index, | ||
| Content: choice.Message.Content, | ||
| Role: choice.Message.Role, | ||
| }) | ||
| } | ||
|
|
||
| llmSpan.LogCompletion(ctx, tlp.Completion{ | ||
| Model: resp.Model, | ||
| Messages: completionMsgs, | ||
| }, tlp.Usage{ | ||
| TotalTokens: resp.Usage.TotalTokens, | ||
| CompletionTokens: resp.Usage.CompletionTokens, | ||
| 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) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -2,4 +2,4 @@ package traceloop | |
|
|
||
| func Version() string { | ||
| return "0.0.2" | ||
| } | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,69 @@ | ||
| package traceloop | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
|
|
||
| semconvai "github.com/traceloop/go-openllmetry/semconv-ai" | ||
| "go.opentelemetry.io/otel/trace" | ||
| ) | ||
|
|
||
| type Workflow struct { | ||
| sdk *Traceloop | ||
| ctx context.Context | ||
| Attributes WorkflowAttributes `json:"workflow_attributes"` | ||
| } | ||
|
|
||
| type Task struct { | ||
| workflow *Workflow | ||
| ctx context.Context | ||
| Name string `json:"name"` | ||
| } | ||
|
|
||
| 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( | ||
| semconvai.TraceloopWorkflowName.String(attrs.Name), | ||
| semconvai.TraceloopSpanKind.String("workflow"), | ||
| semconvai.TraceloopEntityName.String(attrs.Name), | ||
| ) | ||
|
|
||
| return &Workflow{ | ||
| sdk: instance, | ||
| ctx: wCtx, | ||
| Attributes: attrs, | ||
| } | ||
| } | ||
|
|
||
| func (workflow *Workflow) End() { | ||
| trace.SpanFromContext(workflow.ctx).End() | ||
| } | ||
|
|
||
| func (workflow *Workflow) LogPrompt(prompt Prompt) (LLMSpan, error) { | ||
| return workflow.sdk.LogPrompt(workflow.ctx, prompt, workflow.Attributes) | ||
| } | ||
|
|
||
| func (workflow *Workflow) NewTask(name string) *Task { | ||
| tCtx, span := workflow.sdk.getTracer().Start(workflow.ctx, fmt.Sprintf("%s.task", name)) | ||
|
|
||
| span.SetAttributes( | ||
| semconvai.TraceloopWorkflowName.String(workflow.Attributes.Name), | ||
| semconvai.TraceloopSpanKind.String("task"), | ||
| semconvai.TraceloopEntityName.String(name), | ||
| ) | ||
|
|
||
| return &Task{ | ||
| workflow: workflow, | ||
| ctx: tCtx, | ||
| Name: name, | ||
| } | ||
| } | ||
|
|
||
| func (task *Task) End() { | ||
| trace.SpanFromContext(task.ctx).End() | ||
| } | ||
|
|
||
| func (task *Task) LogPrompt(prompt Prompt) (LLMSpan, error) { | ||
| return task.workflow.sdk.LogPrompt(task.ctx, prompt, task.workflow.Attributes) | ||
| } |