Skip to content

Commit 3f52b3f

Browse files
nina-kollmannina-kollman
committed
added context
1 parent 735f968 commit 3f52b3f

File tree

10 files changed

+105
-35
lines changed

10 files changed

+105
-35
lines changed

‎sample-app/generate_joke_workflow_example.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,9 @@ func translateJokeToPirate(ctx context.Context, workflow *sdk.Workflow, client *
8181
},
8282
}
8383

84-
agent := workflow.NewAgent("joke_translation")
84+
agent := workflow.NewAgent("joke_translation", map[string]string{
85+
"translation_type": "pirate",
86+
})
8587
defer agent.End()
8688

8789
llmSpan := agent.LogPrompt(prompt)

‎sample-app/main.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ func workflowExample() {
7171
}
7272

7373
// Log the prompt
74+
workflowName := "example-workflow"
7475
llmSpan := traceloop.LogPrompt(
7576
ctx,
7677
sdk.Prompt{
@@ -79,8 +80,8 @@ func workflowExample() {
7980
Model: request.Model,
8081
Messages: promptMsgs,
8182
},
82-
&sdk.WorkflowAttributes{
83-
Name: "example-workflow",
83+
sdk.ContextAttributes{
84+
WorkflowName: &workflowName,
8485
AssociationProperties: map[string]string{
8586
"user_id": "demo-user",
8687
},

‎sample-app/tool_calling.go‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -81,17 +81,18 @@ func runToolCallingExample() {
8181
Tools: tools,
8282
}
8383

84-
workflowAttrs := sdk.WorkflowAttributes{
85-
Name: "tool-calling-example",
84+
workflowName := "tool-calling-example"
85+
contextAttrs := sdk.ContextAttributes{
86+
WorkflowName: &workflowName,
8687
AssociationProperties: map[string]string{
8788
"user_id": "demo-user",
8889
},
8990
}
9091

9192
fmt.Printf("User: %s\n", userPrompt)
92-
93+
9394
// Log the prompt
94-
llmSpan := traceloop.LogPrompt(ctx, prompt, &workflowAttrs)
95+
llmSpan := traceloop.LogPrompt(ctx, prompt, contextAttrs)
9596

9697
// Make API call to OpenAI
9798
startTime := time.Now()

‎traceloop-sdk/agent.go‎

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,30 +13,45 @@ type Agent struct {
1313
sdk *Traceloop
1414
workflow *Workflow
1515
ctx context.Context
16-
Name string `json:"name"`
16+
Attributes AgentAttributes `json:"agent_attributes"`
1717
}
1818

1919
func (agent *Agent) End() {
2020
trace.SpanFromContext(agent.ctx).End()
2121
}
2222

2323
func (agent *Agent) LogPrompt(prompt Prompt) LLMSpan {
24+
// Merge workflow and agent association properties
25+
contextAttrs := ContextAttributes{
26+
AssociationProperties: make(map[string]string),
27+
}
28+
29+
// Start with workflow properties if available
2430
if agent.workflow != nil {
25-
return agent.sdk.LogPrompt(agent.ctx, prompt, &agent.workflow.Attributes)
31+
contextAttrs.WorkflowName = &agent.workflow.Attributes.Name
32+
for key, value := range agent.workflow.Attributes.AssociationProperties {
33+
contextAttrs.AssociationProperties[key] = value
34+
}
2635
}
27-
return agent.sdk.LogPrompt(agent.ctx, prompt, nil)
36+
37+
// Agent properties override workflow properties
38+
for key, value := range agent.Attributes.AssociationProperties {
39+
contextAttrs.AssociationProperties[key] = value
40+
}
41+
42+
return agent.sdk.LogPrompt(agent.ctx, prompt, contextAttrs)
2843
}
2944

3045
func (agent *Agent) NewTool(name string, toolType string, toolFunction ToolFunction) *Tool {
3146
toolCtx, span := agent.sdk.getTracer().Start(agent.ctx, fmt.Sprintf("%s.tool", name))
3247
span.SetAttributes(
33-
semconvai.LLMAgentName.String(agent.Name),
48+
semconvai.LLMAgentName.String(agent.Attributes.Name),
3449
semconvai.TraceloopSpanKind.String(string(model.SpanKindTool)),
3550
semconvai.TraceloopEntityName.String(name),
3651
)
3752

3853
return &Tool{
39-
agent: agent,
54+
agent: *agent,
4055
ctx: toolCtx,
4156
Name: name,
4257
Type: toolType,

‎traceloop-sdk/sdk.go‎

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -163,12 +163,15 @@ func (instance *Traceloop) NewAgent(ctx context.Context, name string, associatio
163163
sdk: instance,
164164
workflow: nil,
165165
ctx: aCtx,
166-
Name: name,
166+
Attributes: AgentAttributes{
167+
Name: name,
168+
AssociationProperties: associationProperties,
169+
},
167170
}
168171
}
169172

170173
// New workflow-based API
171-
func (instance *Traceloop) LogPrompt(ctx context.Context, prompt Prompt, workflowAttrs *WorkflowAttributes) LLMSpan {
174+
func (instance *Traceloop) LogPrompt(ctx context.Context, prompt Prompt, contextAttrs ContextAttributes) LLMSpan {
172175
spanName := fmt.Sprintf("%s.%s", prompt.Vendor, prompt.Mode)
173176
_, span := instance.getTracer().Start(ctx, spanName)
174177

@@ -178,14 +181,19 @@ func (instance *Traceloop) LogPrompt(ctx context.Context, prompt Prompt, workflo
178181
semconvai.LLMRequestType.String(prompt.Mode),
179182
}
180183

181-
if workflowAttrs != nil {
182-
attrs = append(attrs, semconvai.TraceloopWorkflowName.String(workflowAttrs.Name))
184+
if contextAttrs.WorkflowName != nil {
185+
attrs = append(attrs, semconvai.TraceloopWorkflowName.String(*contextAttrs.WorkflowName))
186+
}
183187

184-
// Add association properties if provided
185-
for key, value := range workflowAttrs.AssociationProperties {
186-
attrs = append(attrs, attribute.String("traceloop.association.properties."+key, value))
187-
}
188+
if contextAttrs.AgentName != nil {
189+
attrs = append(attrs, semconvai.LLMAgentName.String(*contextAttrs.AgentName))
190+
}
191+
192+
// Add association properties
193+
for key, value := range contextAttrs.AssociationProperties {
194+
attrs = append(attrs, attribute.String("traceloop.association.properties."+key, value))
188195
}
196+
189197

190198
span.SetAttributes(attrs...)
191199
setMessagesAttribute(span, "llm.prompts", prompt.Messages)

‎traceloop-sdk/sdk_test.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,10 @@ func TestLogPromptSpanAttributes(t *testing.T) {
7070
}
7171

7272
// Log the prompt using new workflow API
73-
llmSpan := tl.LogPrompt(context.Background(), prompt, workflowAttrs)
73+
llmSpan := tl.LogPrompt(context.Background(), prompt, ContextAttributes{
74+
WorkflowName: &workflowAttrs.Name,
75+
AssociationProperties: workflowAttrs.AssociationProperties,
76+
})
7477

7578
// Log completion with tool calls
7679
completion := Completion{

‎traceloop-sdk/task.go‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,5 +17,9 @@ func (task *Task) End() {
1717
}
1818

1919
func (task *Task) LogPrompt(prompt Prompt) LLMSpan {
20-
return task.workflow.sdk.LogPrompt(task.ctx, prompt, &task.workflow.Attributes)
20+
contextAttrs := ContextAttributes{
21+
WorkflowName: &task.workflow.Attributes.Name,
22+
AssociationProperties: task.workflow.Attributes.AssociationProperties,
23+
}
24+
return task.workflow.sdk.LogPrompt(task.ctx, prompt, contextAttrs)
2125
}

‎traceloop-sdk/tool.go‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import (
77
)
88

99
type Tool struct {
10-
agent *Agent
10+
agent Agent
1111
ctx context.Context
1212
Name string `json:"name"`
1313
Type string `json:"type"`
@@ -19,8 +19,23 @@ func (tool *Tool) End() {
1919
}
2020

2121
func (tool *Tool) LogPrompt(prompt Prompt) LLMSpan {
22+
// Merge workflow and agent association properties
23+
contextAttrs := ContextAttributes{
24+
AssociationProperties: make(map[string]string),
25+
}
26+
27+
// Start with workflow properties if available
2228
if tool.agent.workflow != nil {
23-
return tool.agent.sdk.LogPrompt(tool.ctx, prompt, &tool.agent.workflow.Attributes)
29+
contextAttrs.WorkflowName = &tool.agent.workflow.Attributes.Name
30+
for key, value := range tool.agent.workflow.Attributes.AssociationProperties {
31+
contextAttrs.AssociationProperties[key] = value
32+
}
2433
}
25-
return tool.agent.sdk.LogPrompt(tool.ctx, prompt, nil)
34+
35+
// Agent properties override workflow properties
36+
for key, value := range tool.agent.Attributes.AssociationProperties {
37+
contextAttrs.AssociationProperties[key] = value
38+
}
39+
40+
return tool.agent.sdk.LogPrompt(tool.ctx, prompt, contextAttrs)
2641
}

‎traceloop-sdk/tracing_types.go‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,12 @@ type WorkflowAttributes struct {
3030
AssociationProperties map[string]string `json:"association_properties"`
3131
}
3232

33+
type ContextAttributes struct {
34+
WorkflowName *string `json:"workflow_name,omitempty"`
35+
AgentName *string `json:"agent_name,omitempty"`
36+
AssociationProperties map[string]string `json:"association_properties,omitempty"`
37+
}
38+
3339
type Usage struct {
3440
TotalTokens int `json:"total_tokens"`
3541
CompletionTokens int `json:"completion_tokens"`
@@ -42,7 +48,6 @@ type ToolFunction struct {
4248
Parameters interface{} `json:"parameters"`
4349
}
4450

45-
4651
type ToolCall struct {
4752
ID string `json:"id"`
4853
Type string `json:"type"`
@@ -59,5 +64,6 @@ type ToolCallAttributes struct {
5964
}
6065

6166
type AgentAttributes struct {
62-
Name string `json:"name"`
67+
Name string `json:"agent_name"`
68+
AssociationProperties map[string]string `json:"association_properties"`
6369
}

‎traceloop-sdk/workflow.go‎

Lines changed: 23 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66

77
"github.com/traceloop/go-openllmetry/traceloop-sdk/model"
88
semconvai "github.com/traceloop/go-openllmetry/semconv-ai"
9+
"go.opentelemetry.io/otel/attribute"
910
"go.opentelemetry.io/otel/trace"
1011
)
1112

@@ -36,7 +37,11 @@ func (workflow *Workflow) End() {
3637
}
3738

3839
func (workflow *Workflow) LogPrompt(prompt Prompt) LLMSpan {
39-
return workflow.sdk.LogPrompt(workflow.ctx, prompt, &workflow.Attributes)
40+
contextAttrs := ContextAttributes{
41+
WorkflowName: &workflow.Attributes.Name,
42+
AssociationProperties: workflow.Attributes.AssociationProperties,
43+
}
44+
return workflow.sdk.LogPrompt(workflow.ctx, prompt, contextAttrs)
4045
}
4146

4247
func (workflow *Workflow) NewTask(name string) *Task {
@@ -55,20 +60,30 @@ func (workflow *Workflow) NewTask(name string) *Task {
5560
}
5661
}
5762

58-
func (workflow *Workflow) NewAgent(name string) *Agent {
63+
func (workflow *Workflow) NewAgent(name string, associationProperties map[string]string) *Agent {
5964
aCtx, span := workflow.sdk.getTracer().Start(workflow.ctx, fmt.Sprintf("%s.agent", name))
6065

61-
span.SetAttributes(
66+
attrs := []attribute.KeyValue{
6267
semconvai.TraceloopWorkflowName.String(workflow.Attributes.Name),
6368
semconvai.TraceloopSpanKind.String(string(model.SpanKindAgent)),
6469
semconvai.TraceloopEntityName.String(name),
65-
)
70+
}
71+
72+
// Add agent-specific association properties to the span
73+
for key, value := range associationProperties {
74+
attrs = append(attrs, attribute.String("traceloop.association.properties."+key, value))
75+
}
76+
77+
span.SetAttributes(attrs...)
6678

6779
return &Agent{
68-
sdk: workflow.sdk,
69-
workflow: workflow,
70-
ctx: aCtx,
71-
Name: name,
80+
sdk: workflow.sdk,
81+
workflow: workflow,
82+
ctx: aCtx,
83+
Attributes: AgentAttributes{
84+
Name: name,
85+
AssociationProperties: associationProperties,
86+
},
7287
}
7388
}
7489

0 commit comments

Comments
 (0)