tm_back/pkg/ollama/sdk.go

package ollama

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

// SDK provides a high-level interface for the Ollama service
type SDK struct {
	service OllamaService
	client  *Client
	config  *Config
}

// New creates a new Ollama SDK with the provided configuration and logger
func New(config *Config, logger logger.Logger) (*SDK, error) {
	if config == nil {
		config = DefaultConfig()
	}

	client, err := NewClient(config, logger)
	if err != nil {
		return nil, fmt.Errorf("failed to create ollama client: %w", err)
	}

	service := NewService(client)

	return &SDK{
		service: service,
		client:  client,
		config:  config,
	}, nil
}

// NewWithClient creates a new Ollama SDK with a custom HTTP client
func NewWithClient(config *Config, httpClient *http.Client, logger logger.Logger) (*SDK, error) {
	if config == nil {
		config = DefaultConfig()
	}

	client, err := NewClientWithHTTPClient(config, httpClient, logger)
	if err != nil {
		return nil, fmt.Errorf("failed to create ollama client: %w", err)
	}

	service := NewService(client)

	return &SDK{
		service: service,
		client:  client,
		config:  config,
	}, nil
}

// Generate generates text using a specified model and prompt
func (s *SDK) Generate(ctx context.Context, req *GenerateRequest) (*GenerateResponse, error) {
	return s.service.Generate(ctx, req)
}

// GenerateWithOptions generates text with additional options
func (s *SDK) GenerateWithOptions(ctx context.Context, req *GenerateRequest, opts *GenerateOptions) (*GenerateResponse, error) {
	return s.service.GenerateWithOptions(ctx, req, opts)
}

// Chat performs a conversational chat with the model
func (s *SDK) Chat(ctx context.Context, req *ChatRequest) (*ChatResponse, error) {
	return s.service.Chat(ctx, req)
}

// Embed generates embeddings for the given text
func (s *SDK) Embed(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) {
	return s.service.Embed(ctx, req)
}

// PullModel pulls a model from the registry
func (s *SDK) PullModel(ctx context.Context, model string) (*PullResponse, error) {
	return s.service.PullModel(ctx, model)
}

// ListModels lists available models
func (s *SDK) ListModels(ctx context.Context) (*ListModelsResponse, error) {
	return s.service.ListModels(ctx)
}

// DeleteModel deletes a model
func (s *SDK) DeleteModel(ctx context.Context, model string) (*DeleteResponse, error) {
	return s.service.DeleteModel(ctx, model)
}

// ShowModel shows model information
func (s *SDK) ShowModel(ctx context.Context, model string) (*ShowModelResponse, error) {
	return s.service.ShowModel(ctx, model)
}

// Health checks if the Ollama service is available
func (s *SDK) Health(ctx context.Context) error {
	return s.service.Health(ctx)
}

// NewChatBuilder creates a new chat builder for conversational interactions
func (s *SDK) NewChatBuilder() ChatBuilder {
	return s.service.NewChatBuilder()
}

// QuickGenerate is a convenience method for simple text generation
func (s *SDK) QuickGenerate(ctx context.Context, prompt string) (string, error) {
	req := &GenerateRequest{
		Model:  s.config.DefaultModel,
		Prompt: prompt,
	}

	resp, err := s.Generate(ctx, req)
	if err != nil {
		return "", err
	}

	return resp.Response, nil
}

// QuickGenerateWithModel is a convenience method for text generation with a specific model
func (s *SDK) QuickGenerateWithModel(ctx context.Context, model, prompt string) (string, error) {
	req := &GenerateRequest{
		Model:  model,
		Prompt: prompt,
	}

	resp, err := s.Generate(ctx, req)
	if err != nil {
		return "", err
	}

	return resp.Response, nil
}

// QuickChat is a convenience method for simple chat interactions
func (s *SDK) QuickChat(ctx context.Context, userMessage string) (string, error) {
	req := &ChatRequest{
		Model: s.config.DefaultModel,
		Messages: []Message{
			{
				Role:    MessageRoleUser,
				Content: userMessage,
			},
		},
	}

	resp, err := s.Chat(ctx, req)
	if err != nil {
		return "", err
	}

	return resp.Message.Content, nil
}

// QuickChatWithSystem is a convenience method for chat with system prompt
func (s *SDK) QuickChatWithSystem(ctx context.Context, systemPrompt, userMessage string) (string, error) {
	req := &ChatRequest{
		Model: s.config.DefaultModel,
		Messages: []Message{
			{
				Role:    MessageRoleSystem,
				Content: systemPrompt,
			},
			{
				Role:    MessageRoleUser,
				Content: userMessage,
			},
		},
	}

	resp, err := s.Chat(ctx, req)
	if err != nil {
		return "", err
	}

	return resp.Message.Content, nil
}

// QuickEmbed is a convenience method for generating embeddings
func (s *SDK) QuickEmbed(ctx context.Context, text string) ([]float64, error) {
	req := &EmbedRequest{
		Model:  s.config.DefaultModel,
		Prompt: text,
	}

	resp, err := s.Embed(ctx, req)
	if err != nil {
		return nil, err
	}

	return resp.Embedding, nil
}

// StreamGenerate performs streaming text generation
func (s *SDK) StreamGenerate(ctx context.Context, req *GenerateRequest, onChunk func(string) error) error {
	// Force streaming
	stream := true
	req.Stream = &stream

	return s.client.StreamGenerate(ctx, req, func(chunk *StreamChunk) error {
		if chunk.Response != "" {
			return onChunk(chunk.Response)
		}
		return nil
	})
}

// StreamChat performs streaming chat
func (s *SDK) StreamChat(ctx context.Context, req *ChatRequest, onChunk func(string) error) error {
	// Force streaming
	stream := true
	req.Stream = &stream

	return s.client.StreamChat(ctx, req, func(chunk *StreamChunk) error {
		if chunk.Message != nil && chunk.Message.Content != "" {
			return onChunk(chunk.Message.Content)
		}
		return nil
	})
}

// GetService returns the underlying service for advanced usage
func (s *SDK) GetService() OllamaService {
	return s.service
}

// GetConfig returns the current configuration
func (s *SDK) GetConfig() *Config {
	return s.config
}