500 lines
14 KiB
Go
500 lines
14 KiB
Go
// Package ollama provides an Adapter that calls a local Ollama server.
|
|
package ollama
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
"strings"
|
|
"time"
|
|
|
|
"go.uber.org/zap"
|
|
|
|
"iop/apps/node/internal/runtime"
|
|
"iop/packages/go/config"
|
|
)
|
|
|
|
const Name = "ollama"
|
|
|
|
const runtimeMetadataOpenAIToolCalls = "openai_tool_calls"
|
|
|
|
type Ollama struct {
|
|
instanceName string
|
|
baseURL string
|
|
contextSize int
|
|
capacity int
|
|
maxQueue int
|
|
queueTimeoutMS int
|
|
requestTimeoutMS int
|
|
client *http.Client
|
|
logger *zap.Logger
|
|
}
|
|
|
|
// New creates an Ollama adapter. The optional instanceName is the registry
|
|
// instance key used to disambiguate multiple Ollama instances on the same node.
|
|
func New(cfg config.OllamaConf, logger *zap.Logger, instanceName ...string) *Ollama {
|
|
baseURL := strings.TrimRight(cfg.BaseURL, "/")
|
|
if baseURL == "" {
|
|
baseURL = "http://localhost:11434"
|
|
}
|
|
name := Name
|
|
if len(instanceName) > 0 && instanceName[0] != "" {
|
|
name = instanceName[0]
|
|
}
|
|
return &Ollama{
|
|
instanceName: name,
|
|
baseURL: baseURL,
|
|
contextSize: cfg.ContextSize,
|
|
capacity: cfg.Capacity,
|
|
maxQueue: cfg.MaxQueue,
|
|
queueTimeoutMS: cfg.QueueTimeoutMS,
|
|
requestTimeoutMS: cfg.RequestTimeoutMS,
|
|
client: &http.Client{},
|
|
logger: logger,
|
|
}
|
|
}
|
|
|
|
func (o *Ollama) Name() string { return Name }
|
|
|
|
func (o *Ollama) Capabilities(ctx context.Context) (runtime.Capabilities, error) {
|
|
probeCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
|
|
defer cancel()
|
|
targets, err := o.fetchTargets(probeCtx)
|
|
status := runtime.ProviderStatusAvailable
|
|
if err != nil {
|
|
status = runtime.ProviderStatusUnavailable
|
|
}
|
|
return runtime.Capabilities{
|
|
AdapterName: Name,
|
|
InstanceKey: o.instanceName,
|
|
Targets: targets,
|
|
MaxConcurrency: effectiveCapacity(o.capacity, 4),
|
|
MaxQueue: o.maxQueue,
|
|
QueueTimeoutMS: o.queueTimeoutMS,
|
|
RequestTimeoutMS: o.requestTimeoutMS,
|
|
ProviderStatus: runtime.NormalizeProviderStatus(status),
|
|
}, nil
|
|
}
|
|
|
|
func effectiveCapacity(capacity, defaultVal int) int {
|
|
if capacity <= 0 {
|
|
return defaultVal
|
|
}
|
|
return capacity
|
|
}
|
|
|
|
func (o *Ollama) Execute(ctx context.Context, spec runtime.ExecutionSpec, sink runtime.EventSink) error {
|
|
model := strings.TrimSpace(spec.Target)
|
|
if model == "" {
|
|
model = stringInput(spec.Input, "model")
|
|
}
|
|
if model == "" {
|
|
return fmt.Errorf("ollama adapter: target/model is required")
|
|
}
|
|
messages := ollamaMessagesFromInput(spec.Input)
|
|
if len(messages) == 0 {
|
|
return fmt.Errorf("ollama adapter: prompt/messages are required")
|
|
}
|
|
|
|
if err := sink.Emit(ctx, runtime.RuntimeEvent{
|
|
RunID: spec.RunID,
|
|
Type: runtime.EventTypeStart,
|
|
Timestamp: time.Now(),
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
|
|
chatReq := ollamaChatRequest{
|
|
Model: model,
|
|
Messages: messages,
|
|
Stream: true,
|
|
}
|
|
if options := ollamaOptionsFromInput(spec.Input, o.contextSize); len(options) > 0 {
|
|
chatReq.Options = options
|
|
}
|
|
if v, ok := spec.Input["format"]; ok {
|
|
chatReq.Format = v
|
|
}
|
|
if v, ok := spec.Input["keep_alive"]; ok {
|
|
chatReq.KeepAlive = v
|
|
}
|
|
if v, ok := spec.Input["think"]; ok {
|
|
chatReq.Think = v
|
|
}
|
|
if v, ok := spec.Input["tools"]; ok {
|
|
chatReq.Tools = v
|
|
}
|
|
body, err := json.Marshal(chatReq)
|
|
if err != nil {
|
|
return fmt.Errorf("ollama adapter: marshal request: %w", err)
|
|
}
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, joinURL(o.baseURL, "/api/chat"), bytes.NewReader(body))
|
|
if err != nil {
|
|
return fmt.Errorf("ollama adapter: build request: %w", err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
|
|
o.logger.Info("ollama adapter executing",
|
|
zap.String("run_id", spec.RunID),
|
|
zap.String("target", model),
|
|
zap.String("base_url", o.baseURL),
|
|
zap.Int("context_size", o.contextSize),
|
|
)
|
|
resp, err := o.client.Do(req)
|
|
if err != nil {
|
|
_ = emitError(ctx, sink, spec.RunID, fmt.Sprintf("ollama request failed: %v", err))
|
|
return fmt.Errorf("ollama adapter: request: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
msg := readLimited(resp.Body, 4096)
|
|
if msg == "" {
|
|
msg = resp.Status
|
|
}
|
|
_ = emitError(ctx, sink, spec.RunID, fmt.Sprintf("ollama returned %s: %s", resp.Status, msg))
|
|
return fmt.Errorf("ollama adapter: non-2xx response: %s", resp.Status)
|
|
}
|
|
|
|
decoder := json.NewDecoder(resp.Body)
|
|
outputTokens := 0
|
|
var toolCalls []any
|
|
for {
|
|
var chunk ollamaChatChunk
|
|
if err := decoder.Decode(&chunk); err != nil {
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
_ = emitError(ctx, sink, spec.RunID, fmt.Sprintf("ollama stream decode failed: %v", err))
|
|
return fmt.Errorf("ollama adapter: decode stream: %w", err)
|
|
}
|
|
if chunk.Error != "" {
|
|
_ = emitError(ctx, sink, spec.RunID, chunk.Error)
|
|
return fmt.Errorf("ollama adapter: %s", chunk.Error)
|
|
}
|
|
if chunk.Message.Thinking != "" {
|
|
if err := sink.Emit(ctx, runtime.RuntimeEvent{
|
|
RunID: spec.RunID,
|
|
Type: runtime.EventTypeReasoningDelta,
|
|
Delta: chunk.Message.Thinking,
|
|
Timestamp: time.Now(),
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
if chunk.Message.Content != "" {
|
|
outputTokens += len(strings.Fields(chunk.Message.Content))
|
|
if err := sink.Emit(ctx, runtime.RuntimeEvent{
|
|
RunID: spec.RunID,
|
|
Type: runtime.EventTypeDelta,
|
|
Delta: chunk.Message.Content,
|
|
Timestamp: time.Now(),
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
if len(chunk.Message.ToolCalls) > 0 {
|
|
toolCalls = append(toolCalls, chunk.Message.ToolCalls...)
|
|
}
|
|
if chunk.Done {
|
|
inputTokens := chunk.PromptEvalCount
|
|
if inputTokens == 0 {
|
|
inputTokens = countMessageWords(messages)
|
|
}
|
|
if chunk.EvalCount > 0 {
|
|
outputTokens = chunk.EvalCount
|
|
}
|
|
var metadata map[string]string
|
|
if len(toolCalls) > 0 {
|
|
if encoded, err := json.Marshal(toolCalls); err == nil {
|
|
metadata = map[string]string{
|
|
"finish_reason": "tool_calls",
|
|
runtimeMetadataOpenAIToolCalls: string(encoded),
|
|
}
|
|
}
|
|
}
|
|
return sink.Emit(ctx, runtime.RuntimeEvent{
|
|
RunID: spec.RunID,
|
|
Type: runtime.EventTypeComplete,
|
|
Message: "ollama chat complete",
|
|
Usage: &runtime.UsageStats{
|
|
InputTokens: inputTokens,
|
|
OutputTokens: outputTokens,
|
|
},
|
|
Metadata: metadata,
|
|
Timestamp: time.Now(),
|
|
})
|
|
}
|
|
}
|
|
_ = emitError(ctx, sink, spec.RunID, "ollama stream ended without done=true")
|
|
return fmt.Errorf("ollama adapter: stream ended without done=true")
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
found := false
|
|
for _, t := range targets {
|
|
if t == target || (strings.Index(target, ":") == -1 && t == target+":latest") {
|
|
found = true
|
|
break
|
|
}
|
|
}
|
|
|
|
if found {
|
|
result.Status = runtime.NormalizeProviderStatus(runtime.ProviderStatusAvailable)
|
|
} else {
|
|
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
|
|
}
|
|
|
|
func (o *Ollama) HandleCommand(ctx context.Context, req runtime.CommandRequest) (runtime.CommandResponse, error) {
|
|
if req.Type != runtime.CommandTypeOllamaAPI {
|
|
return runtime.CommandResponse{}, fmt.Errorf("ollama adapter: unsupported command %q", req.Type)
|
|
}
|
|
method := strings.ToUpper(strings.TrimSpace(req.Metadata["ollama_method"]))
|
|
if method == "" {
|
|
method = http.MethodGet
|
|
}
|
|
if method != http.MethodGet && method != http.MethodPost && method != http.MethodDelete {
|
|
return runtime.CommandResponse{}, fmt.Errorf("ollama adapter: unsupported method %q", method)
|
|
}
|
|
path := strings.TrimSpace(req.Metadata["ollama_path"])
|
|
if !strings.HasPrefix(path, "/api/") {
|
|
return runtime.CommandResponse{}, fmt.Errorf("ollama adapter: unsupported path %q", path)
|
|
}
|
|
body := req.Metadata["ollama_body"]
|
|
var reader io.Reader
|
|
if body != "" {
|
|
reader = strings.NewReader(body)
|
|
}
|
|
httpReq, err := http.NewRequestWithContext(ctx, method, joinURLRequestURI(o.baseURL, path), reader)
|
|
if err != nil {
|
|
return runtime.CommandResponse{}, fmt.Errorf("ollama adapter: build passthrough request: %w", err)
|
|
}
|
|
if body != "" {
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
|
}
|
|
resp, err := o.client.Do(httpReq)
|
|
if err != nil {
|
|
return runtime.CommandResponse{}, fmt.Errorf("ollama adapter: passthrough request: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
respBody := readLimited(resp.Body, 16<<20)
|
|
return runtime.CommandResponse{
|
|
RequestID: req.RequestID,
|
|
Type: req.Type,
|
|
Adapter: req.Adapter,
|
|
Target: req.Target,
|
|
SessionID: req.SessionID,
|
|
Result: map[string]string{
|
|
"status_code": fmt.Sprintf("%d", resp.StatusCode),
|
|
"content_type": resp.Header.Get("Content-Type"),
|
|
"body": respBody,
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
func ollamaMessagesFromInput(input map[string]any) []ollamaMessage {
|
|
if raw, ok := input["messages"].([]any); ok {
|
|
out := make([]ollamaMessage, 0, len(raw))
|
|
for _, item := range raw {
|
|
m, ok := item.(map[string]any)
|
|
if !ok {
|
|
continue
|
|
}
|
|
role, _ := m["role"].(string)
|
|
content, _ := m["content"].(string)
|
|
thinking, _ := m["thinking"].(string)
|
|
toolName, _ := m["tool_name"].(string)
|
|
role = strings.TrimSpace(role)
|
|
content = strings.TrimSpace(content)
|
|
if role == "" || (content == "" && thinking == "") {
|
|
continue
|
|
}
|
|
out = append(out, ollamaMessage{
|
|
Role: role,
|
|
Content: content,
|
|
Thinking: thinking,
|
|
Images: anySlice(m["images"]),
|
|
ToolCalls: anySlice(m["tool_calls"]),
|
|
ToolName: toolName,
|
|
})
|
|
}
|
|
if len(out) > 0 {
|
|
return out
|
|
}
|
|
}
|
|
prompt := strings.TrimSpace(stringInput(input, "prompt"))
|
|
if prompt == "" {
|
|
return nil
|
|
}
|
|
return []ollamaMessage{{Role: "user", Content: prompt}}
|
|
}
|
|
|
|
func ollamaOptionsFromInput(input map[string]any, contextSize int) map[string]any {
|
|
options := map[string]any{}
|
|
if raw, ok := input["options"].(map[string]any); ok {
|
|
for k, v := range raw {
|
|
options[k] = v
|
|
}
|
|
}
|
|
if contextSize > 0 {
|
|
options["num_ctx"] = contextSize // Edge-owned: always overrides request
|
|
}
|
|
if len(options) == 0 {
|
|
return nil
|
|
}
|
|
return options
|
|
}
|
|
|
|
func anySlice(v any) []any {
|
|
if items, ok := v.([]any); ok {
|
|
return items
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func stringInput(input map[string]any, key string) string {
|
|
if input == nil {
|
|
return ""
|
|
}
|
|
if v, ok := input[key].(string); ok {
|
|
return v
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func countMessageWords(messages []ollamaMessage) int {
|
|
n := 0
|
|
for _, msg := range messages {
|
|
n += len(strings.Fields(msg.Content))
|
|
}
|
|
return n
|
|
}
|
|
|
|
func emitError(ctx context.Context, sink runtime.EventSink, runID, msg string) error {
|
|
return sink.Emit(ctx, runtime.RuntimeEvent{
|
|
RunID: runID,
|
|
Type: runtime.EventTypeError,
|
|
Error: msg,
|
|
Timestamp: time.Now(),
|
|
})
|
|
}
|
|
|
|
func joinURL(baseURL, path string) string {
|
|
u, err := url.Parse(baseURL)
|
|
if err != nil {
|
|
return strings.TrimRight(baseURL, "/") + path
|
|
}
|
|
u.Path = strings.TrimRight(u.Path, "/") + path
|
|
return u.String()
|
|
}
|
|
|
|
func joinURLRequestURI(baseURL, requestURI string) string {
|
|
u, err := url.Parse(baseURL)
|
|
if err != nil {
|
|
return strings.TrimRight(baseURL, "/") + requestURI
|
|
}
|
|
rel, err := url.Parse(requestURI)
|
|
if err != nil {
|
|
u.Path = strings.TrimRight(u.Path, "/") + requestURI
|
|
return u.String()
|
|
}
|
|
u.Path = strings.TrimRight(u.Path, "/") + rel.Path
|
|
u.RawQuery = rel.RawQuery
|
|
return u.String()
|
|
}
|
|
|
|
func readLimited(r io.Reader, limit int64) string {
|
|
b, _ := io.ReadAll(io.LimitReader(r, limit))
|
|
return strings.TrimSpace(string(b))
|
|
}
|
|
|
|
type ollamaChatRequest struct {
|
|
Model string `json:"model"`
|
|
Messages []ollamaMessage `json:"messages"`
|
|
Stream bool `json:"stream"`
|
|
Options map[string]any `json:"options,omitempty"`
|
|
Format any `json:"format,omitempty"`
|
|
KeepAlive any `json:"keep_alive,omitempty"`
|
|
Think any `json:"think,omitempty"`
|
|
Tools any `json:"tools,omitempty"`
|
|
}
|
|
|
|
type ollamaMessage struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
Thinking string `json:"thinking,omitempty"`
|
|
Images []any `json:"images,omitempty"`
|
|
ToolCalls []any `json:"tool_calls,omitempty"`
|
|
ToolName string `json:"tool_name,omitempty"`
|
|
}
|
|
|
|
type ollamaChatChunk struct {
|
|
Message ollamaMessage `json:"message"`
|
|
Done bool `json:"done"`
|
|
PromptEvalCount int `json:"prompt_eval_count"`
|
|
EvalCount int `json:"eval_count"`
|
|
Error string `json:"error"`
|
|
}
|
|
|
|
type ollamaTagsResponse struct {
|
|
Models []struct {
|
|
Name string `json:"name"`
|
|
} `json:"models"`
|
|
}
|