git.giacomolli.uk - onechat/blob - chain.go

   1 package main
   2
   3 import (
   4         "context"
   5         "encoding/json"
   6         "errors"
   7         "fmt"
   8         "strings"
   9 )
  10
  11 // errChainLogWrite marks an error of a log write, which ends the chain.
  12 var errChainLogWrite = errors.New("log write failed")
  13
  14 // ChainClient sends model calls.
  15 type ChainClient interface {
  16         Create(ctx context.Context, request ModelRequest) (ModelResponse, error)
  17 }
  18
  19 // Chain is one request chain of a turn or of a subagent. It sends request with
  20 // the input of the first request, runs the tool calls of each response, and
  21 // sends their results back, until a response has no tool call: the final
  22 // answer. The replies, the tool calls, and the results go to the log of tools
  23 // as they happen. The thoughts never do.
  24 //
  25 // Each later request begins with the whole input of the request before it,
  26 // then its output items, then the tool results. The output items hold the
  27 // reasoning with its encrypted content, because the calls do not store it. So
  28 // each request begins with all of the previous one and reads it from the prompt
  29 // cache. The messages of take join each later request after the tool results.
  30 type Chain struct {
  31         client  ChainClient
  32         tools   *Toolbox
  33         request ModelRequest
  34         take    func() []Message
  35 }
  36
  37 // Run runs the chain from input, the input of its first request, and returns
  38 // the text of the final answer. An error of a log write wraps
  39 // errChainLogWrite.
  40 func (c *Chain) Run(ctx context.Context, input []any) (string, error) {
  41         request := c.request
  42         for request.Turn = "first"; ; request.Turn = "next" {
  43                 request.Input = input
  44                 response, err := c.client.Create(ctx, request)
  45                 if err != nil {
  46                         return "", fmt.Errorf("create model response: %w", err)
  47                 }
  48
  49                 var text string
  50
  51                 input, text, err = c.handle(ctx, input, response)
  52                 if err != nil {
  53                         return "", err
  54                 }
  55
  56                 if len(input) == 0 {
  57                         return text, nil
  58                 }
  59         }
  60 }
  61
  62 // handle logs the replies of response and runs its tool calls. It returns the
  63 // input of the next request: input, the output items, the tool results, and
  64 // the messages of take. After the final answer it returns no input, and the
  65 // replies of the answer as its text.
  66 func (c *Chain) handle(ctx context.Context, input []any, response ModelResponse) ([]any, string, error) {
  67         var (
  68                 outputs []any
  69                 replies []string
  70         )
  71
  72         for _, raw := range response.Output {
  73                 item, err := chainInputItem(raw)
  74                 if err != nil {
  75                         return nil, "", err
  76                 }
  77
  78                 input = append(input, item)
  79
  80                 var call ChainFunctionCall
  81
  82                 err = json.Unmarshal(raw, &call)
  83                 if err != nil {
  84                         return nil, "", fmt.Errorf("decode output item: %w", err)
  85                 }
  86
  87                 switch call.Type {
  88                 case "message":
  89                         text, err := modelOutputText([]json.RawMessage{raw})
  90                         if err != nil {
  91                                 return nil, "", err
  92                         }
  93
  94                         text = strings.TrimSpace(text)
  95                         if text == "" {
  96                                 continue
  97                         }
  98
  99                         _, err = c.tools.log.Append("reply", text)
 100                         if err != nil {
 101                                 return nil, "", fmt.Errorf("%w: %w", errChainLogWrite, err)
 102                         }
 103
 104                         replies = append(replies, text)
 105                 case "function_call":
 106                         output, err := c.tools.Call(ctx, call.Name, call.Arguments)
 107                         if err != nil {
 108                                 return nil, "", fmt.Errorf("%w: %w", errChainLogWrite, err)
 109                         }
 110
 111                         outputs = append(outputs, ChainFunctionCallOutput{
 112                                 Type:   "function_call_output",
 113                                 CallID: call.CallID,
 114                                 Output: output,
 115                         })
 116                 }
 117         }
 118
 119         if len(outputs) == 0 {
 120                 return nil, strings.Join(replies, "\n\n"), nil
 121         }
 122
 123         input = append(input, outputs...)
 124         for _, message := range c.take() {
 125                 input = append(input, modelUserMessage(message.Text))
 126         }
 127
 128         return input, "", nil
 129 }
 130
 131 // ChainFunctionCall holds the fields of a function call output item.
 132 type ChainFunctionCall struct {
 133         Type      string `json:"type"`
 134         Name      string `json:"name"`
 135         Arguments string `json:"arguments"`
 136         CallID    string `json:"call_id"`
 137 }
 138
 139 // ChainFunctionCallOutput is the input item with the result of a function
 140 // call.
 141 type ChainFunctionCallOutput struct {
 142         Type   string `json:"type"`
 143         CallID string `json:"call_id"`
 144         Output string `json:"output"`
 145 }
 146
 147 // chainInputItem returns an output item as an input item for the next request.
 148 // The calls do not store their items, so the item goes back whole, without its
 149 // id.
 150 func chainInputItem(raw json.RawMessage) (map[string]any, error) {
 151         var item map[string]any
 152
 153         err := json.Unmarshal(raw, &item)
 154         if err != nil {
 155                 return nil, fmt.Errorf("decode output item: %w", err)
 156         }
 157
 158         delete(item, "id")
 159
 160         return item, nil
 161 }