]> git.giacomolli.uk - onechat/blob - chat.go
Name and order the ChatGPT login code
[onechat] / chat.go
1 package main
2
3 import (
4         "bufio"
5         "context"
6         "encoding/json"
7         "errors"
8         "fmt"
9         "io"
10         "net"
11         "path/filepath"
12         "strings"
13 )
14
15 const (
16         // stopCommand is the input line that interrupts the running turn.
17         stopCommand = "/stop"
18
19         // resultLines is the most lines of a tool result that the client shows.
20         resultLines = 10
21 )
22
23 // chat runs "onechat chat": the terminal client of the daemon. Each line of
24 // input is a user message, and the line /stop interrupts the running turn. The
25 // client shows the replies, the tool calls, and their results as they happen.
26 // It ends at the end of the input or when the daemon closes the connection.
27 func chat(ctx context.Context, in io.Reader, out io.Writer) error {
28         dir, err := dataDir()
29         if err != nil {
30                 return err
31         }
32
33         path := filepath.Join(dir, socketFileName)
34         err = checkSocketPath(path)
35         if err != nil {
36                 return err
37         }
38
39         var dialer net.Dialer
40         conn, err := dialer.DialContext(ctx, "unix", path)
41         if err != nil {
42                 return fmt.Errorf(`connect to the daemon (start it with "onechat daemon"): %w`, err)
43         }
44
45         return chatSession(ctx, conn, in, out)
46 }
47
48 // errDaemonClosed tells that the daemon closed the connection.
49 var errDaemonClosed = errors.New("the daemon closed the connection")
50
51 // chatSession runs the client on conn until the input ends, the daemon closes
52 // the connection, or ctx ends, and closes conn. A read of the input that waits
53 // for a line can stay behind when the daemon closes the connection.
54 func chatSession(ctx context.Context, conn net.Conn, in io.Reader, out io.Writer) error {
55         sent := make(chan error, 1)
56         go func() {
57                 sent <- sendLines(ctx, conn, in)
58         }()
59
60         shown := make(chan error, 1)
61         go func() {
62                 shown <- showEvents(conn, out)
63         }()
64
65         var err error
66         select {
67         case err = <-sent:
68                 err = errors.Join(err, conn.Close())
69                 <-shown
70         case err = <-shown:
71                 err = errors.Join(cmpError(err, errDaemonClosed), conn.Close())
72         case <-ctx.Done():
73                 err = conn.Close()
74                 <-shown
75         }
76
77         return err
78 }
79
80 // cmpError returns err, or fallback when err is nil.
81 func cmpError(err, fallback error) error {
82         if err == nil {
83                 return fallback
84         }
85
86         return err
87 }
88
89 // sendLines sends each line of in as a request until the input ends.
90 func sendLines(ctx context.Context, conn net.Conn, in io.Reader) error {
91         encoder := json.NewEncoder(conn)
92         scanner := bufio.NewScanner(in)
93         scanner.Buffer(make([]byte, 64*1024), requestBytes)
94         for ctx.Err() == nil && scanner.Scan() {
95                 line := scanner.Text()
96                 request := clientRequest{
97                         Type: "send",
98                         Text: line,
99                 }
100
101                 switch strings.TrimSpace(line) {
102                 case "":
103                         continue
104                 case stopCommand:
105                         request = clientRequest{
106                                 Type: "interrupt",
107                         }
108                 }
109
110                 err := encoder.Encode(request)
111                 if err != nil {
112                         return fmt.Errorf("send to the daemon: %w", err)
113                 }
114         }
115
116         return scanner.Err()
117 }
118
119 // showEvents writes each event from the daemon to out until the connection
120 // closes.
121 func showEvents(conn net.Conn, out io.Writer) error {
122         scanner := bufio.NewScanner(conn)
123         scanner.Buffer(make([]byte, 64*1024), requestBytes)
124         for scanner.Scan() {
125                 var event Event
126                 err := json.Unmarshal(scanner.Bytes(), &event)
127                 if err != nil {
128                         return fmt.Errorf("decode event: %w", err)
129                 }
130
131                 text := eventText(event)
132                 if text == "" {
133                         continue
134                 }
135
136                 _, err = io.WriteString(out, text+"\n")
137                 if err != nil {
138                         return err
139                 }
140         }
141
142         return scanner.Err()
143 }
144
145 // eventText returns the text that the client shows for event, or "" for none.
146 // The client shows a user message only from another source, such as telegram,
147 // because the user typed the others. A tool result shows its first resultLines
148 // lines.
149 func eventText(event Event) string {
150         switch event.Kind {
151         case "ready":
152                 return "[connected to " + event.Text + "]"
153         case "user":
154                 if event.Source == terminalSource {
155                         return ""
156                 }
157
158                 return "[" + event.Source + "] " + event.Text
159         case "reply":
160                 return event.Text
161         case "tool":
162                 return "[tool] " + event.Text
163         case "echo":
164                 return "[result] " + firstLines(event.Text, resultLines)
165         case "busy":
166                 return "[working]"
167         case "idle":
168                 return "[done]"
169         case "error":
170                 return "[error] " + event.Text
171         default:
172                 return ""
173         }
174 }
175
176 // firstLines returns the first count lines of text, with the number of the
177 // lines left out.
178 func firstLines(text string, count int) string {
179         text = strings.TrimRight(text, "\n")
180         lines := strings.Split(text, "\n")
181         if len(lines) <= count {
182                 return text
183         }
184
185         return strings.Join(lines[:count], "\n") + fmt.Sprintf("\n[%d more lines]", len(lines)-count)
186 }