216 lines
8.3 KiB
Go
216 lines
8.3 KiB
Go
package models
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"net/http"
|
|
"net/url"
|
|
"strings"
|
|
"time"
|
|
|
|
"git.cin.su/m/Assistant/contracts/capability"
|
|
"git.cin.su/m/Assistant/contracts/cognitive"
|
|
)
|
|
|
|
type OpenAICompatibleLocalAdapter struct {
|
|
base *url.URL
|
|
client *http.Client
|
|
apiKey string
|
|
}
|
|
|
|
func NewOpenAICompatibleLocalAdapter(raw string, transport http.RoundTripper) (*OpenAICompatibleLocalAdapter, error) {
|
|
return NewOpenAICompatibleLocalAdapterWithAPIKey(raw, "", transport)
|
|
}
|
|
|
|
// NewOpenAICompatibleLocalAdapterWithAPIKey creates an adapter for a gateway
|
|
// whose bearer credential was resolved by the trusted runtime boundary.
|
|
func NewOpenAICompatibleLocalAdapterWithAPIKey(raw, apiKey string, transport http.RoundTripper) (*OpenAICompatibleLocalAdapter, error) {
|
|
return NewOpenAICompatibleLocalAdapterWithAPIKeyAndTimeout(raw, apiKey, transport, 10*time.Second)
|
|
}
|
|
|
|
// NewOpenAICompatibleLocalAdapterWithAPIKeyAndTimeout creates a local gateway
|
|
// adapter with an explicit invocation deadline. The request context remains
|
|
// authoritative for per-call cancellation; this client deadline prevents a
|
|
// provider from holding a worker indefinitely.
|
|
func NewOpenAICompatibleLocalAdapterWithAPIKeyAndTimeout(raw, apiKey string, transport http.RoundTripper, timeout time.Duration) (*OpenAICompatibleLocalAdapter, error) {
|
|
base, err := url.Parse(raw)
|
|
if err != nil || base.Scheme != "http" || base.Host == "" || strings.TrimSuffix(base.Path, "/") != "/v1" || base.RawQuery != "" || base.Fragment != "" || base.User != nil {
|
|
return nil, errors.New("invalid local compatible endpoint")
|
|
}
|
|
if timeout <= 0 {
|
|
return nil, errors.New("local compatible timeout must be positive")
|
|
}
|
|
if transport == nil {
|
|
transport = http.DefaultTransport
|
|
}
|
|
return &OpenAICompatibleLocalAdapter{base: base, client: &http.Client{Transport: transport, Timeout: timeout, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}, apiKey: strings.TrimSpace(apiKey)}, nil
|
|
}
|
|
|
|
type compatibleModelsResponse struct {
|
|
Data []struct {
|
|
ID string `json:"id"`
|
|
Created int64 `json:"created"`
|
|
OwnedBy string `json:"owned_by"`
|
|
Meta struct {
|
|
Context int64 `json:"n_ctx"`
|
|
TrainingContext int64 `json:"n_ctx_train"`
|
|
Parameters int64 `json:"n_params"`
|
|
Size int64 `json:"size"`
|
|
FileType string `json:"ftype"`
|
|
} `json:"meta"`
|
|
} `json:"data"`
|
|
}
|
|
|
|
func (a *OpenAICompatibleLocalAdapter) Probe(ctx context.Context) ProviderProbe {
|
|
_, probe := a.Discover(ctx)
|
|
return probe
|
|
}
|
|
|
|
func (a *OpenAICompatibleLocalAdapter) Discover(ctx context.Context) ([]DiscoveredModel, ProviderProbe) {
|
|
var response compatibleModelsResponse
|
|
status, err := a.requestJSON(ctx, http.MethodGet, "/models", nil, &response)
|
|
if err != nil {
|
|
return nil, probeFailure(err)
|
|
}
|
|
if status != http.StatusOK {
|
|
return nil, ProviderProbe{Health: capability.HealthUnavailable, Failure: mapCompatibleStatus(status)}
|
|
}
|
|
models := make([]DiscoveredModel, 0, len(response.Data))
|
|
for _, item := range response.Data {
|
|
if item.ID == "" {
|
|
continue
|
|
}
|
|
stableMetadata := fmt.Sprintf("%s\x00%s\x00%d\x00%d\x00%d\x00%d\x00%s", item.ID, item.OwnedBy, item.Meta.Context, item.Meta.TrainingContext, item.Meta.Parameters, item.Meta.Size, item.Meta.FileType)
|
|
digest := sha256.Sum256([]byte(stableMetadata))
|
|
contextBytes := item.Meta.Context * 4
|
|
models = append(models, DiscoveredModel{ProviderModelID: capability.ProviderModelID(item.ID), Fingerprint: capability.ModelFingerprint(hex.EncodeToString(digest[:])), MaxContextBytes: contextBytes, Modalities: []capability.ModelModality{capability.ModalityText}, Capabilities: []capability.ModelCapability{capability.CapabilityUnknown}, SizeBytes: item.Meta.Size, Runtime: item.OwnedBy, Quantization: item.Meta.FileType})
|
|
}
|
|
return models, ProviderProbe{Health: capability.HealthAvailable}
|
|
}
|
|
|
|
type compatibleChatRequest struct {
|
|
Model string `json:"model"`
|
|
Messages []compatibleMessage `json:"messages"`
|
|
Temperature int `json:"temperature"`
|
|
Stream bool `json:"stream"`
|
|
MaxTokens int `json:"max_tokens"`
|
|
}
|
|
type compatibleMessage struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
}
|
|
type compatibleChatResponse struct {
|
|
Choices []struct {
|
|
FinishReason string `json:"finish_reason"`
|
|
Message compatibleMessage `json:"message"`
|
|
} `json:"choices"`
|
|
Usage struct {
|
|
Prompt int64 `json:"prompt_tokens"`
|
|
Completion int64 `json:"completion_tokens"`
|
|
} `json:"usage"`
|
|
}
|
|
|
|
func (a *OpenAICompatibleLocalAdapter) Advise(ctx context.Context, request cognitive.AdvisoryRequest) (cognitive.AdvisoryResponse, error) {
|
|
if err := ValidateAdvisoryRequest(request); err != nil || request.ProviderModelID == "" || len(request.Context) != 0 {
|
|
return cognitive.AdvisoryResponse{}, ErrInvalidAdvisory
|
|
}
|
|
payload := compatibleChatRequest{Model: string(request.ProviderModelID), Messages: []compatibleMessage{{Role: "system", Content: "Provide non-canonical advisory text only. Never claim that an action occurred."}, {Role: "user", Content: request.Message}}, Temperature: 0, Stream: false, MaxTokens: request.MaxOutputBytes / 4}
|
|
var result compatibleChatResponse
|
|
status, err := a.requestJSON(ctx, http.MethodPost, "/chat/completions", payload, &result)
|
|
if err != nil {
|
|
return advisoryFailure(request, probeFailure(err).Failure), nil
|
|
}
|
|
if status != http.StatusOK {
|
|
return advisoryFailure(request, mapCompatibleStatus(status)), nil
|
|
}
|
|
if len(result.Choices) != 1 || result.Choices[0].Message.Content == "" {
|
|
return advisoryFailure(request, cognitive.FinishInvalidProviderResponse), nil
|
|
}
|
|
finish := cognitive.FinishCompleted
|
|
if result.Choices[0].FinishReason == "length" {
|
|
finish = cognitive.FinishLengthLimit
|
|
}
|
|
if finish != cognitive.FinishCompleted {
|
|
return advisoryFailure(request, finish), nil
|
|
}
|
|
response := cognitive.AdvisoryResponse{Version: request.Version, RequestID: request.RequestID, ProviderBindingID: request.ProviderBindingID, ProviderBindingRevision: request.ProviderBindingRevision, AdvisoryText: result.Choices[0].Message.Content, FinishReason: finish, Usage: &cognitive.ProviderUsage{InputUnits: result.Usage.Prompt, OutputUnits: result.Usage.Completion}}
|
|
if ValidateAdvisoryResponse(request, response) != nil {
|
|
return advisoryFailure(request, cognitive.FinishInvalidProviderResponse), nil
|
|
}
|
|
return response, nil
|
|
}
|
|
|
|
func (a *OpenAICompatibleLocalAdapter) requestJSON(ctx context.Context, method, path string, input, output any) (int, error) {
|
|
var body io.Reader
|
|
if input != nil {
|
|
encoded, err := json.Marshal(input)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
body = bytes.NewReader(encoded)
|
|
}
|
|
target := *a.base
|
|
target.Path = strings.TrimSuffix(a.base.Path, "/") + path
|
|
request, err := http.NewRequestWithContext(ctx, method, target.String(), body)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
request.Header.Set("Accept", "application/json")
|
|
if a.apiKey != "" {
|
|
request.Header.Set("Authorization", "Bearer "+a.apiKey)
|
|
}
|
|
if input != nil {
|
|
request.Header.Set("Content-Type", "application/json")
|
|
}
|
|
response, err := a.client.Do(request)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
defer response.Body.Close()
|
|
if response.StatusCode != http.StatusOK {
|
|
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 4096))
|
|
return response.StatusCode, nil
|
|
}
|
|
decoder := json.NewDecoder(io.LimitReader(response.Body, 2<<20))
|
|
if err := decoder.Decode(output); err != nil {
|
|
return response.StatusCode, fmt.Errorf("decode provider response: %w", err)
|
|
}
|
|
return response.StatusCode, nil
|
|
}
|
|
|
|
func probeFailure(err error) ProviderProbe {
|
|
finish := cognitive.FinishProviderUnavailable
|
|
if errors.Is(err, context.DeadlineExceeded) {
|
|
finish = cognitive.FinishTimeout
|
|
}
|
|
var network net.Error
|
|
if errors.As(err, &network) && network.Timeout() {
|
|
finish = cognitive.FinishTimeout
|
|
}
|
|
return ProviderProbe{Health: capability.HealthUnavailable, Failure: finish}
|
|
}
|
|
|
|
func mapCompatibleStatus(status int) cognitive.FinishReason {
|
|
switch status {
|
|
case 401, 403:
|
|
return cognitive.FinishAuthenticationFailed
|
|
case 404:
|
|
return cognitive.FinishModelUnavailable
|
|
case 408, 504:
|
|
return cognitive.FinishTimeout
|
|
case 413:
|
|
return cognitive.FinishContextTooLarge
|
|
case 429:
|
|
return cognitive.FinishRateLimited
|
|
default:
|
|
return cognitive.FinishProviderUnavailable
|
|
}
|
|
}
|