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