Assistant/adapters/models/ollama.go

package models

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/url"
	"time"

	"git.cin.su/m/Assistant/contracts/capability"
	"git.cin.su/m/Assistant/contracts/cognitive"
)

type OllamaAdapter struct {
	base   *url.URL
	client *http.Client
}

func NewOllamaAdapter(raw string, transport http.RoundTripper) (*OllamaAdapter, error) {
	base, err := url.Parse(raw)
	if err != nil || base.Scheme != "http" || base.Host == "" || base.Path != "" || base.User != nil || base.RawQuery != "" || base.Fragment != "" {
		return nil, errors.New("invalid Ollama endpoint")
	}
	if transport == nil {
		transport = http.DefaultTransport
	}
	return &OllamaAdapter{base: base, client: &http.Client{Transport: transport, Timeout: 10 * time.Second, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}}, nil
}

type ollamaTags struct {
	Models []struct {
		Name    string `json:"name"`
		Model   string `json:"model"`
		Size    int64  `json:"size"`
		Digest  string `json:"digest"`
		Details struct {
			Format        string `json:"format"`
			Family        string `json:"family"`
			ParameterSize string `json:"parameter_size"`
			Quantization  string `json:"quantization_level"`
			ContextLength int64  `json:"context_length"`
		} `json:"details"`
		Capabilities []string `json:"capabilities"`
	} `json:"models"`
}
type ollamaVersion struct {
	Version string `json:"version"`
}

func (a *OllamaAdapter) Probe(ctx context.Context) ProviderProbe {
	var version ollamaVersion
	status, err := a.request(ctx, http.MethodGet, "/api/version", nil, &version)
	if err != nil {
		return probeFailure(err)
	}
	if status != http.StatusOK || version.Version == "" {
		return ProviderProbe{Health: capability.HealthUnavailable, Failure: mapCompatibleStatus(status)}
	}
	return ProviderProbe{Health: capability.HealthAvailable, Version: version.Version}
}

func (a *OllamaAdapter) Discover(ctx context.Context) ([]DiscoveredModel, ProviderProbe) {
	probe := a.Probe(ctx)
	if probe.Health != capability.HealthAvailable {
		return nil, probe
	}
	var tags ollamaTags
	status, err := a.request(ctx, http.MethodGet, "/api/tags", nil, &tags)
	if err != nil {
		return nil, probeFailure(err)
	}
	if status != http.StatusOK {
		return nil, ProviderProbe{Health: capability.HealthUnavailable, Failure: mapCompatibleStatus(status)}
	}
	result := make([]DiscoveredModel, 0, len(tags.Models))
	for _, item := range tags.Models {
		if item.Name == "" || item.Digest == "" {
			continue
		}
		capabilities := []capability.ModelCapability{capability.CapabilityUnknown}
		for _, declared := range item.Capabilities {
			if declared == "completion" {
				capabilities = []capability.ModelCapability{capability.CapabilityAdvisory}
				break
			}
		}
		result = append(result, DiscoveredModel{ProviderModelID: capability.ProviderModelID(item.Name), Fingerprint: capability.ModelFingerprint(item.Digest), MaxContextBytes: item.Details.ContextLength * 4, Modalities: []capability.ModelModality{capability.ModalityText}, Capabilities: capabilities, SizeBytes: item.Size, Runtime: "ollama/" + probe.Version, Quantization: item.Details.Quantization})
	}
	return result, probe
}

type ollamaChatRequest struct {
	Model    string              `json:"model"`
	Messages []compatibleMessage `json:"messages"`
	Stream   bool                `json:"stream"`
	Think    bool                `json:"think"`
	Options  struct {
		Temperature int   `json:"temperature"`
		NumCtx      int64 `json:"num_ctx,omitempty"`
	} `json:"options"`
}
type ollamaChatResponse struct {
	Message         compatibleMessage `json:"message"`
	Done            bool              `json:"done"`
	DoneReason      string            `json:"done_reason"`
	PromptEvalCount int64             `json:"prompt_eval_count"`
	EvalCount       int64             `json:"eval_count"`
}

func (a *OllamaAdapter) 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 := ollamaChatRequest{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}}, Stream: false, Think: false}
	var result ollamaChatResponse
	status, err := a.request(ctx, http.MethodPost, "/api/chat", payload, &result)
	if err != nil {
		return advisoryFailure(request, probeFailure(err).Failure), nil
	}
	if status != http.StatusOK {
		return advisoryFailure(request, mapCompatibleStatus(status)), nil
	}
	if !result.Done || result.Message.Content == "" {
		return advisoryFailure(request, cognitive.FinishInvalidProviderResponse), nil
	}
	response := cognitive.AdvisoryResponse{Version: request.Version, RequestID: request.RequestID, ProviderBindingID: request.ProviderBindingID, ProviderBindingRevision: request.ProviderBindingRevision, AdvisoryText: result.Message.Content, FinishReason: cognitive.FinishCompleted, Usage: &cognitive.ProviderUsage{InputUnits: result.PromptEvalCount, OutputUnits: result.EvalCount}}
	if ValidateAdvisoryResponse(request, response) != nil {
		return advisoryFailure(request, cognitive.FinishInvalidProviderResponse), nil
	}
	return response, nil
}

func (a *OllamaAdapter) request(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 = path
	request, err := http.NewRequestWithContext(ctx, method, target.String(), body)
	if err != nil {
		return 0, err
	}
	request.Header.Set("Accept", "application/json")
	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, err
	}
	return response.StatusCode, nil
}