Skip to content

feat: add trace_id and ToolContext to workflow types #103

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
saurabhjain1592 wants to merge 1 commit into from
Closed

feat: add trace_id and ToolContext to workflow types #103

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/main/java/com/getaxonflow/sdk/AxonFlow.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4303,6 +4303,9 @@ public com.getaxonflow.sdk.types.workflow.WorkflowTypes.ListWorkflowsResponse li
if (options.getOffset() > 0) {
appendQueryParam(query, "offset", String.valueOf(options.getOffset()));
}
if (options.getTraceId() != null) {
appendQueryParam(query, "trace_id", options.getTraceId());
}
}

if (query.length() > 0) {
Expand Down
131 changes: 123 additions & 8 deletions src/main/java/com/getaxonflow/sdk/types/workflow/WorkflowTypes.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@

import java.time.Instant;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
Expand Down Expand Up @@ -230,16 +231,21 @@ public static final class CreateWorkflowRequest {
@JsonProperty("metadata")
private final Map<String, Object> metadata;

@JsonProperty("trace_id")
private final String traceId;

@JsonCreator
public CreateWorkflowRequest(
@JsonProperty("workflow_name") String workflowName,
@JsonProperty("source") WorkflowSource source,
@JsonProperty("total_steps") Integer totalSteps,
@JsonProperty("metadata") Map<String, Object> metadata) {
@JsonProperty("metadata") Map<String, Object> metadata,
@JsonProperty("trace_id") String traceId) {
this.workflowName = Objects.requireNonNull(workflowName, "workflowName is required");
this.source = source != null ? source : WorkflowSource.EXTERNAL;
this.totalSteps = totalSteps;
this.metadata = metadata != null ? Collections.unmodifiableMap(metadata) : Collections.emptyMap();
this.traceId = traceId;
}

public String getWorkflowName() {
Expand All @@ -258,6 +264,10 @@ public Map<String, Object> getMetadata() {
return metadata;
}

public String getTraceId() {
return traceId;
}

public static Builder builder() {
return new Builder();
}
Expand All @@ -267,6 +277,7 @@ public static final class Builder {
private WorkflowSource source = WorkflowSource.EXTERNAL;
private Integer totalSteps;
private Map<String, Object> metadata;
private String traceId;

public Builder workflowName(String workflowName) {
this.workflowName = workflowName;
Expand All @@ -288,8 +299,13 @@ public Builder metadata(Map<String, Object> metadata) {
return this;
}

public Builder traceId(String traceId) {
this.traceId = traceId;
return this;
}

public CreateWorkflowRequest build() {
return new CreateWorkflowRequest(workflowName, source, totalSteps, metadata);
return new CreateWorkflowRequest(workflowName, source, totalSteps, metadata, traceId);
}
}
}
Expand All @@ -315,18 +331,23 @@ public static final class CreateWorkflowResponse {
@JsonProperty("created_at")
private final Instant createdAt;

@JsonProperty("trace_id")
private final String traceId;

@JsonCreator
public CreateWorkflowResponse(
@JsonProperty("workflow_id") String workflowId,
@JsonProperty("workflow_name") String workflowName,
@JsonProperty("source") WorkflowSource source,
@JsonProperty("status") WorkflowStatus status,
@JsonProperty("created_at") Instant createdAt) {
@JsonProperty("created_at") Instant createdAt,
@JsonProperty("trace_id") String traceId) {
this.workflowId = workflowId;
this.workflowName = workflowName;
this.source = source;
this.status = status;
this.createdAt = createdAt;
this.traceId = traceId;
}

public String getWorkflowId() {
Expand All @@ -348,6 +369,64 @@ public WorkflowStatus getStatus() {
public Instant getCreatedAt() {
return createdAt;
}

public String getTraceId() {
return traceId;
}
}

/**
* Tool-level context for per-tool governance within tool_call steps.
*/
@JsonIgnoreProperties(ignoreUnknown = true)
public static final class ToolContext {

@JsonProperty("tool_name")
private final String toolName;

@JsonProperty("tool_type")
private final String toolType;

@JsonProperty("tool_input")
private final Map<String, Object> toolInput;

private ToolContext(Builder builder) {
this.toolName = builder.toolName;
this.toolType = builder.toolType;
this.toolInput = builder.toolInput != null ? Collections.unmodifiableMap(new HashMap<>(builder.toolInput)) : null;
}

@JsonCreator
public ToolContext(
@JsonProperty("tool_name") String toolName,
@JsonProperty("tool_type") String toolType,
@JsonProperty("tool_input") Map<String, Object> toolInput) {
this.toolName = toolName;
this.toolType = toolType;
this.toolInput = toolInput != null ? Collections.unmodifiableMap(new HashMap<>(toolInput)) : null;
}

public String getToolName() { return toolName; }
public String getToolType() { return toolType; }
public Map<String, Object> getToolInput() { return toolInput; }

public static Builder builder(String toolName) {
return new Builder(toolName);
}

public static final class Builder {
private final String toolName;
private String toolType;
private Map<String, Object> toolInput;

public Builder(String toolName) {
this.toolName = Objects.requireNonNull(toolName, "toolName must not be null");
}

public Builder toolType(String toolType) { this.toolType = toolType; return this; }
public Builder toolInput(Map<String, Object> toolInput) { this.toolInput = toolInput; return this; }
public ToolContext build() { return new ToolContext(this); }
}
}

/**
Expand All @@ -371,18 +450,23 @@ public static final class StepGateRequest {
@JsonProperty("provider")
private final String provider;

@JsonProperty("tool_context")
private final ToolContext toolContext;

@JsonCreator
public StepGateRequest(
@JsonProperty("step_name") String stepName,
@JsonProperty("step_type") StepType stepType,
@JsonProperty("step_input") Map<String, Object> stepInput,
@JsonProperty("model") String model,
@JsonProperty("provider") String provider) {
@JsonProperty("provider") String provider,
@JsonProperty("tool_context") ToolContext toolContext) {
this.stepName = stepName;
this.stepType = Objects.requireNonNull(stepType, "stepType is required");
this.stepInput = stepInput != null ? Collections.unmodifiableMap(stepInput) : Collections.emptyMap();
this.model = model;
this.provider = provider;
this.toolContext = toolContext;
}

public String getStepName() {
Expand All @@ -405,6 +489,10 @@ public String getProvider() {
return provider;
}

public ToolContext getToolContext() {
return toolContext;
}

public static Builder builder() {
return new Builder();
}
Expand All @@ -415,6 +503,7 @@ public static final class Builder {
private Map<String, Object> stepInput;
private String model;
private String provider;
private ToolContext toolContext;

public Builder stepName(String stepName) {
this.stepName = stepName;
Expand All @@ -441,8 +530,13 @@ public Builder provider(String provider) {
return this;
}

public Builder toolContext(ToolContext toolContext) {
this.toolContext = toolContext;
return this;
}

public StepGateRequest build() {
return new StepGateRequest(stepName, stepType, stepInput, model, provider);
return new StepGateRequest(stepName, stepType, stepInput, model, provider, toolContext);
}
}
}
Expand Down Expand Up @@ -679,6 +773,9 @@ public static final class WorkflowStatusResponse {
@JsonProperty("steps")
private final List<WorkflowStepInfo> steps;

@JsonProperty("trace_id")
private final String traceId;

@JsonCreator
public WorkflowStatusResponse(
@JsonProperty("workflow_id") String workflowId,
Expand All @@ -689,7 +786,8 @@ public WorkflowStatusResponse(
@JsonProperty("total_steps") Integer totalSteps,
@JsonProperty("started_at") Instant startedAt,
@JsonProperty("completed_at") Instant completedAt,
@JsonProperty("steps") List<WorkflowStepInfo> steps) {
@JsonProperty("steps") List<WorkflowStepInfo> steps,
@JsonProperty("trace_id") String traceId) {
this.workflowId = workflowId;
this.workflowName = workflowName;
this.source = source;
Expand All @@ -699,6 +797,7 @@ public WorkflowStatusResponse(
this.startedAt = startedAt;
this.completedAt = completedAt;
this.steps = steps != null ? Collections.unmodifiableList(steps) : Collections.emptyList();
this.traceId = traceId;
}

public String getWorkflowId() {
Expand Down Expand Up @@ -737,6 +836,10 @@ public List<WorkflowStepInfo> getSteps() {
return steps;
}

public String getTraceId() {
return traceId;
}

public boolean isTerminal() {
return status == WorkflowStatus.COMPLETED ||
status == WorkflowStatus.ABORTED ||
Expand All @@ -753,12 +856,14 @@ public static final class ListWorkflowsOptions {
private final WorkflowSource source;
private final int limit;
private final int offset;
private final String traceId;

public ListWorkflowsOptions(WorkflowStatus status, WorkflowSource source, int limit, int offset) {
public ListWorkflowsOptions(WorkflowStatus status, WorkflowSource source, int limit, int offset, String traceId) {
this.status = status;
this.source = source;
this.limit = limit > 0 ? limit : 50;
this.offset = Math.max(offset, 0);
this.traceId = traceId;
}

public WorkflowStatus getStatus() {
Expand All @@ -777,6 +882,10 @@ public int getOffset() {
return offset;
}

public String getTraceId() {
return traceId;
}

public static Builder builder() {
return new Builder();
}
Expand All @@ -786,6 +895,7 @@ public static final class Builder {
private WorkflowSource source;
private int limit = 50;
private int offset = 0;
private String traceId;

public Builder status(WorkflowStatus status) {
this.status = status;
Expand All @@ -807,8 +917,13 @@ public Builder offset(int offset) {
return this;
}

public Builder traceId(String traceId) {
this.traceId = traceId;
return this;
}

public ListWorkflowsOptions build() {
return new ListWorkflowsOptions(status, source, limit, offset);
return new ListWorkflowsOptions(status, source, limit, offset, traceId);
}
}
}
Expand Down
Loading