diff --git a/cmd/worker/main.go b/cmd/worker/main.go index 0f96bb4..93ffad6 100644 --- a/cmd/worker/main.go +++ b/cmd/worker/main.go @@ -53,6 +53,7 @@ func main() { w.RegisterWorkflow(workflow.TestWorkflow) w.RegisterWorkflow(workflow.RoutingWorkflow) w.RegisterWorkflow(workflow.WorkflowGraphQuery) + w.RegisterWorkflow(workflow.LLMTestWorkflow) // Register all activities w.RegisterActivity(activity.CloneRepoActivity) @@ -72,6 +73,8 @@ func main() { // Routing workflow activities w.RegisterActivity(activity.LLMRouterActivity) + w.RegisterActivity(activity.LLMInferenceActivity) + w.RegisterActivity(activity.LLMBatchInferenceActivity) w.RegisterActivity(activity.ValidateWorkflowSpecActivity) w.RegisterActivity(activity.ValidateCronWorkflowSpecActivity) diff --git a/workflow/llm_test_workflow.go b/workflow/llm_test_workflow.go new file mode 100644 index 0000000..dd9ffe1 --- /dev/null +++ b/workflow/llm_test_workflow.go @@ -0,0 +1,35 @@ +package workflow + +import ( + "time" + "go.temporal.io/sdk/workflow" + "github.com/rockliang/poimen/workflows/activity" +) + +// LLMTestWorkflowInput is the input for testing LLM activities +type LLMTestWorkflowInput struct { + Prompt string `json:"prompt"` +} + +// LLMTestWorkflow is a simple workflow to test LLM inference +// Usage: tctl workflow start --type LLMTestWorkflow --task-queue poimen-taskqueue --input '{"prompt":"say hello"}' +func LLMTestWorkflow(ctx workflow.Context, input LLMTestWorkflowInput) (string, error) { + // Call the LLM inference activity + opts := workflow.ActivityOptions{ + StartToCloseTimeout: 60 * time.Second, + } + actCtx := workflow.WithActivityOptions(ctx, opts) + + actInput := activity.LLMInferenceInput{ + Model: "reasoning", + UserPrompt: input.Prompt, + } + + var result activity.LLMInferenceOutput + err := workflow.ExecuteActivity(actCtx, activity.LLMInferenceActivity, actInput).Get(actCtx, &result) + if err != nil { + return "", err + } + + return result.Response, nil +}