diff --git a/README.md b/README.md index e76e02b..40269d0 100644 --- a/README.md +++ b/README.md @@ -20,10 +20,10 @@ whiskr is a private, self-hosted web chat interface for interacting with AI mode - Reasoning effort control - Web search tool - Structured JSON output +- Statistics for messages (provider, ttft, tps and token count) ## TODO -- Statistics for messages (tps, token count, etc.) - Retry button for assistant messages - Import and export of chats - Image and file attachments diff --git a/chat.go b/chat.go index 825497c..e09590f 100644 --- a/chat.go +++ b/chat.go @@ -119,6 +119,9 @@ func HandleChat(w http.ResponseWriter, r *http.Request) { } request.Stream = true + request.Usage = &openrouter.IncludeUsage{ + Include: true, + } // DEBUG dump(request) @@ -145,6 +148,8 @@ func HandleChat(w http.ResponseWriter, r *http.Request) { return } + var id string + for { chunk, err := stream.Recv() if err != nil { @@ -157,6 +162,12 @@ func HandleChat(w http.ResponseWriter, r *http.Request) { return } + if id == "" { + id = chunk.ID + + response.Send(IDChunk(id)) + } + if len(chunk.Choices) == 0 { continue } diff --git a/main.go b/main.go index feaf343..03096e5 100644 --- a/main.go +++ b/main.go @@ -43,6 +43,7 @@ func main() { }) }) + r.Get("/-/stats/{id}", HandleStats) r.Post("/-/chat", HandleChat) if !NoOpen { diff --git a/openrouter.go b/openrouter.go index 5df1551..5fb6630 100644 --- a/openrouter.go +++ b/openrouter.go @@ -2,10 +2,44 @@ package main import ( "context" + "encoding/json" + "errors" + "fmt" + "net/http" "github.com/revrost/go-openrouter" ) +type Generation struct { + ID string `json:"id"` + TotalCost float64 `json:"total_cost"` + CreatedAt string `json:"created_at"` + Model string `json:"model"` + Origin string `json:"origin"` + Usage float64 `json:"usage"` + IsBYOK bool `json:"is_byok"` + UpstreamID *string `json:"upstream_id"` + CacheDiscount *float64 `json:"cache_discount"` + UpstreamInferenceCost *float64 `json:"upstream_inference_cost"` + AppID *int `json:"app_id"` + Streamed *bool `json:"streamed"` + Cancelled *bool `json:"cancelled"` + ProviderName *string `json:"provider_name"` + Latency *int `json:"latency"` + ModerationLatency *int `json:"moderation_latency"` + GenerationTime *int `json:"generation_time"` + FinishReason *string `json:"finish_reason"` + NativeFinishReason *string `json:"native_finish_reason"` + TokensPrompt *int `json:"tokens_prompt"` + TokensCompletion *int `json:"tokens_completion"` + NativeTokensPrompt *int `json:"native_tokens_prompt"` + NativeTokensCompletion *int `json:"native_tokens_completion"` + NativeTokensReasoning *int `json:"native_tokens_reasoning"` + NumMediaPrompt *int `json:"num_media_prompt"` + NumMediaCompletion *int `json:"num_media_completion"` + NumSearchResults *int `json:"num_search_results"` +} + func OpenRouterClient() *openrouter.Client { return openrouter.NewClient(OpenRouterToken) } @@ -20,3 +54,34 @@ func OpenRouterStartStream(ctx context.Context, request openrouter.ChatCompletio return stream, nil } + +func OpenRouterGetGeneration(ctx context.Context, id string) (*Generation, error) { + req, err := http.NewRequest("GET", fmt.Sprintf("https://openrouter.ai/api/v1/generation?id=%s", id), nil) + if err != nil { + return nil, err + } + + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", OpenRouterToken)) + + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, err + } + + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, errors.New(resp.Status) + } + + var response struct { + Data Generation `json:"data"` + } + + err = json.NewDecoder(resp.Body).Decode(&response) + if err != nil { + return nil, err + } + + return &response.Data, nil +} diff --git a/static/css/chat.css b/static/css/chat.css index eb1ab88..e15aab0 100644 --- a/static/css/chat.css +++ b/static/css/chat.css @@ -146,7 +146,6 @@ body.loading #version { width: max-content; padding-top: 28px; background: #363a4f; - overflow: hidden; flex-shrink: 0; } @@ -170,13 +169,16 @@ body.loading #version { left: 6px; } +.statistics .provider::after, +.statistics .ttft::after, +.statistics .tps::after, .message .tags::before { content: ""; position: absolute; top: 7px; left: -10px; height: 2px; - width: 5px; + width: 6px; background: #939ab7; } @@ -284,6 +286,11 @@ body.loading #version { display: none; } +#messages .message div.text { + border-bottom-left-radius: 6px; + border-bottom-right-radius: 6px; +} + .message.has-reasoning .text { padding-top: 4px; } @@ -357,6 +364,61 @@ body.loading #version { content: ". . ."; } +.statistics { + position: absolute; + transition: 150ms; + top: calc(100% + 5px); + left: 8px; + display: flex; + align-items: center; + gap: 20px; + white-space: nowrap; + font-size: 13px; + line-height: 13px; + pointer-events: none; +} + +.statistics .provider, +.statistics .ttft, +.statistics .tps, +.statistics .tokens { + position: relative; + display: flex; + align-items: center; + gap: 3px; +} + +.statistics .provider::after, +.statistics .ttft::after, +.statistics .tps::after { + left: unset; + right: -14px; +} + +.statistics .provider::before { + background-image: url(icons/provider.svg); +} + +.statistics .ttft::before { + background-image: url(icons/ttft.svg); +} + +.statistics .tps::before { + background-image: url(icons/tps.svg); +} + +.statistics .tokens::before { + background-image: url(icons/amount.svg); +} + +.message:not(:hover) .statistics { + opacity: 0; +} + +.message:not(.has-statistics) .statistics { + display: none; +} + #chat { display: flex; position: relative; @@ -451,16 +513,20 @@ body.loading #version, #messages .message .role::before, #messages .message .tag-json, #messages .message .tag-search, +#messages .message .copy, +#messages .message .edit, +#messages .message .delete, +.pre-copy, +.message .statistics .provider::before, +.message .statistics .ttft::before, +.message .statistics .tps::before, +.message .statistics .tokens::before, #json, #search, #scrolling, #clear, #add, #send, -.pre-copy, -#messages .message .copy, -#messages .message .edit, -.message .delete, #chat .option label { display: block; width: 20px; @@ -470,6 +536,10 @@ body.loading #version, background-repeat: no-repeat; } +.message .statistics .provider::before, +.message .statistics .ttft::before, +.message .statistics .tps::before, +.message .statistics .tokens::before, #messages .message .tag-json, #messages .message .tag-search, #messages .message .role::before { @@ -500,7 +570,7 @@ input.invalid { background-image: url(icons/save.svg); } -.message .delete { +#messages .message .delete { background-image: url(icons/delete.svg); } diff --git a/static/css/icons/provider.svg b/static/css/icons/provider.svg new file mode 100644 index 0000000..f483e23 --- /dev/null +++ b/static/css/icons/provider.svg @@ -0,0 +1,7 @@ + + + + + + + \ No newline at end of file diff --git a/static/css/icons/tps.svg b/static/css/icons/tps.svg new file mode 100644 index 0000000..7e8a247 --- /dev/null +++ b/static/css/icons/tps.svg @@ -0,0 +1,7 @@ + + + + + + + \ No newline at end of file diff --git a/static/css/icons/ttft.svg b/static/css/icons/ttft.svg new file mode 100644 index 0000000..6e3161c --- /dev/null +++ b/static/css/icons/ttft.svg @@ -0,0 +1,7 @@ + + + + + + + \ No newline at end of file diff --git a/static/js/chat.js b/static/js/chat.js index dbd7b0c..5df79ff 100644 --- a/static/js/chat.js +++ b/static/js/chat.js @@ -45,6 +45,7 @@ #text; #tags = []; + #statistics; #error = false; #editing = false; @@ -60,6 +61,7 @@ #_reasoning; #_text; #_edit; + #_statistics; constructor(role, reasoning, text) { this.#id = uid(); @@ -202,6 +204,11 @@ this.delete(); }); + // statistics + this.#_statistics = make("div", "statistics"); + + this.#_message.appendChild(this.#_statistics); + // add to dom $messages.appendChild(this.#_message); @@ -268,13 +275,42 @@ #render(only = false, noScroll = false) { if (!only || only === "tags") { - this.#_tags.innerHTML = this.#tags - .map((tag) => `
`) - .join(""); + const tags = this.#tags.map( + (tag) => `
`, + ); + + this.#_tags.innerHTML = tags.join(""); this.#_message.classList.toggle("has-tags", this.#tags.length > 0); } + if (!only || only === "statistics") { + let html = ""; + + if (this.#statistics) { + const { provider, ttft, time, input, output } = this.#statistics; + + const tps = output / (time / 1000); + + html = [ + provider ? `
${provider}
` : "", + `
${formatMilliseconds(ttft)}
`, + `
${fixed(tps, 2)} t/s
`, + `
+
${input}
+ + +
${output}
+ = +
${input + output}
+
`, + ].join(""); + } + + this.#_statistics.innerHTML = html; + + this.#_message.classList.toggle("has-statistics", !!html); + } + if (this.#error) { return; } @@ -329,9 +365,39 @@ data.tags = this.#tags; } + if (this.#statistics) { + data.statistics = this.#statistics; + } + return data; } + setStatistics(statistics) { + this.#statistics = statistics; + + this.#render("statistics"); + this.#save(); + } + + async loadGenerationData(generationID) { + if (!generationID) { + return; + } + + try { + const response = await fetch(`/-/stats/${generationID}`), + data = await response.json(); + + if (!data || data.error) { + throw new Error(data?.error || response.statusText); + } + + this.setStatistics(data); + } catch (err) { + console.error(err); + } + } + addTag(tag) { if (this.#tags.includes(tag)) { return; @@ -586,6 +652,10 @@ obj.showError(message.error); } + if (message.statistics) { + obj.setStatistics(message.statistics); + } + if (message.tags) { message.tags.forEach((tag) => obj.addTag(tag)); } @@ -804,6 +874,8 @@ message.addTag("search"); } + let generationID; + stream( "/-/chat", { @@ -822,10 +894,18 @@ $chat.classList.remove("completing"); + setTimeout(() => { + message.loadGenerationData(generationID); + }, 750); + return; } switch (chunk.type) { + case "id": + generationID = chunk.text; + + break; case "reason": message.setState("reasoning"); message.addReasoning(chunk.text); diff --git a/static/js/lib.js b/static/js/lib.js index 1469f6a..fb0eca7 100644 --- a/static/js/lib.js +++ b/static/js/lib.js @@ -60,3 +60,17 @@ function escapeHtml(text) { .replace(//g, ">"); } + +function formatMilliseconds(ms) { + if (ms < 1000) { + return `${ms}ms`; + } else if (ms < 10000) { + return `${(ms / 1000).toFixed(1)}s`; + } + + return `${Math.round(ms / 1000)}s`; +} + +function fixed(num, decimals = 0) { + return num.toFixed(decimals).replace(/\.?0+$/m, ""); +} diff --git a/stats.go b/stats.go new file mode 100644 index 0000000..535cc18 --- /dev/null +++ b/stats.go @@ -0,0 +1,67 @@ +package main + +import ( + "net/http" + "strings" + + "github.com/go-chi/chi/v5" +) + +type Statistics struct { + Provider *string `json:"provider,omitempty"` + Model string `json:"model"` + Cost float64 `json:"cost"` + TTFT int `json:"ttft"` + Time int `json:"time"` + InputTokens int `json:"input"` + OutputTokens int `json:"output"` +} + +func HandleStats(w http.ResponseWriter, r *http.Request) { + id := chi.URLParam(r, "id") + + if id == "" || !strings.HasPrefix(id, "gen-") { + RespondJson(w, http.StatusBadRequest, map[string]any{ + "error": "invalid id", + }) + + return + } + + generation, err := OpenRouterGetGeneration(r.Context(), id) + if err != nil { + RespondJson(w, http.StatusInternalServerError, map[string]any{ + "error": err.Error(), + }) + + return + } + + statistics := Statistics{ + Provider: generation.ProviderName, + Model: generation.Model, + Cost: generation.TotalCost, + TTFT: Nullable(generation.Latency, 0), + Time: Nullable(generation.GenerationTime, 0), + } + + nativeIn := Nullable(generation.NativeTokensPrompt, 0) + normalIn := Nullable(generation.TokensPrompt, 0) + + statistics.InputTokens = max(nativeIn, normalIn) + + nativeOut := Nullable(generation.NativeTokensCompletion, 0) + Nullable(generation.NativeTokensReasoning, 0) + normalOut := Nullable(generation.TokensCompletion, 0) + + statistics.OutputTokens = max(nativeOut, normalOut) + + RespondJson(w, http.StatusOK, statistics) +} + +func Nullable[T any](ptr *T, def T) T { + if ptr == nil { + return def + } + + return *ptr +} diff --git a/stream.go b/stream.go index eba0e35..db129b5 100644 --- a/stream.go +++ b/stream.go @@ -64,6 +64,13 @@ func TextChunk(text string) Chunk { } } +func IDChunk(id string) Chunk { + return Chunk{ + Type: "id", + Text: id, + } +} + func ErrorChunk(err error) Chunk { return Chunk{ Type: "error",