76 lines
2.1 KiB
Go
76 lines
2.1 KiB
Go
package ollama
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"iop/apps/node/internal/runtime"
|
|
)
|
|
|
|
func (o *Ollama) ProbeProvider(ctx context.Context, target string) (runtime.ProviderProbeResult, error) {
|
|
probeCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
|
|
defer cancel()
|
|
|
|
targets, err := o.fetchTargets(probeCtx)
|
|
result := runtime.ProviderProbeResult{
|
|
AdapterName: Name,
|
|
InstanceKey: o.instanceName,
|
|
Target: target,
|
|
}
|
|
if err != nil {
|
|
result.Status = runtime.NormalizeProviderStatus(runtime.ProviderStatusUnavailable)
|
|
result.Detail = err.Error()
|
|
return result, nil
|
|
}
|
|
|
|
result.Targets = targets
|
|
if target == "" {
|
|
result.Status = runtime.NormalizeProviderStatus(runtime.ProviderStatusAvailable)
|
|
return result, nil
|
|
}
|
|
for _, candidate := range targets {
|
|
if candidate == target || (!strings.Contains(target, ":") && candidate == target+":latest") {
|
|
result.Status = runtime.NormalizeProviderStatus(runtime.ProviderStatusAvailable)
|
|
return result, nil
|
|
}
|
|
}
|
|
result.Status = runtime.NormalizeProviderStatus(runtime.ProviderStatusUnavailable)
|
|
result.Detail = fmt.Sprintf("target model %q not found in provider models", target)
|
|
return result, nil
|
|
}
|
|
|
|
func (o *Ollama) fetchTargets(ctx context.Context) ([]string, error) {
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, joinURL(o.baseURL, "/api/tags"), nil)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("build request: %w", err)
|
|
}
|
|
resp, err := o.client.Do(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("request failed: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
return nil, fmt.Errorf("status code %d", resp.StatusCode)
|
|
}
|
|
var tags ollamaTagsResponse
|
|
if err := json.NewDecoder(resp.Body).Decode(&tags); err != nil {
|
|
return nil, fmt.Errorf("decode response: %w", err)
|
|
}
|
|
out := make([]string, 0, len(tags.Models))
|
|
for _, model := range tags.Models {
|
|
if model.Name != "" {
|
|
out = append(out, model.Name)
|
|
}
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
type ollamaTagsResponse struct {
|
|
Models []struct {
|
|
Name string `json:"name"`
|
|
} `json:"models"`
|
|
}
|