poimen-workflows/activity/llm_inference.go

package activity

import (
	"context"
	"fmt"

	"github.com/rockliang/poimen/workflows/activity/llm"
	"github.com/rockliang/poimen/workflows/pkg/types"
)

type LLMInferenceInput struct {
	Model        string  `json:"model"`
	SystemPrompt string  `json:"system_prompt"`
	UserPrompt   string  `json:"user_prompt"`
	Temperature  float64 `json:"temperature,omitempty"`
	MaxTokens    int     `json:"max_tokens,omitempty"`
	AuthToken    string  `json:"auth_token,omitempty"`
}

type LLMInferenceOutput struct {
	Response     string `json:"response"`
	Model        string `json:"model"`
	StopReason   string `json:"stop_reason"`
	TokensUsed   int    `json:"tokens_used"`
	ErrorMessage string `json:"error,omitempty"`
}

func LLMInferenceActivity(ctx context.Context, in LLMInferenceInput) (LLMInferenceOutput, error) {
	logger := newActivityLogger(ctx)
	output := LLMInferenceOutput{Model: in.Model}

	if in.Model == "" {
		return output, fmt.Errorf("model not specified")
	}
	if in.UserPrompt == "" {
		return output, fmt.Errorf("user_prompt not specified")
	}

	logger.Info("Starting LLM inference", "model", in.Model)

	client, err := llm.NewClient()
	if err != nil {
		output.ErrorMessage = err.Error()
		return output, fmt.Errorf("failed to create LLM client: %w", err)
	}

	response, err := client.CreateMessage(ctx, llm.MessageInput{
		Model:        types.ModelSpec{ModelID: in.Model},
		SystemPrompt: in.SystemPrompt,
		Messages:     []llm.MessageParam{{Role: "user", Content: in.UserPrompt}},
		AuthToken:    in.AuthToken,
	})
	if err != nil {
		output.ErrorMessage = err.Error()
		logger.Warn("LLM API call failed", "error", err)
		return output, fmt.Errorf("LLM inference failed: %w", err)
	}

	output.Response = response
	output.StopReason = "stop_sequence"
	logger.Info("LLM inference completed", "response_len", len(response))
	return output, nil
}

type LLMBatchInferenceInput struct {
	Model        string   `json:"model"`
	SystemPrompt string   `json:"system_prompt"`
	Prompts      []string `json:"prompts"`
	Temperature  float64  `json:"temperature,omitempty"`
	AuthToken    string   `json:"auth_token,omitempty"`
}

type LLMBatchInferenceOutput struct {
	Responses []string `json:"responses"`
	Model     string   `json:"model"`
	Errors    []string `json:"errors,omitempty"`
}

func LLMBatchInferenceActivity(ctx context.Context, in LLMBatchInferenceInput) (LLMBatchInferenceOutput, error) {
	logger := newActivityLogger(ctx)
	output := LLMBatchInferenceOutput{Model: in.Model, Responses: []string{}, Errors: []string{}}

	if in.Model == "" {
		return output, fmt.Errorf("model not specified")
	}
	if len(in.Prompts) == 0 {
		return output, fmt.Errorf("no prompts provided")
	}

	logger.Info("Starting batch inference", "model", in.Model, "count", len(in.Prompts))

	client, err := llm.NewClient()
	if err != nil {
		return output, fmt.Errorf("failed to create LLM client: %w", err)
	}

	for i, prompt := range in.Prompts {
		response, err := client.CreateMessage(ctx, llm.MessageInput{
			Model:        types.ModelSpec{ModelID: in.Model},
			SystemPrompt: in.SystemPrompt,
			Messages:     []llm.MessageParam{{Role: "user", Content: prompt}},
		})
		if err != nil {
			output.Errors = append(output.Errors, fmt.Sprintf("prompt %d: %v", i, err))
			output.Responses = append(output.Responses, "")
			logger.Warn("Failed prompt", "index", i, "error", err)
		} else {
			output.Responses = append(output.Responses, response)
		}
	}

	logger.Info("Batch inference completed", "responses", len(output.Responses), "errors", len(output.Errors))
	return output, nil
}