tm_back/pkg/ollama/service.go

package ollama

import (
	"context"
	"fmt"
	"tm/pkg/logger"
)

// Service implements the OllamaService interface
type Service struct {
	client *Client
	config *Config
	logger logger.Logger
}

// NewService creates a new Ollama service
func NewService(client *Client) OllamaService {
	return &Service{
		client: client,
		config: client.config,
		logger: client.logger,
	}
}

// Generate generates text using a specified model and prompt
func (s *Service) Generate(ctx context.Context, req *GenerateRequest) (*GenerateResponse, error) {
	if s.logger != nil {
		s.logger.Info("Generating text", map[string]interface{}{
			"model":         req.Model,
			"prompt_length": len(req.Prompt),
		})
	}

	// Set defaults
	if req.Model == "" {
		req.Model = s.config.DefaultModel
	}

	var resp GenerateResponse
	err := s.client.Post(ctx, "/api/generate", req, &resp)
	if err != nil {
		if s.logger != nil {
			s.logger.Error("Text generation failed", map[string]interface{}{
				"model": req.Model,
				"error": err.Error(),
			})
		}
		return nil, ErrGenerationFailed(err.Error())
	}

	if s.logger != nil {
		s.logger.Info("Text generation completed", map[string]interface{}{
			"model":           req.Model,
			"response_length": len(resp.Response),
			"eval_count":      resp.EvalCount,
		})
	}

	return &resp, nil
}

// GenerateWithOptions generates text with additional options
func (s *Service) GenerateWithOptions(ctx context.Context, req *GenerateRequest, opts *GenerateOptions) (*GenerateResponse, error) {
	if s.logger != nil {
		s.logger.Info("Generating text with options", map[string]interface{}{
			"model":         req.Model,
			"prompt_length": len(req.Prompt),
		})
	}

	// Apply options to request
	if opts != nil {
		if req.Options == nil {
			req.Options = make(map[string]interface{})
		}

		if opts.Temperature != nil {
			req.Options["temperature"] = *opts.Temperature
		}
		if opts.TopP != nil {
			req.Options["top_p"] = *opts.TopP
		}
		if opts.TopK != nil {
			req.Options["top_k"] = *opts.TopK
		}
		if opts.RepeatPenalty != nil {
			req.Options["repeat_penalty"] = *opts.RepeatPenalty
		}
		if opts.Seed != nil {
			req.Options["seed"] = *opts.Seed
		}
		if opts.Stop != nil {
			req.Options["stop"] = opts.Stop
		}
		if opts.Format != "" {
			req.Format = opts.Format
		}

		// Merge additional options
		for k, v := range opts.Options {
			req.Options[k] = v
		}
	}

	return s.Generate(ctx, req)
}

// Chat performs a conversational chat with the model
func (s *Service) Chat(ctx context.Context, req *ChatRequest) (*ChatResponse, error) {
	if s.logger != nil {
		s.logger.Info("Starting chat", map[string]interface{}{
			"model":         req.Model,
			"message_count": len(req.Messages),
		})
	}

	// Set defaults
	if req.Model == "" {
		req.Model = s.config.DefaultModel
	}

	var resp ChatResponse
	err := s.client.Post(ctx, "/api/chat", req, &resp)
	if err != nil {
		if s.logger != nil {
			s.logger.Error("Chat failed", map[string]interface{}{
				"model": req.Model,
				"error": err.Error(),
			})
		}
		return nil, ErrChatFailed(err.Error())
	}

	if s.logger != nil {
		s.logger.Info("Chat completed", map[string]interface{}{
			"model":           req.Model,
			"response_length": len(resp.Message.Content),
			"eval_count":      resp.EvalCount,
		})
	}

	return &resp, nil
}

// Embed generates embeddings for the given text
func (s *Service) Embed(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) {
	if s.logger != nil {
		s.logger.Info("Generating embeddings", map[string]interface{}{
			"model":       req.Model,
			"text_length": len(req.Prompt),
		})
	}

	var resp EmbedResponse
	err := s.client.Post(ctx, "/api/embeddings", req, &resp)
	if err != nil {
		if s.logger != nil {
			s.logger.Error("Embedding generation failed", map[string]interface{}{
				"model": req.Model,
				"error": err.Error(),
			})
		}
		return nil, ErrEmbeddingFailed(err.Error())
	}

	if s.logger != nil {
		s.logger.Info("Embeddings generated", map[string]interface{}{
			"model":               req.Model,
			"embedding_dimension": len(resp.Embedding),
		})
	}

	return &resp, nil
}

// PullModel pulls a model from the registry
func (s *Service) PullModel(ctx context.Context, model string) (*PullResponse, error) {
	if s.logger != nil {
		s.logger.Info("Pulling model", map[string]interface{}{
			"model": model,
		})
	}

	req := &PullRequest{
		Name: model,
	}

	var resp PullResponse
	err := s.client.Post(ctx, "/api/pull", req, &resp)
	if err != nil {
		if s.logger != nil {
			s.logger.Error("Model pull failed", map[string]interface{}{
				"model": model,
				"error": err.Error(),
			})
		}
		return nil, ErrModelPullFailed(model, err.Error())
	}

	if s.logger != nil {
		s.logger.Info("Model pulled successfully", map[string]interface{}{
			"model":  model,
			"status": resp.Status,
		})
	}

	return &resp, nil
}

// ListModels lists available models
func (s *Service) ListModels(ctx context.Context) (*ListModelsResponse, error) {
	if s.logger != nil {
		s.logger.Debug("Listing models", map[string]interface{}{})
	}

	var resp ListModelsResponse
	err := s.client.Get(ctx, "/api/tags", &resp)
	if err != nil {
		if s.logger != nil {
			s.logger.Error("Failed to list models", map[string]interface{}{
				"error": err.Error(),
			})
		}
		return nil, fmt.Errorf("failed to list models: %w", err)
	}

	if s.logger != nil {
		s.logger.Debug("Models listed", map[string]interface{}{
			"model_count": len(resp.Models),
		})
	}

	return &resp, nil
}

