iop/apps/node/internal/adapters/ollama/chat.go

package ollama

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"go.uber.org/zap"

	"iop/apps/node/internal/runtime"
)

func (o *Ollama) Execute(ctx context.Context, spec runtime.ExecutionSpec, sink runtime.EventSink) error {
	chatReq, messages, err := o.buildChatRequest(spec)
	if err != nil {
		return err
	}
	if err := emitRuntimeEvent(ctx, sink, spec.RunID, runtime.EventTypeStart, ""); err != nil {
		return err
	}

	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", chatReq.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 {
		return emitHTTPStatusError(ctx, sink, spec.RunID, resp)
	}

	session := ollamaStreamSession{
		runID:    spec.RunID,
		messages: messages,
		sink:     sink,
	}
	return session.consume(ctx, resp.Body)
}

func (o *Ollama) buildChatRequest(spec runtime.ExecutionSpec) (ollamaChatRequest, []ollamaMessage, error) {
	model := strings.TrimSpace(spec.Target)
	if model == "" {
		model = stringInput(spec.Input, "model")
	}
	if model == "" {
		return ollamaChatRequest{}, nil, fmt.Errorf("ollama adapter: target/model is required")
	}
	messages := ollamaMessagesFromInput(spec.Input)
	if len(messages) == 0 {
		return ollamaChatRequest{}, nil, fmt.Errorf("ollama adapter: prompt/messages are required")
	}

	req := ollamaChatRequest{
		Model:    model,
		Messages: messages,
		Stream:   true,
	}
	if options := ollamaOptionsFromInput(spec.Input, o.contextSize); len(options) > 0 {
		req.Options = options
	}
	req.Format = spec.Input["format"]
	req.KeepAlive = spec.Input["keep_alive"]
	req.Think = spec.Input["think"]
	req.Tools = spec.Input["tools"]
	return req, messages, nil
}

func emitHTTPStatusError(
	ctx context.Context,
	sink runtime.EventSink,
	runID string,
	resp *http.Response,
) error {
	msg := readLimited(resp.Body, 4096)
	if msg == "" {
		msg = resp.Status
	}
	_ = emitError(ctx, sink, runID, fmt.Sprintf("ollama returned %s: %s", resp.Status, msg))
	return fmt.Errorf("ollama adapter: non-2xx response: %s", resp.Status)
}

type ollamaStreamSession struct {
	runID        string
	messages     []ollamaMessage
	sink         runtime.EventSink
	outputTokens int
	toolCalls    []any
}

func (s *ollamaStreamSession) consume(ctx context.Context, body io.Reader) error {
	decoder := json.NewDecoder(body)
	for {
		var chunk ollamaChatChunk
		if err := decoder.Decode(&chunk); err != nil {
			if err == io.EOF {
				break
			}
			_ = emitError(ctx, s.sink, s.runID, fmt.Sprintf("ollama stream decode failed: %v", err))
			return fmt.Errorf("ollama adapter: decode stream: %w", err)
		}
		done, err := s.handleChunk(ctx, chunk)
		if err != nil {
			return err
		}
		if done {
			return nil
		}
	}
	_ = emitError(ctx, s.sink, s.runID, "ollama stream ended without done=true")
	return fmt.Errorf("ollama adapter: stream ended without done=true")
}

func (s *ollamaStreamSession) handleChunk(ctx context.Context, chunk ollamaChatChunk) (bool, error) {
	if chunk.Error != "" {
		_ = emitError(ctx, s.sink, s.runID, chunk.Error)
		return false, fmt.Errorf("ollama adapter: %s", chunk.Error)
	}
	if chunk.Message.Thinking != "" {
		if err := emitRuntimeEvent(
			ctx,
			s.sink,
			s.runID,
			runtime.EventTypeReasoningDelta,
			chunk.Message.Thinking,
		); err != nil {
			return false, err
		}
	}
	if chunk.Message.Content != "" {
		s.outputTokens += len(strings.Fields(chunk.Message.Content))
		if err := emitRuntimeEvent(
			ctx,
			s.sink,
			s.runID,
			runtime.EventTypeDelta,
			chunk.Message.Content,
		); err != nil {
			return false, err
		}
	}
	if len(chunk.Message.ToolCalls) > 0 {
		s.toolCalls = append(s.toolCalls, chunk.Message.ToolCalls...)
	}
	if !chunk.Done {
		return false, nil
	}
	return true, s.emitComplete(ctx, chunk)
}

func (s *ollamaStreamSession) emitComplete(ctx context.Context, chunk ollamaChatChunk) error {
	inputTokens := chunk.PromptEvalCount
	if inputTokens == 0 {
		inputTokens = countMessageWords(s.messages)
	}
	if chunk.EvalCount > 0 {
		s.outputTokens = chunk.EvalCount
	}
	var metadata map[string]string
	if len(s.toolCalls) > 0 {
		if encoded, err := json.Marshal(s.toolCalls); err == nil {
			metadata = map[string]string{
				"finish_reason":                "tool_calls",
				runtimeMetadataOpenAIToolCalls: string(encoded),
			}
		}
	}
	return s.sink.Emit(ctx, runtime.RuntimeEvent{
		RunID:   s.runID,
		Type:    runtime.EventTypeComplete,
		Message: "ollama chat complete",
		Usage: &runtime.UsageStats{
			InputTokens:  inputTokens,
			OutputTokens: s.outputTokens,
		},
		Metadata:  metadata,
		Timestamp: time.Now(),
	})
}

func emitRuntimeEvent(
	ctx context.Context,
	sink runtime.EventSink,
	runID string,
	eventType runtime.EventType,
	delta string,
) error {
	return sink.Emit(ctx, runtime.RuntimeEvent{
		RunID:     runID,
		Type:      eventType,
		Delta:     delta,
		Timestamp: time.Now(),
	})
}

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(),
	})
}

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"`
}