go-llm/v2/generate.go

package llm

import (
	"context"
	"encoding/json"
	"fmt"

	"gitea.stevedudenhoeffer.com/steve/go-llm/v2/internal/schema"
)

const structuredOutputToolName = "structured_output"

// Generate sends a single user prompt to the model and parses the response into T.
// T must be a struct. The model is forced to return structured output matching T's schema
// by using a hidden tool call internally.
func Generate[T any](ctx context.Context, model *Model, prompt string, opts ...RequestOption) (T, error) {
	return GenerateWith[T](ctx, model, []Message{UserMessage(prompt)}, opts...)
}

// GenerateWith sends the given messages to the model and parses the response into T.
// T must be a struct. The model is forced to return structured output matching T's schema
// by using a hidden tool call internally.
func GenerateWith[T any](ctx context.Context, model *Model, messages []Message, opts ...RequestOption) (T, error) {
	var zero T

	s := schema.FromStruct(zero)

	tool := Tool{
		Name:        structuredOutputToolName,
		Description: "Return your response as structured data using this function. You MUST call this function with your response.",
		Schema:      s,
	}

	// Append WithTools as the last option so it overrides any user-provided tools.
	opts = append(opts, WithTools(NewToolBox(tool)))

	resp, err := model.Complete(ctx, messages, opts...)
	if err != nil {
		return zero, err
	}

	// Find the structured_output tool call in the response.
	for _, tc := range resp.ToolCalls {
		if tc.Name == structuredOutputToolName {
			var result T
			if err := json.Unmarshal([]byte(tc.Arguments), &result); err != nil {
				return zero, fmt.Errorf("failed to parse structured output: %w", err)
			}
			return result, nil
		}
	}

	return zero, ErrNoStructuredOutput
}