// DeleteModel deletes a model
func (s *Service) DeleteModel(ctx context.Context, model string) (*DeleteResponse, error) {
	if s.logger != nil {
		s.logger.Info("Deleting model", map[string]interface{}{
			"model": model,
		})
	}

	req := &DeleteRequest{
		Name: model,
	}

	var resp DeleteResponse
	err := s.client.Delete(ctx, "/api/delete", req, &resp)
	if err != nil {
		if s.logger != nil {
			s.logger.Error("Model deletion failed", map[string]interface{}{
				"model": model,
				"error": err.Error(),
			})
		}
		return nil, ErrModelDeleteFailed(model, err.Error())
	}

	if s.logger != nil {
		s.logger.Info("Model deleted successfully", map[string]interface{}{
			"model":  model,
			"status": resp.Status,
		})
	}

	return &resp, nil
}

// ShowModel shows model information
func (s *Service) ShowModel(ctx context.Context, model string) (*ShowModelResponse, error) {
	if s.logger != nil {
		s.logger.Debug("Showing model information", map[string]interface{}{
			"model": model,
		})
	}

	req := &ShowModelRequest{
		Name: model,
	}

	var resp ShowModelResponse
	err := s.client.Post(ctx, "/api/show", req, &resp)
	if err != nil {
		if s.logger != nil {
			s.logger.Error("Failed to show model", map[string]interface{}{
				"model": model,
				"error": err.Error(),
			})
		}
		return nil, fmt.Errorf("failed to show model %s: %w", model, err)
	}

	if s.logger != nil {
		s.logger.Debug("Model information retrieved", map[string]interface{}{
			"model": model,
		})
	}

	return &resp, nil
}

// Health checks if the Ollama service is available
func (s *Service) Health(ctx context.Context) error {
	if s.logger != nil {
		s.logger.Debug("Checking Ollama health", map[string]interface{}{})
	}

	var resp HealthResponse
	err := s.client.Get(ctx, "/api/version", &resp)
	if err != nil {
		if s.logger != nil {
			s.logger.Error("Health check failed", map[string]interface{}{
				"error": err.Error(),
			})
		}
		return ErrConnectionFailed(err.Error())
	}

	if s.logger != nil {
		s.logger.Debug("Health check passed", map[string]interface{}{})
	}

	return nil
}

// NewChatBuilder creates a new chat builder for conversational interactions
func (s *Service) NewChatBuilder() ChatBuilder {
	return &chatBuilder{
		service: s,
		request: &ChatRequest{
			Model:    s.config.DefaultModel,
			Messages: make([]Message, 0),
			Options: ChatOptions{
				Temperature: &s.config.DefaultTemperature,
				TopP:        &s.config.DefaultTopP,
			},
		},
	}
}

// chatBuilder implements the ChatBuilder interface
type chatBuilder struct {
	service OllamaService
	request *ChatRequest
}

func (b *chatBuilder) SetModel(model string) ChatBuilder {
	b.request.Model = model
	return b
}

func (b *chatBuilder) SetSystem(system string) ChatBuilder {
	// Add or update system message
	for i, msg := range b.request.Messages {
		if msg.Role == MessageRoleSystem {
			b.request.Messages[i].Content = system
			return b
		}
	}

	// Add new system message at the beginning
	systemMsg := Message{
		Role:    MessageRoleSystem,
		Content: system,
	}
	b.request.Messages = append([]Message{systemMsg}, b.request.Messages...)
	return b
}

func (b *chatBuilder) AddMessage(role MessageRole, content string) ChatBuilder {
	msg := Message{
		Role:    role,
		Content: content,
	}
	b.request.Messages = append(b.request.Messages, msg)
	return b
}

func (b *chatBuilder) SetTemperature(temperature float64) ChatBuilder {
	b.request.Options.Temperature = &temperature
	return b
}

func (b *chatBuilder) SetTopP(topP float64) ChatBuilder {
	b.request.Options.TopP = &topP
	return b
}

func (b *chatBuilder) SetMaxTokens(maxTokens int) ChatBuilder {
	b.request.Options.NumPredict = &maxTokens
	return b
}

func (b *chatBuilder) SetStream(stream bool) ChatBuilder {
	b.request.Stream = &stream
	return b
}

func (b *chatBuilder) SetFormat(format string) ChatBuilder {
	b.request.Format = format
	return b
}

func (b *chatBuilder) SetOptions(options map[string]interface{}) ChatBuilder {
	if b.request.Options.Options == nil {
		b.request.Options.Options = make(map[string]interface{})
	}
	for k, v := range options {
		b.request.Options.Options[k] = v
	}
	return b
}

func (b *chatBuilder) Build() (*ChatRequest, error) {
	if len(b.request.Messages) == 0 {
		return nil, ErrInvalidRequest("at least one message is required")
	}
	return b.request, nil
}

func (b *chatBuilder) Send(ctx context.Context) (*ChatResponse, error) {
	req, err := b.Build()
	if err != nil {
		return nil, err
	}
	return b.service.Chat(ctx, req)
}

func (b *chatBuilder) Continue(ctx context.Context, userMessage string) (*ChatResponse, error) {
	// Add user message
	b.AddMessage(MessageRoleUser, userMessage)

	// Send chat request
	resp, err := b.Send(ctx)
	if err != nil {
		return nil, err
	}

	// Add assistant response to conversation history
	b.AddMessage(MessageRoleAssistant, resp.Message.Content)

	return resp, nil
}