1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
|
// Package chatrun holds the shared chat-running primitives that used to be
// duplicated across the three Hexai surfaces: the CLI (internal/hexaicli), the
// LSP server (internal/lsp) and the tmux code-action tool (internal/hexaiaction).
//
// Every surface had to do the same two things when talking to an LLM:
//
// 1. Invoke the client, preferring streaming (llm.Streamer) when the provider
// implements it and falling back to a single Chat call otherwise, while
// collecting the full response text.
// 2. Account for the exchange: count the bytes sent (sum of message contents)
// and received (response length) and feed them into the stats package.
//
// Both pieces lived in slightly diverging copies (runChatRequest/runStreaming
// Chat/runSimpleChat in the CLI, chatWithStats in the LSP, runOnce in the
// action tool). This package centralises them so all three share one
// implementation and behaviour stays identical.
package chatrun
import (
"context"
"fmt"
"io"
"strings"
"codeberg.org/snonux/hexai/internal/llm"
"codeberg.org/snonux/hexai/internal/stats"
)
// Chatter is the minimal client capability Invoke needs: a single Chat call.
// It is deliberately narrower than llm.Client (no Name/DefaultModel) so callers
// that only have a chat-doer abstraction — such as the code-action surface —
// can use Invoke without widening their own interface. Clients that also
// implement llm.Streamer get streaming automatically via a type assertion.
type Chatter interface {
Chat(ctx context.Context, messages []llm.Message, opts ...llm.RequestOption) (string, error)
}
// Invoke sends msgs to the client and returns the full assistant response.
//
// When the client implements llm.Streamer the response is streamed and, if out
// is non-nil, each chunk is forwarded to out as it arrives (this is how the CLI
// renders incremental output). Otherwise a single Chat call is made and, when
// out is non-nil, the whole response is written to it once.
//
// Passing a nil out collects the response without writing it anywhere, which is
// what the LSP and code-action surfaces need (they post-process the text before
// applying it to a document).
func Invoke(ctx context.Context, client Chatter, msgs []llm.Message, opts []llm.RequestOption, out io.Writer) (string, error) {
if streamer, ok := client.(llm.Streamer); ok {
return invokeStreaming(ctx, streamer, msgs, opts, out)
}
return invokeSimple(ctx, client, msgs, opts, out)
}
// invokeStreaming drives ChatStream, accumulating the full text while
// optionally mirroring each chunk to out. A write error to out aborts further
// writes and is returned once streaming finishes.
func invokeStreaming(ctx context.Context, client llm.Streamer, msgs []llm.Message, opts []llm.RequestOption, out io.Writer) (string, error) {
var output strings.Builder
var writeErr error
err := client.ChatStream(ctx, msgs, func(chunk string) {
output.WriteString(chunk)
if out == nil || writeErr != nil {
return
}
if _, werr := io.WriteString(out, chunk); werr != nil {
writeErr = werr
}
}, opts...)
if err != nil {
return "", err
}
if writeErr != nil {
return "", writeErr
}
return output.String(), nil
}
// invokeSimple performs a single Chat call and, when out is non-nil, writes the
// whole response to it.
func invokeSimple(ctx context.Context, client Chatter, msgs []llm.Message, opts []llm.RequestOption, out io.Writer) (string, error) {
output, err := client.Chat(ctx, msgs, opts...)
if err != nil {
return "", err
}
if out != nil {
if _, werr := fmt.Fprint(out, output); werr != nil {
return "", werr
}
}
return output, nil
}
// SentBytes returns the total number of content bytes across msgs. This is the
// "sent" figure every surface reports and feeds into the stats package.
func SentBytes(msgs []llm.Message) int {
sent := 0
for _, m := range msgs {
sent += len(m.Content)
}
return sent
}
// Account records a completed exchange in the stats package and returns the
// sent/received byte counts so callers can build their own summaries. The
// stats.Update error is intentionally swallowed because none of the surfaces
// treat a stats failure as fatal; callers that want to log it can call
// stats.Update directly instead.
func Account(ctx context.Context, provider, model string, msgs []llm.Message, output string) (sent, recv int) {
sent = SentBytes(msgs)
recv = len(output)
_ = stats.Update(ctx, provider, model, sent, recv)
return sent, recv
}
|
