answer/pkg/answer/answer.go

package answer

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net/url"
	"strings"

	"gitea.stevedudenhoeffer.com/steve/answer/pkg/cache"
	"gitea.stevedudenhoeffer.com/steve/answer/pkg/extractor"
	"gitea.stevedudenhoeffer.com/steve/answer/pkg/search"
	gollm "gitea.stevedudenhoeffer.com/steve/go-llm"
)

var ErrMaxTries = errors.New("maximum number of pages tried reached")
var ErrMaxAnswers = errors.New("maximum number of answers parsed reached")
var ErrTooManyArguments = errors.New("too many arguments")

type Question struct {
	// Question is the question to answer
	Question string

	Model gollm.ChatCompletion

	Search search.Search

	Cache cache.Cache
}

// Answers is a list of answers to a question
type Answers []string

type Options struct {
	// MaxSearches is the maximum possible number of searches to execute for this question.  If this is set to 5, the function could
	// search up to 5 possible times to find an answer.
	MaxSearches int

	// MaxThinks is the maximum number of times to think about a question.  A "Think" is different than a search in that
	// the LLM just breaks the question down into smaller parts and tries to answer them.  This is useful for complex
	// questions that are hard to answer since LLMs are better at answering smaller questions.
	MaxThinks int

	// MaxTries is the absolute maximum number of pages to try to get an answer from.  For instance, if MaxSearches is 5 and
	// 5 pages are tried and no answers are found, the function will return ErrMaxTries.
	MaxTries int

	// OnNewFunction is a callback that, if non-nil, will be called when a new function is called by the LLM.
	// The "answer" and "no_answer" functions are not included in this callback.
	// Return an error to stop the function from being called.
	OnNewFunction func(ctx context.Context, funcName string, question string, parameter string) error
}

var DefaultOptions = Options{
	MaxSearches: 5,
	MaxThinks:   10,
	MaxTries:    5,
}

type Result struct {
	Result string
	Error  error
}

func fanExecuteToolCalls(ctx context.Context, toolBox *gollm.ToolBox, calls []gollm.ToolCall) []Result {
	var results []Result
	var resultsOutput = make(chan Result, len(calls))

	fnCall := func(call gollm.ToolCall) Result {
		str, err := toolBox.Execute(ctx, call)
		if err != nil {
			return Result{
				Error: err,
			}
		}

		return Result{
			Result: str,
		}
	}

	for _, call := range calls {
		go func(call gollm.ToolCall) {
			resultsOutput <- fnCall(call)
		}(call)
	}

	for i := 0; i < len(calls); i++ {
		result := <-resultsOutput
		results = append(results, result)
	}

	close(resultsOutput)

	return results
}

type article struct {
	URL   string
	Title string
	Body  string
}

func extractArticle(ctx context.Context, c cache.Cache, u *url.URL) (res article, err error) {
	defer func() {
		e := recover()

		if e != nil {
			if e, ok := e.(error); ok {
				err = fmt.Errorf("panic: %w", e)
			} else {
				err = fmt.Errorf("panic: %v", e)
			}
		}
	}()

	extractors := extractor.MultiExtractor(
		extractor.CacheExtractor{
			Cache:     c,
			Tag:       "goose",
			Extractor: extractor.GooseExtractor{},
		},
		extractor.CacheExtractor{
			Cache:     c,
			Tag:       "playwright",
			Extractor: extractor.PlaywrightExtractor{},
		},
	)

	a, err := extractors.Extract(ctx, u.String())

	if err != nil {
		return article{
			URL:   "",
			Title: "",
			Body:  "",
		}, err
	}

	return article{
		URL:   a.URL,
		Title: a.Title,
		Body:  a.Body,
	}, nil
}

func doesTextAnswerQuestion(ctx context.Context, q Question, text string) (string, error) {
	fnAnswer := gollm.NewFunction(
		"answer",
		"The answer from the given text that answers the question.",
		func(ctx context.Context, args struct {
			Answer string `description:"the answer to the question, the answer should come from the text"`
		}) (string, error) {
			return args.Answer, nil
		})

	fnNoAnswer := gollm.NewFunction(
		"no_answer",
		"Indicate that the text does not answer the question.",
		func(ctx context.Context, args struct {
			Ignored string `description:"ignored, just here to make sure the function is called.  Fill with anything."`
		}) (string, error) {
			return "", nil
		})

	req := gollm.Request{
		Messages: []gollm.Message{
			{
				Role: gollm.RoleSystem,
				Text: "Evaluate the given text to see if it answers the question from the user.  The text is as follows:",
			},
			{
				Role: gollm.RoleSystem,
				Text: text,
			},
			{
				Role: gollm.RoleUser,
				Text: q.Question,
			},
		},
		Toolbox: gollm.NewToolBox(fnAnswer, fnNoAnswer),
	}

	res, err := q.Model.ChatComplete(ctx, req)

	if err != nil {
		return "", err
	}

	if len(res.Choices) == 0 {
		return "", nil
	}

	if len(res.Choices[0].Calls) == 0 {
		return "", nil
	}

	return req.Toolbox.Execute(ctx, res.Choices[0].Calls[0])
}

