109 lines
2.9 KiB
Go
109 lines
2.9 KiB
Go
package llm
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
|
|
"go.maximumdirect.net/promptkit/internal/domain"
|
|
)
|
|
|
|
const maxAuthenticationResponseBytes = 1 << 20
|
|
|
|
// AuthenticationEvidence describes the observable result of a key probe.
|
|
type AuthenticationEvidence struct {
|
|
HTTP HTTPExecutionEvidence
|
|
Accepted *bool
|
|
Overflow bool
|
|
Unrecognized bool
|
|
}
|
|
|
|
// OpenRouterKeyEndpoint derives the default key endpoint from a base endpoint.
|
|
func OpenRouterKeyEndpoint(base string) (string, error) {
|
|
return url.JoinPath(base, "key")
|
|
}
|
|
|
|
// ProbeOpenRouterKey executes and interprets an OpenRouter key request.
|
|
func (c *OpenAICompatibleClient) ProbeOpenRouterKey(ctx context.Context, endpoint string, timeoutSeconds int, target domain.ExecutionTarget) (AuthenticationEvidence, error) {
|
|
parsed, err := url.Parse(endpoint)
|
|
if err != nil {
|
|
return AuthenticationEvidence{}, err
|
|
}
|
|
parsed.User = nil
|
|
response, evidence, requestErr := c.GetWithAuthenticationEvidence(ctx, parsed.String(), timeoutSeconds, target)
|
|
result := AuthenticationEvidence{HTTP: evidence}
|
|
if response != nil && response.Body != nil {
|
|
defer response.Body.Close()
|
|
}
|
|
if !evidence.ResponseReceived && response == nil {
|
|
return result, requestErr
|
|
}
|
|
status := evidence.HTTPStatus
|
|
if status == nil && response != nil {
|
|
status = &response.StatusCode
|
|
}
|
|
if status == nil {
|
|
result.Unrecognized = true
|
|
return result, nil
|
|
}
|
|
switch *status {
|
|
case http.StatusUnauthorized:
|
|
accepted := false
|
|
result.Accepted = &accepted
|
|
case http.StatusOK:
|
|
if response == nil || response.Body == nil {
|
|
result.Unrecognized = true
|
|
return result, nil
|
|
}
|
|
contents, err := io.ReadAll(io.LimitReader(response.Body, maxAuthenticationResponseBytes+1))
|
|
if err != nil && ctx.Err() != nil {
|
|
return result, ctx.Err()
|
|
}
|
|
if len(contents) > maxAuthenticationResponseBytes {
|
|
result.Overflow = true
|
|
return result, nil
|
|
}
|
|
if err != nil {
|
|
return result, err
|
|
}
|
|
if recognizedOpenRouterKeyResponse(contents) {
|
|
accepted := true
|
|
result.Accepted = &accepted
|
|
} else {
|
|
result.Unrecognized = true
|
|
}
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func recognizedOpenRouterKeyResponse(contents []byte) bool {
|
|
var envelope map[string]json.RawMessage
|
|
if err := json.Unmarshal(contents, &envelope); err != nil || envelope == nil {
|
|
return false
|
|
}
|
|
var data map[string]json.RawMessage
|
|
if err := json.Unmarshal(envelope["data"], &data); err != nil || data == nil {
|
|
return false
|
|
}
|
|
label := bytes.TrimSpace(data["label"])
|
|
var labelValue string
|
|
if len(label) == 0 || label[0] != '"' || json.Unmarshal(label, &labelValue) != nil {
|
|
return false
|
|
}
|
|
freeTier := bytes.TrimSpace(data["is_free_tier"])
|
|
if !bytes.Equal(freeTier, []byte("true")) && !bytes.Equal(freeTier, []byte("false")) {
|
|
return false
|
|
}
|
|
decoder := json.NewDecoder(bytes.NewReader(data["usage"]))
|
|
decoder.UseNumber()
|
|
var number any
|
|
if decoder.Decode(&number) != nil {
|
|
return false
|
|
}
|
|
_, ok := number.(json.Number)
|
|
return ok
|
|
}
|