X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/97274864799ac2eea84bf80c9727a810b695db55..153f3f9b4b09d521412eacca30eddeed395c3f41:/probe.go?ds=inline

diff --git a/probe.go b/probe.go
index 9d272ec..32a3ba9 100644
--- a/probe.go
+++ b/probe.go
@@ -98,7 +98,7 @@ func probe(ctx context.Context, args []string, out io.Writer) (err error) {
 	}
 
 	usagePath := filepath.Join(root, "usage.jsonl")
-	client := newClient(address, auth, "onechat-"+probeName, usagePath)
+	client := newModelClient(address, auth, "onechat-"+probeName, usagePath)
 	client.LimitCalls(probeCallLimit)
 	summarizer := slotSummarizer{
 		inner: newCompactor(client, probePrompt, config.Compactor),
@@ -140,7 +140,7 @@ func probe(ctx context.Context, args []string, out io.Writer) (err error) {
 // by the rules of probeStop.
 type probeRun struct {
 	memory    *Memory
-	client    *Client
+	client    *ModelClient
 	usagePath string
 }
 
@@ -295,7 +295,7 @@ func readProbeLog(path string, count int) (entries []probeEntry, err error) {
 // probeStop returns why a probe run must stop: a failed call, or a quota window
 // whose used percent grew by more than probeQuotaLimit points since the first
 // call. It returns "" when the run can go on.
-func probeStop(records []usageRecord) string {
+func probeStop(records []ModelUsageRecord) string {
 	for _, record := range records {
 		if record.Error != "" {
 			return "call failed: " + record.Error
@@ -375,7 +375,7 @@ func builtByModel(memory *Memory, key nodeKey) (bool, error) {
 // readUsageFile reads the records of usage.jsonl at path, in the order of the
 // start of their calls. The file has them in the order of the end of their
 // calls, which differs when calls run at once.
-func readUsageFile(path string) (records []usageRecord, err error) {
+func readUsageFile(path string) (records []ModelUsageRecord, err error) {
 	file, err := os.Open(path)
 	if errors.Is(err, os.ErrNotExist) {
 		return nil, nil
@@ -392,7 +392,7 @@ func readUsageFile(path string) (records []usageRecord, err error) {
 	scanner := bufio.NewScanner(file)
 	scanner.Buffer(make([]byte, 64*1024), probeLineBytes)
 	for scanner.Scan() {
-		var record usageRecord
+		var record ModelUsageRecord
 		err = json.Unmarshal(scanner.Bytes(), &record)
 		if err != nil {
 			return nil, fmt.Errorf("decode usage: %w", err)
@@ -406,7 +406,7 @@ func readUsageFile(path string) (records []usageRecord, err error) {
 		return nil, fmt.Errorf("read usage file: %w", err)
 	}
 
-	slices.SortStableFunc(records, func(a, b usageRecord) int {
+	slices.SortStableFunc(records, func(a, b ModelUsageRecord) int {
 		return callStart(a).Compare(callStart(b))
 	})
 
@@ -415,7 +415,7 @@ func readUsageFile(path string) (records []usageRecord, err error) {
 
 // callStart returns the start of the call of record, or the zero time when its
 // date is not valid.
-func callStart(record usageRecord) time.Time {
+func callStart(record ModelUsageRecord) time.Time {
 	start, err := time.Parse(time.RFC3339Nano, record.Date)
 	if err != nil {
 		return time.Time{}
@@ -426,7 +426,7 @@ func callStart(record usageRecord) time.Time {
 
 // usedPercents returns the used percent of each quota window of record, by
 // header name.
-func usedPercents(record usageRecord) map[string]float64 {
+func usedPercents(record ModelUsageRecord) map[string]float64 {
 	used := make(map[string]float64)
 	for name, value := range record.Quota {
 		if !strings.HasSuffix(name, "-used-percent") {
@@ -443,7 +443,7 @@ func usedPercents(record usageRecord) map[string]float64 {
 }
 
 // writeProbeReport writes the result of a probe run.
-func writeProbeReport(out io.Writer, appended int, stop string, records []usageRecord, sizes []int) error {
+func writeProbeReport(out io.Writer, appended int, stop string, records []ModelUsageRecord, sizes []int) error {
 	var input, cached, output, failed int
 	for _, record := range records {
 		input += record.InputTokens
@@ -502,7 +502,7 @@ func sizeText(sizes []int) string {
 
 // timeText describes the time per call and how many calls ran at once. A call
 // runs from its date for its seconds.
-func timeText(records []usageRecord) string {
+func timeText(records []ModelUsageRecord) string {
 	type event struct {
 		at    time.Time
 		delta int
@@ -550,7 +550,7 @@ func timeText(records []usageRecord) string {
 }
 
 // quotaText returns the quota headers of record, sorted by name.
-func quotaText(record usageRecord) string {
+func quotaText(record ModelUsageRecord) string {
 	names := slices.Sorted(maps.Keys(record.Quota))
 	parts := make([]string, len(names))
 	for n, name := range names {