func functionSearch(ctx context.Context, q Question, searchTerm string) (string, error) {

	slog.Info("searching", "search", searchTerm, "question", q)
	res, err := q.Search.Search(ctx, searchTerm)
	if err != nil {
		return "", err
	}

	if len(res) == 0 {
		return "", nil
	}

	// first pass try to see if any provide the result without needing archive
	for _, r := range res {
		trimmed := strings.TrimSpace(r.URL)
		if trimmed == "" {

		}

		slog.Info("extracting article", "url", trimmed)
		u, err := url.Parse(trimmed)
		if err != nil {
			continue
		}

		a, err := extractArticle(ctx, q.Cache, u)

		if err != nil {
			continue
		}

		slog.Info("extracted article", "url", a.URL, "title", a.Title, "body", a.Body)

		if a.Title != "" && a.Body != "" {
			answer, err := doesTextAnswerQuestion(ctx, q, a.Body)

			if err != nil {
				slog.Error("error checking if text answers question", "question", q.Question, "error", err)
				continue
			}

			if answer != "" {
				return answer, nil
			}
		}
	}

	return "", nil
}

func functionThink(ctx context.Context, q Question) (string, error) {
	fnAnswer := gollm.NewFunction(
		"answer",
		"Answer the question.",
		func(ctx context.Context, args struct {
			Answer string `description:"the answer to the question"`
		}) (string, error) {
			return args.Answer, nil
		})

	var temp float32 = 0.8
	req := gollm.Request{
		Messages: []gollm.Message{
			{
				Role: gollm.RoleSystem,
				Text: "Answer the given question as accurately and concisely as possible using the answer function.",
			},
			{
				Role: gollm.RoleUser,
				Text: q.Question,
			},
		},
		Toolbox:     gollm.NewToolBox(fnAnswer),
		Temperature: &temp,
	}

	res, err := q.Model.ChatComplete(ctx, req)

	if err != nil {
		return "", err
	}

	if len(res.Choices) == 0 {
		return "", nil
	}

	if len(res.Choices[0].Calls) == 0 {
		return "", nil
	}

	return req.Toolbox.Execute(ctx, res.Choices[0].Calls[0])
}

func (o Options) Answer(ctx context.Context, q Question) (Answers, error) {
	fnSearch := gollm.NewFunction(
		"search",
		"Search the web for an answer to a question.  You can call this function up to "+fmt.Sprint(o.MaxSearches)+" times.",
		func(ctx context.Context, args struct {
			SearchQuery string `description:"what to search the web for for this question"`
			Question    string `description:"what question(s) you are trying to answer with this search"`
		}) (string, error) {
			q2 := q
			q2.Question = args.Question

			return functionSearch(ctx, q2, args.SearchQuery)
		})

	fnThink := gollm.NewFunction(
		"think",
		"Think about a question.  This is useful for breaking down complex questions into smaller parts that are easier to answer.",
		func(ctx context.Context, args struct {
			Question string `json:"question" description:"the question to think about"`
		}) (string, error) {
			q2 := q
			q2.Question = args.Question

			return functionThink(ctx, q2)
		})

	fnAnswer := gollm.NewFunction(
		"answer",
		"You definitively answer a question, if you call this it means you know the answer and do not need to search for it or use any other function to find it",
		func(ctx context.Context, args struct {
			Answer string `json:"answer" description:"the answer to the question"`
		}) (string, error) {
			return args.Answer, nil
		})

	var funcs = []*gollm.Function{fnAnswer}

	if o.MaxSearches > 0 {
		funcs = append(funcs, fnSearch)
	}

	if o.MaxThinks > 0 {
		funcs = append(funcs, fnThink)
	}

	var temp float32 = 0.8

	req := gollm.Request{
		Messages: []gollm.Message{
			{
				Role: gollm.RoleSystem,
				Text: "You are being asked to answer a question.  You must respond with a function.  You can answer it if you know the answer, or if some functions exist you can use those to help you find the answer.",
			},
			{
				Role: gollm.RoleUser,
				Text: q.Question,
			},
		},
		Toolbox:     gollm.NewToolBox(funcs...),
		Temperature: &temp,
	}

	res, err := q.Model.ChatComplete(ctx, req)

	if err != nil {
		return nil, err
	}

	if len(res.Choices) == 0 {
		return nil, nil
	}

	if len(res.Choices) > o.MaxSearches {
		res.Choices = res.Choices[:o.MaxSearches]
	}

	var answers []Result
	for _, choice := range res.Choices {
		fnChoice := func(choice gollm.ResponseChoice) []Result {
			var calls []Result
			var callsOutput = make(chan Result, len(choice.Calls))
			fnCall := func(call gollm.ToolCall) Result {
				str, err := req.Toolbox.Execute(ctx, call)

				if err != nil {
					return Result{
						Error: err,
					}
				}

				return Result{
					Result: str,
				}
			}

			for _, call := range choice.Calls {
				go func(call gollm.ToolCall) {
					if o.OnNewFunction != nil {
						err := o.OnNewFunction(ctx, call.FunctionCall.Name, q.Question, call.FunctionCall.Arguments)
						if err != nil {
							callsOutput <- Result{
								Error: err,
							}
							return
						}
					}
					callsOutput <- fnCall(call)
				}(call)
			}

			for i := 0; i < len(choice.Calls); i++ {
				result := <-callsOutput
				calls = append(calls, result)
			}

			close(callsOutput)

			return calls
		}

		answers = append(answers, fnChoice(choice)...)
	}

	var errs []error
	var results []string

	for _, answer := range answers {
		if answer.Error != nil {
			errs = append(errs, answer.Error)
			continue
		}

		results = append(results, answer.Result)
	}

	if len(errs) > 0 {
		return nil, errors.Join(errs...)
	}

	return results, nil

}

func Answer(ctx context.Context, q Question) (Answers, error) {
	return DefaultOptions.Answer(ctx, q)
}