poimen-workflows/action/llm_inference.go

package action

import (
	"context"
	"fmt"

	"github.com/rockliang/poimen/workflows/action/llm"
	"github.com/rockliang/poimen/workflows/statemachine"
)

// LLMInferenceInput is input for LLMInferenceActivity
type LLMInferenceInput struct {
	Model         string                 `json:"model"`          // Model ID (reasoning, ornith:35b, etc)
	SystemPrompt  string                 `json:"system_prompt"`  // System instruction
	UserPrompt    string                 `json:"user_prompt"`    // User message
	Temperature   float64                `json:"temperature,omitempty"`   // LLM temperature (0-1)
	MaxTokens     int                    `json:"max_tokens,omitempty"`    // Max output tokens
	AuthToken     string                 `json:"auth_token,omitempty"`    // JWT token for authenticated endpoints
}

// LLMInferenceOutput is output from LLMInferenceActivity
type LLMInferenceOutput struct {
	Response     string `json:"response"`      // LLM response text
	Model        string `json:"model"`         // Model used
	StopReason   string `json:"stop_reason"`   // How inference stopped (stop_sequence, length, etc)
	TokensUsed   int    `json:"tokens_used"`   // Total tokens consumed
	ErrorMessage string `json:"error,omitempty"`
}

// LLMInferenceActivity calls LLM API with given prompt and returns response
func LLMInferenceActivity(ctx context.Context, in LLMInferenceInput) (LLMInferenceOutput, error) {
	logger := newActivityLogger(ctx)

	output := LLMInferenceOutput{
		Model: in.Model,
	}

	// Validate input
	if in.Model == "" {
		return output, fmt.Errorf("model not specified")
	}

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

	logger.logf("info", "Starting LLM inference with model: %s", in.Model)

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

	// Call LLM
	logger.logf("info", "Calling LLM API (model=%s, prompt_len=%d, auth=%v)", in.Model, len(in.UserPrompt), in.AuthToken != "")

	response, err := client.CreateMessage(ctx, llm.MessageInput{
		Model: statemachine.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.logf("error", "LLM API call failed: %v", err)
		return output, fmt.Errorf("LLM inference failed: %w", err)
	}

	output.Response = response
	output.StopReason = "stop_sequence"

	logger.logf("info", "LLM inference completed (response_len=%d)", len(response))

	return output, nil
}

// LLMBatchInferenceInput is input for batch inference
type LLMBatchInferenceInput struct {
	Model        string                   `json:"model"`
	SystemPrompt string                   `json:"system_prompt"`
	Prompts      []string                 `json:"prompts"`       // List of user prompts
	Temperature  float64                  `json:"temperature,omitempty"`
	AuthToken    string                   `json:"auth_token,omitempty"`    // JWT token for authenticated endpoints
}

// LLMBatchInferenceOutput is output from batch inference
type LLMBatchInferenceOutput struct {
	Responses []string `json:"responses"`     // LLM responses (parallel to input Prompts)
	Model     string   `json:"model"`
	Errors    []string `json:"errors,omitempty"`
}

// LLMBatchInferenceActivity calls LLM multiple times in sequence
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.logf("info", "Starting batch LLM inference (model=%s, count=%d)", in.Model, len(in.Prompts))

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

	// Process each prompt
	for i, prompt := range in.Prompts {
		logger.logf("info", "Processing prompt %d/%d", i+1, len(in.Prompts))

		response, err := client.CreateMessage(ctx, llm.MessageInput{
			Model: statemachine.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.logf("warn", "Failed to process prompt %d: %v", i, err)
		} else {
			output.Responses = append(output.Responses, response)
		}
	}

	logger.logf("info", "Batch inference completed (responses=%d, errors=%d)", 
		len(output.Responses), len(output.Errors))

	return output, nil
}