go-llm/function.go

package go_llm

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

	"github.com/sashabaranov/go-openai"
	"github.com/sashabaranov/go-openai/jsonschema"

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

type Function struct {
	Name        string      `json:"name"`
	Description string      `json:"description,omitempty"`
	Strict      bool        `json:"strict,omitempty"`
	Parameters  schema.Type `json:"parameters"`

	Forced bool `json:"forced,omitempty"`

	// Timeout is the maximum time to wait for the function to complete
	Timeout time.Duration `json:"-"`

	// fn is the function to call, only set if this is constructed with NewFunction
	fn reflect.Value

	paramType reflect.Type

	// definition is a cache of the openaiImpl jsonschema definition
	definition *jsonschema.Definition
}

func (f *Function) Execute(ctx *Context, input string) (any, error) {
	if !f.fn.IsValid() {
		return "", fmt.Errorf("function %s is not implemented", f.Name)
	}

	// first, we need to parse the input into the struct
	p := reflect.New(f.paramType)
	fmt.Println("Function.Execute", f.Name, "input:", input)
	//m := map[string]any{}
	err := json.Unmarshal([]byte(input), p.Interface())
	if err != nil {
		return "", fmt.Errorf("failed to unmarshal input: %w (input: %s)", err, input)
	}

	// now we can call the function
	exec := func(ctx *Context) (any, error) {
		out := f.fn.Call([]reflect.Value{reflect.ValueOf(ctx), p.Elem()})

		if len(out) != 2 {
			return "", fmt.Errorf("function %s must return two values, got %d", f.Name, len(out))
		}

		if out[1].IsNil() {
			return out[0].Interface(), nil
		}

		return "", out[1].Interface().(error)
	}

	var cancel context.CancelFunc
	if f.Timeout > 0 {
		ctx, cancel = ctx.WithTimeout(f.Timeout)
		defer cancel()
	}

	return exec(ctx)
}

func (f *Function) toOpenAIFunction() *openai.FunctionDefinition {
	return &openai.FunctionDefinition{
		Name:        f.Name,
		Description: f.Description,
		Strict:      f.Strict,
		Parameters:  f.Parameters,
	}
}
func (f *Function) toOpenAIDefinition() jsonschema.Definition {
	if f.definition == nil {
		def := f.Parameters.Definition()
		f.definition = &def
	}

	return *f.definition
}

type FunctionCall struct {
	Name      string `json:"name,omitempty"`
	Arguments string `json:"arguments,omitempty"`
}

func (fc *FunctionCall) toRaw() map[string]any {
	res := map[string]interface{}{
		"name": fc.Name,
	}

	if fc.Arguments != "" {
		res["arguments"] = fc.Arguments
	}

	return res
}