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