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, then
  26 // its output items, then the tool results. The output items hold the reasoning
  27 // with its encrypted content, because the calls do not store it. So each
  28 // 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         images  *ImageStore
  36 }
  37
  38 // ChainFunctionCall holds the fields of a function call output item.
  39 type ChainFunctionCall struct {
  40         Type      string `json:"type"`
  41         Name      string `json:"name"`
  42         Arguments string `json:"arguments"`
  43         CallID    string `json:"call_id"`
  44 }
  45
  46 // ChainFunctionCallOutput is the input item with the result of a function call.
  47 type ChainFunctionCallOutput struct {
  48         Type   string `json:"type"`
  49         CallID string `json:"call_id"`
  50         Output any    `json:"output"`
  51 }
  52
  53 // Run runs the chain from input, the input of its first request, and returns
  54 // the text of the final answer. An error of a log write wraps errChainLogWrite.
  55 func (c *Chain) Run(ctx context.Context, input []any) (string, error) {
  56         request := c.request
  57         for request.Turn = "first"; ; request.Turn = "next" {
  58                 request.Input = input
  59                 response, err := c.client.Create(ctx, request)
  60                 if err != nil {
  61                         return "", fmt.Errorf("create model response: %w", err)
  62                 }
  63
  64                 var text string
  65
  66                 input, text, err = c.handle(ctx, input, response)
  67                 if err != nil {
  68                         return "", fmt.Errorf("handle model response: %w", err)
  69                 }
  70
  71                 if len(input) == 0 {
  72                         return text, nil
  73                 }
  74         }
  75 }
  76
  77 // handle logs the replies of response and runs its tool calls. It returns the
  78 // input of the next request: input, the output items, the tool results, and the
  79 // messages of take. After the final answer it returns no input, and the replies
  80 // of the answer as its text.
  81 func (c *Chain) handle(ctx context.Context, input []any, response ModelResponse) ([]any, string, error) {
  82         var (
  83                 outputs []any
  84                 replies []string
  85         )
  86
  87         for _, raw := range response.Output {
  88                 item, err := chainInputItem(raw)
  89                 if err != nil {
  90                         return nil, "", fmt.Errorf("decode chain input: %w", err)
  91                 }
  92
  93                 input = append(input, item)
  94
  95                 var call ChainFunctionCall
  96
  97                 err = json.Unmarshal(raw, &call)
  98                 if err != nil {
  99                         return nil, "", fmt.Errorf("decode output item: %w", err)
 100                 }
 101
 102                 switch call.Type {
 103                 case "message":
 104                         text, err := modelOutputText([]json.RawMessage{raw})
 105                         if err != nil {
 106                                 return nil, "", fmt.Errorf("read response text: %w", err)
 107                         }
 108
 109                         text = strings.TrimSpace(text)
 110                         if text == "" {
 111                                 continue
 112                         }
 113
 114                         _, err = c.tools.log.Append("reply", text)
 115                         if err != nil {
 116                                 return nil, "", fmt.Errorf("%w: %w", errChainLogWrite, err)
 117                         }
 118
 119                         replies = append(replies, text)
 120                 case "function_call":
 121                         output, err := c.tools.Call(ctx, call.Name, call.Arguments)
 122                         if err != nil {
 123                                 return nil, "", fmt.Errorf("%w: %w", errChainLogWrite, err)
 124                         }
 125
 126                         outputs = append(outputs, ChainFunctionCallOutput{
 127                                 Type:   "function_call_output",
 128                                 CallID: call.CallID,
 129                                 Output: output,
 130                         })
 131                 }
 132         }
 133
 134         if len(outputs) == 0 {
 135                 return nil, strings.Join(replies, "\n\n"), nil
 136         }
 137
 138         input = append(input, outputs...)
 139         for _, message := range c.take() {
 140                 item, err := imageUserMessage(c.images, message.Text)
 141                 if err != nil {
 142                         return nil, "", fmt.Errorf("read queued images: %w", err)
 143                 }
 144
 145                 input = append(input, item)
 146         }
 147
 148         return input, "", nil
 149 }
 150
 151 // chainInputItem returns an output item as an input item for the next request.
 152 // The calls do not store their items, so the item goes back whole, without its
 153 // id.
 154 func chainInputItem(raw json.RawMessage) (map[string]any, error) {
 155         var item map[string]any
 156
 157         err := json.Unmarshal(raw, &item)
 158         if err != nil {
 159                 return nil, fmt.Errorf("decode output item: %w", err)
 160         }
 161
 162         delete(item, "id")
 163
 164         return item, nil
 165 }