diff --git a/pkg/ollama/README.md b/pkg/ollama/README.md new file mode 100644 index 0000000..ad86132 --- /dev/null +++ b/pkg/ollama/README.md @@ -0,0 +1,394 @@ +# Ollama SDK + +A comprehensive Go SDK for interacting with Ollama AI models. This SDK provides a clean, type-safe interface for text generation, chat conversations, embeddings, and model management. + +## Features + +- **Text Generation**: Generate text using various AI models +- **Chat Conversations**: Interactive chat with support for system prompts and conversation history +- **Embeddings**: Generate embeddings for text similarity and search +- **Model Management**: Pull, list, delete, and inspect models +- **Streaming Support**: Real-time streaming of generated text and chat responses +- **Automatic Retries**: Configurable retry mechanism with exponential backoff +- **Structured Logging**: Built-in logging with structured fields +- **Fluent API**: Builder pattern for constructing complex chat conversations +- **Type Safety**: Full type safety with comprehensive validation +- **Error Handling**: Detailed error types and proper error wrapping +- **Configuration**: Flexible configuration via environment variables or code +- **Health Checks**: Built-in health checking capabilities + +## Installation + +The SDK is already part of your `tm` module. Import it as: +```go +import "tm/pkg/ollama" +``` + +## Quick Start + +### Basic Setup + +```go +// 1. Configuration +config := ollama.DefaultConfig() +config.BaseURL = "http://localhost:11434" // Your Ollama server URL +config.DefaultModel = "llama2" // Default model to use + +// 2. Logger (using project's logger) +loggerConfig := &logger.Config{Level: "info", Format: "json", Output: "stdout"} +projectLogger := logger.NewLogger(loggerConfig) + +// 3. Create SDK +sdk, err := ollama.New(config, projectLogger) +if err != nil { + log.Fatal("Failed to create Ollama SDK:", err) +} +``` + +### Environment Variables + +Set environment variables for automatic configuration: + +```bash +export OLLAMA_BASE_URL="http://localhost:11434" +export OLLAMA_TIMEOUT="300s" +export OLLAMA_RETRY_ATTEMPTS="3" +export OLLAMA_RETRY_DELAY="2s" +export OLLAMA_ENABLE_LOGGING="true" +export OLLAMA_DEFAULT_MODEL="llama2" +export OLLAMA_DEFAULT_TEMPERATURE="0.7" +export OLLAMA_DEFAULT_TOP_P="0.9" +export OLLAMA_DEFAULT_MAX_TOKENS="1000" +``` + +## API Reference + +### Configuration + +```go +type Config struct { + BaseURL string `env:"OLLAMA_BASE_URL"` + Timeout time.Duration `env:"OLLAMA_TIMEOUT"` + RetryAttempts int `env:"OLLAMA_RETRY_ATTEMPTS"` + RetryDelay time.Duration `env:"OLLAMA_RETRY_DELAY"` + EnableLogging bool `env:"OLLAMA_ENABLE_LOGGING"` + UserAgent string `env:"OLLAMA_USER_AGENT"` + DefaultModel string `env:"OLLAMA_DEFAULT_MODEL"` + DefaultTemperature float64 `env:"OLLAMA_DEFAULT_TEMPERATURE"` + DefaultTopP float64 `env:"OLLAMA_DEFAULT_TOP_P"` + DefaultMaxTokens int `env:"OLLAMA_DEFAULT_MAX_TOKENS"` + EnableStreaming bool `env:"OLLAMA_ENABLE_STREAMING"` +} +``` + +### Creating SDK Instance + +```go +// With default configuration +sdk, err := ollama.New(nil, logger) + +// With custom configuration +config := &ollama.Config{ + BaseURL: "http://localhost:11434", + Timeout: 300 * time.Second, + RetryAttempts: 3, + RetryDelay: 2 * time.Second, + EnableLogging: true, + DefaultModel: "llama2", + DefaultTemperature: 0.7, + DefaultTopP: 0.9, + DefaultMaxTokens: 1000, +} +sdk, err := ollama.New(config, logger) + +// With custom HTTP client +httpClient := &http.Client{Timeout: 10 * time.Second} +sdk, err := ollama.NewWithClient(config, httpClient, logger) +``` + +## Usage Examples + +### Simple Text Generation + +```go +ctx := context.Background() + +// Quick generation with default model +response, err := sdk.QuickGenerate(ctx, "Write a short story about a robot") +if err != nil { + log.Printf("Generation failed: %v", err) +} +fmt.Println(response) + +// Generation with specific model +response, err := sdk.QuickGenerateWithModel(ctx, "codellama", "Write a Python function to sort a list") +``` + +### Advanced Text Generation + +```go +req := &ollama.GenerateRequest{ + Model: "llama2", + Prompt: "Explain quantum computing in simple terms", + Options: map[string]interface{}{ + "temperature": 0.7, + "top_p": 0.9, + "max_tokens": 500, + }, +} + +resp, err := sdk.Generate(ctx, req) +if err != nil { + log.Printf("Generation failed: %v", err) +} +fmt.Printf("Generated: %s\n", resp.Response) +``` + +### Chat Conversations + +```go +// Simple chat +response, err := sdk.QuickChat(ctx, "What is the capital of France?") +fmt.Println(response) + +// Chat with system prompt +response, err := sdk.QuickChatWithSystem(ctx, + "You are a helpful coding assistant.", + "How do I implement a binary search in Go?") +``` + +### Advanced Chat with Conversation History + +```go +req := &ollama.ChatRequest{ + Model: "llama2", + Messages: []ollama.Message{ + { + Role: ollama.MessageRoleSystem, + Content: "You are a helpful assistant.", + }, + { + Role: ollama.MessageRoleUser, + Content: "Hello, how are you?", + }, + { + Role: ollama.MessageRoleAssistant, + Content: "I'm doing well, thank you! How can I help you today?", + }, + { + Role: ollama.MessageRoleUser, + Content: "Can you help me with a coding problem?", + }, + }, +} + +resp, err := sdk.Chat(ctx, req) +fmt.Printf("Assistant: %s\n", resp.Message.Content) +``` + +### Fluent Chat Builder + +```go +response, err := sdk.NewChatBuilder(). + SetModel("llama2"). + SetSystem("You are a helpful coding assistant."). + AddMessage(ollama.MessageRoleUser, "How do I create a struct in Go?"). + SetTemperature(0.7). + SetMaxTokens(500). + Send(ctx) + +fmt.Printf("Assistant: %s\n", response.Message.Content) + +// Continue the conversation +response, err = sdk.NewChatBuilder(). + SetModel("llama2"). + SetSystem("You are a helpful coding assistant."). + AddMessage(ollama.MessageRoleUser, "How do I create a struct in Go?"). + SetTemperature(0.7). + SetMaxTokens(500). + Continue(ctx, "Can you show me an example with methods?") + +fmt.Printf("Assistant: %s\n", response.Message.Content) +``` + +### Streaming Responses + +```go +// Streaming text generation +err := sdk.StreamGenerate(ctx, &ollama.GenerateRequest{ + Model: "llama2", + Prompt: "Write a poem about programming", +}, func(chunk string) error { + fmt.Print(chunk) + return nil +}) + +// Streaming chat +err := sdk.StreamChat(ctx, &ollama.ChatRequest{ + Model: "llama2", + Messages: []ollama.Message{ + { + Role: ollama.MessageRoleUser, + Content: "Tell me a story", + }, + }, +}, func(chunk string) error { + fmt.Print(chunk) + return nil +}) +``` + +### Embeddings + +```go +// Generate embeddings +embeddings, err := sdk.QuickEmbed(ctx, "This is some text to embed") +if err != nil { + log.Printf("Embedding failed: %v", err) +} + +// Advanced embedding request +req := &ollama.EmbedRequest{ + Model: "llama2", + Prompt: "Text to generate embeddings for", +} + +resp, err := sdk.Embed(ctx, req) +fmt.Printf("Embedding dimension: %d\n", len(resp.Embedding)) +``` + +### Model Management + +```go +// List available models +models, err := sdk.ListModels(ctx) +for _, model := range models.Models { + fmt.Printf("Model: %s, Size: %d bytes\n", model.Name, model.Size) +} + +// Pull a new model +pullResp, err := sdk.PullModel(ctx, "codellama") +fmt.Printf("Pull status: %s\n", pullResp.Status) + +// Show model information +modelInfo, err := sdk.ShowModel(ctx, "llama2") +fmt.Printf("Model details: %+v\n", modelInfo.Details) + +// Delete a model +deleteResp, err := sdk.DeleteModel(ctx, "unwanted-model") +fmt.Printf("Delete status: %s\n", deleteResp.Status) +``` + +### Health Check + +```go +err := sdk.Health(ctx) +if err != nil { + log.Printf("Ollama service is not available: %v", err) +} else { + fmt.Println("Ollama service is healthy") +} +``` + +## Error Handling + +The SDK provides specific error types for different failure scenarios: + +```go +resp, err := sdk.Generate(ctx, req) +if err != nil { + switch e := err.(type) { + case *ollama.ConfigError: + fmt.Printf("Configuration error: %s\n", e.Message) + case *ollama.ModelError: + fmt.Printf("Model error for %s: %s\n", e.Model, e.Message) + case *ollama.GenerationError: + fmt.Printf("Generation error: %s\n", e.Message) + case *ollama.ConnectionError: + fmt.Printf("Connection error: %s\n", e.Message) + case *ollama.TimeoutError: + fmt.Printf("Timeout during %s: %s\n", e.Operation, e.Message) + default: + fmt.Printf("Unexpected error: %v\n", err) + } +} +``` + +## Advanced Usage + +### Custom HTTP Client + +```go +httpClient := &http.Client{ + Timeout: 60 * time.Second, + Transport: &http.Transport{ + MaxIdleConns: 10, + IdleConnTimeout: 30 * time.Second, + DisableCompression: true, + }, +} + +sdk, err := ollama.NewWithClient(config, httpClient, logger) +``` + +### Custom Options + +```go +req := &ollama.GenerateRequest{ + Model: "llama2", + Prompt: "Write a haiku", +} + +opts := &ollama.GenerateOptions{ + Temperature: &[]float64{0.5}[0], + TopP: &[]float64{0.8}[0], + TopK: &[]int{40}[0], + RepeatPenalty: &[]float64{1.1}[0], + Seed: &[]int{42}[0], + Stop: []string{"", "\n\n"}, + Format: "json", +} + +resp, err := sdk.GenerateWithOptions(ctx, req, opts) +``` + +## Best Practices + +1. **Always use context**: Pass context for cancellation and timeouts +2. **Handle errors properly**: Use type assertions for specific error handling +3. **Configure timeouts**: Set appropriate timeouts for your use case +4. **Use streaming for long responses**: Enable streaming for better user experience +5. **Monitor model usage**: Keep track of token usage and costs +6. **Validate inputs**: Ensure prompts and parameters are within limits +7. **Use structured logging**: Leverage the built-in logging for debugging + +## Troubleshooting + +### Common Issues + +1. **Connection refused**: Ensure Ollama server is running on the correct port +2. **Model not found**: Pull the model first using `PullModel()` +3. **Timeout errors**: Increase the timeout configuration +4. **Rate limiting**: Implement backoff strategies for high-volume usage + +### Debug Mode + +Enable debug logging to troubleshoot issues: + +```go +config := ollama.DefaultConfig() +config.EnableLogging = true + +loggerConfig := &logger.Config{ + Level: "debug", + Format: "json", + Output: "stdout", +} +projectLogger := logger.NewLogger(loggerConfig) + +sdk, err := ollama.New(config, projectLogger) +``` + +## License + +This SDK is part of the Tender Management system and follows the same licensing terms. diff --git a/pkg/ollama/USAGE.md b/pkg/ollama/USAGE.md new file mode 100644 index 0000000..ef1dc96 --- /dev/null +++ b/pkg/ollama/USAGE.md @@ -0,0 +1,263 @@ +# Ollama SDK - Quick Usage Guide + +## Installation +The SDK is already part of your `tm` module. Import it as: +```go +import "tm/pkg/ollama" +``` + +## Basic Setup +```go +// 1. Configuration +config := ollama.DefaultConfig() +config.BaseURL = "http://localhost:11434" // Your Ollama server URL +config.DefaultModel = "llama2" // Default model to use + +// 2. Logger (using project's logger) +loggerConfig := &logger.Config{Level: "info", Format: "json", Output: "stdout"} +projectLogger := logger.NewLogger(loggerConfig) + +// 3. Create SDK +sdk, err := ollama.New(config, projectLogger) +if err != nil { + log.Fatal("Failed to create Ollama SDK:", err) +} +``` + +## Environment Variables +```bash +export OLLAMA_BASE_URL="http://localhost:11434" +export OLLAMA_TIMEOUT="300s" +export OLLAMA_RETRY_ATTEMPTS="3" +export OLLAMA_RETRY_DELAY="2s" +export OLLAMA_ENABLE_LOGGING="true" +export OLLAMA_DEFAULT_MODEL="llama2" +export OLLAMA_DEFAULT_TEMPERATURE="0.7" +export OLLAMA_DEFAULT_TOP_P="0.9" +export OLLAMA_DEFAULT_MAX_TOKENS="1000" +``` + +## Quick Examples + +### Simple Text Generation +```go +ctx := context.Background() + +// Generate text with default model +response, err := sdk.QuickGenerate(ctx, "Write a short story about a robot") +if err != nil { + log.Printf("Generation failed: %v", err) +} +fmt.Println(response) + +// Generate with specific model +response, err := sdk.QuickGenerateWithModel(ctx, "codellama", "Write a Python function to sort a list") +``` + +### Chat Conversations +```go +// Simple chat +response, err := sdk.QuickChat(ctx, "What is the capital of France?") +fmt.Println(response) + +// Chat with system prompt +response, err := sdk.QuickChatWithSystem(ctx, + "You are a helpful coding assistant.", + "How do I implement a binary search in Go?") +``` + +### Advanced Chat with Builder +```go +// Create a chat builder +builder := sdk.NewChatBuilder(). + SetModel("llama2"). + SetSystem("You are a helpful coding assistant."). + SetTemperature(0.7). + SetMaxTokens(500) + +// Start conversation +response, err := builder. + AddMessage(ollama.MessageRoleUser, "How do I create a struct in Go?"). + Send(ctx) + +fmt.Printf("Assistant: %s\n", response.Message.Content) + +// Continue the conversation +response, err = builder.Continue(ctx, "Can you show me an example with methods?") +fmt.Printf("Assistant: %s\n", response.Message.Content) +``` + +### Streaming Responses +```go +// Streaming text generation +err := sdk.StreamGenerate(ctx, &ollama.GenerateRequest{ + Model: "llama2", + Prompt: "Write a poem about programming", +}, func(chunk string) error { + fmt.Print(chunk) // Print each chunk as it arrives + return nil +}) + +// Streaming chat +err := sdk.StreamChat(ctx, &ollama.ChatRequest{ + Model: "llama2", + Messages: []ollama.Message{ + { + Role: ollama.MessageRoleUser, + Content: "Tell me a story", + }, + }, +}, func(chunk string) error { + fmt.Print(chunk) + return nil +}) +``` + +### Embeddings +```go +// Generate embeddings +embeddings, err := sdk.QuickEmbed(ctx, "This is some text to embed") +if err != nil { + log.Printf("Embedding failed: %v", err) +} + +fmt.Printf("Embedding dimension: %d\n", len(embeddings)) +``` + +### Model Management +```go +// List available models +models, err := sdk.ListModels(ctx) +for _, model := range models.Models { + fmt.Printf("Model: %s, Size: %d bytes\n", model.Name, model.Size) +} + +// Pull a new model +pullResp, err := sdk.PullModel(ctx, "codellama") +fmt.Printf("Pull status: %s\n", pullResp.Status) + +// Show model information +modelInfo, err := sdk.ShowModel(ctx, "llama2") +fmt.Printf("Model family: %s\n", modelInfo.Details.Family) + +// Delete a model +deleteResp, err := sdk.DeleteModel(ctx, "unwanted-model") +fmt.Printf("Delete status: %s\n", deleteResp.Status) +``` + +### Health Check +```go +err := sdk.Health(ctx) +if err != nil { + log.Printf("Ollama service is not available: %v", err) +} else { + fmt.Println("Ollama service is healthy") +} +``` + +## Error Handling +```go +resp, err := sdk.Generate(ctx, req) +if err != nil { + switch e := err.(type) { + case *ollama.ConfigError: + fmt.Printf("Configuration error: %s\n", e.Message) + case *ollama.ModelError: + fmt.Printf("Model error for %s: %s\n", e.Model, e.Message) + case *ollama.GenerationError: + fmt.Printf("Generation error: %s\n", e.Message) + case *ollama.ConnectionError: + fmt.Printf("Connection error: %s\n", e.Message) + case *ollama.TimeoutError: + fmt.Printf("Timeout during %s: %s\n", e.Operation, e.Message) + default: + fmt.Printf("Unexpected error: %v\n", err) + } +} +``` + +## Advanced Configuration +```go +config := &ollama.Config{ + BaseURL: "http://localhost:11434", + Timeout: 300 * time.Second, + RetryAttempts: 3, + RetryDelay: 2 * time.Second, + EnableLogging: true, + UserAgent: "MyApp-OllamaSDK/1.0", + DefaultModel: "llama2", + DefaultTemperature: 0.7, + DefaultTopP: 0.9, + DefaultMaxTokens: 1000, + EnableStreaming: false, +} + +sdk, err := ollama.New(config, logger) +``` + +## Custom HTTP Client +```go +httpClient := &http.Client{ + Timeout: 60 * time.Second, + Transport: &http.Transport{ + MaxIdleConns: 10, + IdleConnTimeout: 30 * time.Second, + DisableCompression: true, + }, +} + +sdk, err := ollama.NewWithClient(config, httpClient, logger) +``` + +## Integration with Project Logger +```go +// Use the project's logger system +loggerConfig := &logger.Config{ + Level: "info", + Format: "json", + Output: "stdout", +} +projectLogger := logger.NewLogger(loggerConfig) + +sdk, err := ollama.New(config, projectLogger) +``` + +## Best Practices + +1. **Always use context**: Pass context for cancellation and timeouts +2. **Handle errors properly**: Use type assertions for specific error handling +3. **Configure timeouts**: Set appropriate timeouts for your use case +4. **Use streaming for long responses**: Enable streaming for better user experience +5. **Monitor model usage**: Keep track of token usage and costs +6. **Validate inputs**: Ensure prompts and parameters are within limits +7. **Use structured logging**: Leverage the built-in logging for debugging + +## Common Use Cases + +### Document Summarization +```go +response, err := sdk.QuickChatWithSystem(ctx, + "You are a helpful assistant that summarizes documents concisely.", + "Please summarize this document: [document content]") +``` + +### Code Review Assistant +```go +response, err := sdk.QuickChatWithSystem(ctx, + "You are a senior software engineer reviewing code. Provide constructive feedback.", + "Please review this Go code: [code snippet]") +``` + +### Translation Service +```go +response, err := sdk.QuickChatWithSystem(ctx, + "You are a professional translator. Translate the following text accurately.", + "Translate this to French: Hello, how are you?") +``` + +### Question Answering +```go +response, err := sdk.QuickChatWithSystem(ctx, + "You are a helpful assistant that answers questions accurately and concisely.", + "What is the difference between HTTP and HTTPS?") +``` diff --git a/pkg/ollama/client.go b/pkg/ollama/client.go new file mode 100644 index 0000000..a070ed6 --- /dev/null +++ b/pkg/ollama/client.go @@ -0,0 +1,290 @@ +package ollama + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "time" + "tm/pkg/logger" +) + +// Client implements the HTTPClient interface for Ollama API +type Client struct { + config *Config + httpClient *http.Client + logger logger.Logger +} + +// NewClient creates a new Ollama HTTP client +func NewClient(config *Config, logger logger.Logger) (*Client, error) { + if config == nil { + config = DefaultConfig() + } + + if err := config.Validate(); err != nil { + return nil, err + } + + httpClient := &http.Client{ + Timeout: config.Timeout, + } + + return &Client{ + config: config, + httpClient: httpClient, + logger: logger, + }, nil +} + +// NewClientWithHTTPClient creates a new client with a custom HTTP client +func NewClientWithHTTPClient(config *Config, httpClient *http.Client, logger logger.Logger) (*Client, error) { + if config == nil { + config = DefaultConfig() + } + + if err := config.Validate(); err != nil { + return nil, err + } + + return &Client{ + config: config, + httpClient: httpClient, + logger: logger, + }, nil +} + +// Post performs a POST request to the Ollama API +func (c *Client) Post(ctx context.Context, url string, body interface{}, result interface{}) error { + return c.makeRequest(ctx, "POST", url, body, result) +} + +// Get performs a GET request to the Ollama API +func (c *Client) Get(ctx context.Context, url string, result interface{}) error { + return c.makeRequest(ctx, "GET", url, nil, result) +} + +// Delete performs a DELETE request to the Ollama API +func (c *Client) Delete(ctx context.Context, url string, body interface{}, result interface{}) error { + return c.makeRequest(ctx, "DELETE", url, body, result) +} + +// makeRequest performs an HTTP request with retry logic and logging +func (c *Client) makeRequest(ctx context.Context, method, endpoint string, body interface{}, result interface{}) error { + url := c.config.BaseURL + endpoint + + var bodyReader io.Reader + if body != nil { + bodyBytes, err := json.Marshal(body) + if err != nil { + return fmt.Errorf("failed to marshal request body: %w", err) + } + bodyReader = bytes.NewReader(bodyBytes) + } + + var lastErr error + for attempt := 0; attempt <= c.config.RetryAttempts; attempt++ { + if attempt > 0 { + // Wait before retry + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(c.config.RetryDelay): + } + } + + err := c.doRequest(ctx, method, url, bodyReader, result) + if err == nil { + return nil + } + + lastErr = err + + // Don't retry on certain errors + if c.shouldNotRetry(err) { + break + } + + if c.logger != nil { + c.logger.Warn("Request failed, retrying", map[string]interface{}{ + "attempt": attempt + 1, + "max_attempts": c.config.RetryAttempts + 1, + "error": err.Error(), + "url": url, + "method": method, + }) + } + } + + return fmt.Errorf("request failed after %d attempts: %w", c.config.RetryAttempts+1, lastErr) +} + +// doRequest performs a single HTTP request +func (c *Client) doRequest(ctx context.Context, method, url string, body io.Reader, result interface{}) error { + req, err := http.NewRequestWithContext(ctx, method, url, body) + if err != nil { + return fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + req.Header.Set("User-Agent", c.config.UserAgent) + + if c.logger != nil && c.config.EnableLogging { + c.logger.Debug("Making HTTP request", map[string]interface{}{ + "method": method, + "url": url, + "user_agent": c.config.UserAgent, + }) + } + + resp, err := c.httpClient.Do(req) + if err != nil { + return fmt.Errorf("failed to execute request: %w", err) + } + defer resp.Body.Close() + + if c.logger != nil && c.config.EnableLogging { + c.logger.Debug("Received HTTP response", map[string]interface{}{ + "status_code": resp.StatusCode, + "url": url, + }) + } + + if resp.StatusCode >= 400 { + bodyBytes, _ := io.ReadAll(resp.Body) + return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(bodyBytes)) + } + + if result != nil { + if err := json.NewDecoder(resp.Body).Decode(result); err != nil { + return fmt.Errorf("failed to decode response: %w", err) + } + } + + return nil +} + +// shouldNotRetry determines if an error should not be retried +func (c *Client) shouldNotRetry(err error) bool { + // Don't retry on context cancellation + if err == context.Canceled || err == context.DeadlineExceeded { + return true + } + + // Don't retry on certain HTTP status codes + if err.Error() == "HTTP 400" || err.Error() == "HTTP 401" || err.Error() == "HTTP 403" || err.Error() == "HTTP 404" { + return true + } + + return false +} + +// StreamGenerate performs a streaming text generation request +func (c *Client) StreamGenerate(ctx context.Context, req *GenerateRequest, onChunk func(*StreamChunk) error) error { + url := c.config.BaseURL + "/api/generate" + + // Force streaming + stream := true + req.Stream = &stream + + bodyBytes, err := json.Marshal(req) + if err != nil { + return fmt.Errorf("failed to marshal request body: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(bodyBytes)) + if err != nil { + return fmt.Errorf("failed to create request: %w", err) + } + + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("User-Agent", c.config.UserAgent) + + resp, err := c.httpClient.Do(httpReq) + if err != nil { + return fmt.Errorf("failed to execute request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode >= 400 { + bodyBytes, _ := io.ReadAll(resp.Body) + return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(bodyBytes)) + } + + decoder := json.NewDecoder(resp.Body) + for { + var chunk StreamChunk + if err := decoder.Decode(&chunk); err != nil { + if err == io.EOF { + break + } + return fmt.Errorf("failed to decode stream chunk: %w", err) + } + + if err := onChunk(&chunk); err != nil { + return fmt.Errorf("error processing stream chunk: %w", err) + } + + if chunk.Done { + break + } + } + + return nil +} + +// StreamChat performs a streaming chat request +func (c *Client) StreamChat(ctx context.Context, req *ChatRequest, onChunk func(*StreamChunk) error) error { + url := c.config.BaseURL + "/api/chat" + + // Force streaming + stream := true + req.Stream = &stream + + bodyBytes, err := json.Marshal(req) + if err != nil { + return fmt.Errorf("failed to marshal request body: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(bodyBytes)) + if err != nil { + return fmt.Errorf("failed to create request: %w", err) + } + + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("User-Agent", c.config.UserAgent) + + resp, err := c.httpClient.Do(httpReq) + if err != nil { + return fmt.Errorf("failed to execute request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode >= 400 { + bodyBytes, _ := io.ReadAll(resp.Body) + return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(bodyBytes)) + } + + decoder := json.NewDecoder(resp.Body) + for { + var chunk StreamChunk + if err := decoder.Decode(&chunk); err != nil { + if err == io.EOF { + break + } + return fmt.Errorf("failed to decode stream chunk: %w", err) + } + + if err := onChunk(&chunk); err != nil { + return fmt.Errorf("error processing stream chunk: %w", err) + } + + if chunk.Done { + break + } + } + + return nil +} diff --git a/pkg/ollama/config.go b/pkg/ollama/config.go new file mode 100644 index 0000000..2bb1d80 --- /dev/null +++ b/pkg/ollama/config.go @@ -0,0 +1,91 @@ +package ollama + +import ( + "time" +) + +// Config represents the configuration for the Ollama SDK +type Config struct { + // BaseURL is the base URL of the Ollama API server + BaseURL string `env:"OLLAMA_BASE_URL" default:"http://localhost:11434"` + + // Timeout is the timeout for HTTP requests + Timeout time.Duration `env:"OLLAMA_TIMEOUT" default:"300s"` + + // RetryAttempts is the number of retry attempts for failed requests + RetryAttempts int `env:"OLLAMA_RETRY_ATTEMPTS" default:"3"` + + // RetryDelay is the delay between retry attempts + RetryDelay time.Duration `env:"OLLAMA_RETRY_DELAY" default:"2s"` + + // EnableLogging enables request/response logging + EnableLogging bool `env:"OLLAMA_ENABLE_LOGGING" default:"true"` + + // UserAgent is the user agent string for HTTP requests + UserAgent string `env:"OLLAMA_USER_AGENT" default:"TenderManagement-OllamaSDK/1.0"` + + // DefaultModel is the default model to use if none is specified + DefaultModel string `env:"OLLAMA_DEFAULT_MODEL" default:"llama2"` + + // DefaultTemperature is the default temperature for text generation + DefaultTemperature float64 `env:"OLLAMA_DEFAULT_TEMPERATURE" default:"0.7"` + + // DefaultTopP is the default top_p value for text generation + DefaultTopP float64 `env:"OLLAMA_DEFAULT_TOP_P" default:"0.9"` + + // DefaultMaxTokens is the default maximum number of tokens to generate + DefaultMaxTokens int `env:"OLLAMA_DEFAULT_MAX_TOKENS" default:"1000"` + + // EnableStreaming enables streaming responses by default + EnableStreaming bool `env:"OLLAMA_ENABLE_STREAMING" default:"false"` +} + +// DefaultConfig returns a default configuration for the Ollama SDK +func DefaultConfig() *Config { + return &Config{ + BaseURL: "http://localhost:11434", + Timeout: 300 * time.Second, + RetryAttempts: 3, + RetryDelay: 2 * time.Second, + EnableLogging: true, + UserAgent: "TenderManagement-OllamaSDK/1.0", + DefaultModel: "llama2", + DefaultTemperature: 0.7, + DefaultTopP: 0.9, + DefaultMaxTokens: 1000, + EnableStreaming: false, + } +} + +// Validate validates the configuration +func (c *Config) Validate() error { + if c.BaseURL == "" { + return ErrInvalidConfig("BaseURL cannot be empty") + } + + if c.Timeout <= 0 { + return ErrInvalidConfig("Timeout must be positive") + } + + if c.RetryAttempts < 0 { + return ErrInvalidConfig("RetryAttempts cannot be negative") + } + + if c.RetryDelay < 0 { + return ErrInvalidConfig("RetryDelay cannot be negative") + } + + if c.DefaultTemperature < 0 || c.DefaultTemperature > 2 { + return ErrInvalidConfig("DefaultTemperature must be between 0 and 2") + } + + if c.DefaultTopP < 0 || c.DefaultTopP > 1 { + return ErrInvalidConfig("DefaultTopP must be between 0 and 1") + } + + if c.DefaultMaxTokens < 0 { + return ErrInvalidConfig("DefaultMaxTokens cannot be negative") + } + + return nil +} diff --git a/pkg/ollama/entities.go b/pkg/ollama/entities.go new file mode 100644 index 0000000..fb8615e --- /dev/null +++ b/pkg/ollama/entities.go @@ -0,0 +1,182 @@ +package ollama + +import ( + "time" +) + +// GenerateRequest represents a request to generate text +type GenerateRequest struct { + Model string `json:"model" validate:"required"` + Prompt string `json:"prompt" validate:"required"` + Stream *bool `json:"stream,omitempty"` + Format string `json:"format,omitempty"` + Options map[string]interface{} `json:"options,omitempty"` + System string `json:"system,omitempty"` + Template string `json:"template,omitempty"` + Context []int `json:"context,omitempty"` + Raw bool `json:"raw,omitempty"` + KeepAlive string `json:"keep_alive,omitempty"` +} + +// GenerateResponse represents the response from a text generation request +type GenerateResponse struct { + Model string `json:"model"` + CreatedAt time.Time `json:"created_at"` + Response string `json:"response"` + Done bool `json:"done"` + Context []int `json:"context,omitempty"` + TotalDuration int64 `json:"total_duration,omitempty"` + LoadDuration int64 `json:"load_duration,omitempty"` + PromptEvalCount int `json:"prompt_eval_count,omitempty"` + PromptEvalDuration int64 `json:"prompt_eval_duration,omitempty"` + EvalCount int `json:"eval_count,omitempty"` + EvalDuration int64 `json:"eval_duration,omitempty"` +} + +// ChatRequest represents a request for a chat conversation +type ChatRequest struct { + Model string `json:"model" validate:"required"` + Messages []Message `json:"messages" validate:"required,min=1"` + Stream *bool `json:"stream,omitempty"` + Format string `json:"format,omitempty"` + Options ChatOptions `json:"options,omitempty"` + KeepAlive string `json:"keep_alive,omitempty"` +} + +// Message represents a message in a chat conversation +type Message struct { + Role MessageRole `json:"role" validate:"required,oneof=system user assistant"` + Content string `json:"content" validate:"required"` + Images []string `json:"images,omitempty"` +} + +// ChatOptions represents options for chat requests +type ChatOptions struct { + Temperature *float64 `json:"temperature,omitempty"` + TopP *float64 `json:"top_p,omitempty"` + TopK *int `json:"top_k,omitempty"` + RepeatPenalty *float64 `json:"repeat_penalty,omitempty"` + Seed *int `json:"seed,omitempty"` + Stop []string `json:"stop,omitempty"` + NumCtx *int `json:"num_ctx,omitempty"` + NumPredict *int `json:"num_predict,omitempty"` + NumKeep *int `json:"num_keep,omitempty"` + NumThread *int `json:"num_thread,omitempty"` + RepeatLastN *int `json:"repeat_last_n,omitempty"` + Options map[string]interface{} `json:"options,omitempty"` +} + +// ChatResponse represents the response from a chat request +type ChatResponse struct { + Model string `json:"model"` + CreatedAt time.Time `json:"created_at"` + Message Message `json:"message"` + Done bool `json:"done"` + TotalDuration int64 `json:"total_duration,omitempty"` + LoadDuration int64 `json:"load_duration,omitempty"` + PromptEvalCount int `json:"prompt_eval_count,omitempty"` + PromptEvalDuration int64 `json:"prompt_eval_duration,omitempty"` + EvalCount int `json:"eval_count,omitempty"` + EvalDuration int64 `json:"eval_duration,omitempty"` +} + +// EmbedRequest represents a request to generate embeddings +type EmbedRequest struct { + Model string `json:"model" validate:"required"` + Prompt string `json:"prompt" validate:"required"` + Options map[string]interface{} `json:"options,omitempty"` +} + +// EmbedResponse represents the response from an embedding request +type EmbedResponse struct { + Embedding []float64 `json:"embedding"` +} + +// PullRequest represents a request to pull a model +type PullRequest struct { + Name string `json:"name" validate:"required"` + Insecure bool `json:"insecure,omitempty"` + Stream *bool `json:"stream,omitempty"` +} + +// PullResponse represents the response from a pull request +type PullResponse struct { + Status string `json:"status"` + Digest string `json:"digest,omitempty"` + Total int64 `json:"total,omitempty"` + Completed int64 `json:"completed,omitempty"` +} + +// ListModelsResponse represents the response from listing models +type ListModelsResponse struct { + Models []ModelInfo `json:"models"` +} + +// ModelInfo represents information about a model +type ModelInfo struct { + Name string `json:"name"` + ModifiedAt time.Time `json:"modified_at"` + Size int64 `json:"size"` + Digest string `json:"digest"` + Details ModelDetails `json:"details,omitempty"` +} + +// ModelDetails represents detailed information about a model +type ModelDetails struct { + Format string `json:"format,omitempty"` + Family string `json:"family,omitempty"` + Families []string `json:"families,omitempty"` + ParameterSize string `json:"parameter_size,omitempty"` + QuantizationLevel string `json:"quantization_level,omitempty"` + ParentModel string `json:"parent_model,omitempty"` + Templates map[string]interface{} `json:"templates,omitempty"` + System string `json:"system,omitempty"` + Options map[string]interface{} `json:"options,omitempty"` +} + +// DeleteRequest represents a request to delete a model +type DeleteRequest struct { + Name string `json:"name" validate:"required"` +} + +// DeleteResponse represents the response from a delete request +type DeleteResponse struct { + Status string `json:"status"` +} + +// ShowModelRequest represents a request to show model information +type ShowModelRequest struct { + Name string `json:"name" validate:"required"` + Verbose bool `json:"verbose,omitempty"` +} + +// ShowModelResponse represents the response from a show model request +type ShowModelResponse struct { + License string `json:"license,omitempty"` + Modelfile string `json:"modelfile,omitempty"` + Parameters string `json:"parameters,omitempty"` + Template string `json:"template,omitempty"` + Details ModelDetails `json:"details,omitempty"` + Options map[string]interface{} `json:"options,omitempty"` +} + +// HealthResponse represents the health check response +type HealthResponse struct { + Status string `json:"status"` +} + +// StreamChunk represents a chunk in a streaming response +type StreamChunk struct { + Model string `json:"model,omitempty"` + CreatedAt time.Time `json:"created_at,omitempty"` + Response string `json:"response,omitempty"` + Done bool `json:"done"` + Context []int `json:"context,omitempty"` + TotalDuration int64 `json:"total_duration,omitempty"` + LoadDuration int64 `json:"load_duration,omitempty"` + PromptEvalCount int `json:"prompt_eval_count,omitempty"` + PromptEvalDuration int64 `json:"prompt_eval_duration,omitempty"` + EvalCount int `json:"eval_count,omitempty"` + EvalDuration int64 `json:"eval_duration,omitempty"` + Message *Message `json:"message,omitempty"` +} diff --git a/pkg/ollama/errors.go b/pkg/ollama/errors.go new file mode 100644 index 0000000..e121c07 --- /dev/null +++ b/pkg/ollama/errors.go @@ -0,0 +1,146 @@ +package ollama + +import ( + "fmt" +) + +// Error types for the Ollama SDK +var ( + // ErrInvalidConfig is returned when configuration is invalid + ErrInvalidConfig = func(msg string) error { + return &ConfigError{Message: msg} + } + + // ErrInvalidRequest is returned when a request is invalid + ErrInvalidRequest = func(msg string) error { + return &RequestError{Message: msg} + } + + // ErrModelNotFound is returned when a model is not found + ErrModelNotFound = func(model string) error { + return &ModelError{Message: fmt.Sprintf("model %s not found", model), Model: model} + } + + // ErrModelPullFailed is returned when pulling a model fails + ErrModelPullFailed = func(model string, reason string) error { + return &ModelError{Message: fmt.Sprintf("failed to pull model %s: %s", model, reason), Model: model} + } + + // ErrModelDeleteFailed is returned when deleting a model fails + ErrModelDeleteFailed = func(model string, reason string) error { + return &ModelError{Message: fmt.Sprintf("failed to delete model %s: %s", model, reason), Model: model} + } + + // ErrGenerationFailed is returned when text generation fails + ErrGenerationFailed = func(reason string) error { + return &GenerationError{Message: fmt.Sprintf("generation failed: %s", reason)} + } + + // ErrChatFailed is returned when chat fails + ErrChatFailed = func(reason string) error { + return &ChatError{Message: fmt.Sprintf("chat failed: %s", reason)} + } + + // ErrEmbeddingFailed is returned when embedding generation fails + ErrEmbeddingFailed = func(reason string) error { + return &EmbeddingError{Message: fmt.Sprintf("embedding failed: %s", reason)} + } + + // ErrConnectionFailed is returned when connection to Ollama fails + ErrConnectionFailed = func(reason string) error { + return &ConnectionError{Message: fmt.Sprintf("connection failed: %s", reason)} + } + + // ErrTimeout is returned when a request times out + ErrTimeout = func(operation string) error { + return &TimeoutError{Message: fmt.Sprintf("timeout during %s", operation), Operation: operation} + } + + // ErrRateLimitExceeded is returned when rate limit is exceeded + ErrRateLimitExceeded = func() error { + return &RateLimitError{Message: "rate limit exceeded"} + } +) + +// ConfigError represents a configuration error +type ConfigError struct { + Message string +} + +func (e *ConfigError) Error() string { + return fmt.Sprintf("config error: %s", e.Message) +} + +// RequestError represents a request validation error +type RequestError struct { + Message string +} + +func (e *RequestError) Error() string { + return fmt.Sprintf("request error: %s", e.Message) +} + +// ModelError represents a model-related error +type ModelError struct { + Message string + Model string +} + +func (e *ModelError) Error() string { + return fmt.Sprintf("model error: %s", e.Message) +} + +// GenerationError represents a text generation error +type GenerationError struct { + Message string +} + +func (e *GenerationError) Error() string { + return fmt.Sprintf("generation error: %s", e.Message) +} + +// ChatError represents a chat error +type ChatError struct { + Message string +} + +func (e *ChatError) Error() string { + return fmt.Sprintf("chat error: %s", e.Message) +} + +// EmbeddingError represents an embedding generation error +type EmbeddingError struct { + Message string +} + +func (e *EmbeddingError) Error() string { + return fmt.Sprintf("embedding error: %s", e.Message) +} + +// ConnectionError represents a connection error +type ConnectionError struct { + Message string +} + +func (e *ConnectionError) Error() string { + return fmt.Sprintf("connection error: %s", e.Message) +} + +// TimeoutError represents a timeout error +type TimeoutError struct { + Message string + Operation string +} + +func (e *TimeoutError) Error() string { + return fmt.Sprintf("timeout error: %s", e.Message) +} + +// RateLimitError represents a rate limit error +type RateLimitError struct { + Message string +} + +func (e *RateLimitError) Error() string { + return fmt.Sprintf("rate limit error: %s", e.Message) +} diff --git a/pkg/ollama/interfaces.go b/pkg/ollama/interfaces.go new file mode 100644 index 0000000..b9695e1 --- /dev/null +++ b/pkg/ollama/interfaces.go @@ -0,0 +1,121 @@ +package ollama + +import ( + "context" +) + +// OllamaService defines the interface for Ollama operations +type OllamaService interface { + // Generate generates text using a specified model and prompt + Generate(ctx context.Context, req *GenerateRequest) (*GenerateResponse, error) + + // GenerateWithOptions generates text with additional options + GenerateWithOptions(ctx context.Context, req *GenerateRequest, opts *GenerateOptions) (*GenerateResponse, error) + + // Chat performs a conversational chat with the model + Chat(ctx context.Context, req *ChatRequest) (*ChatResponse, error) + + // Embed generates embeddings for the given text + Embed(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) + + // PullModel pulls a model from the registry + PullModel(ctx context.Context, model string) (*PullResponse, error) + + // ListModels lists available models + ListModels(ctx context.Context) (*ListModelsResponse, error) + + // DeleteModel deletes a model + DeleteModel(ctx context.Context, model string) (*DeleteResponse, error) + + // ShowModel shows model information + ShowModel(ctx context.Context, model string) (*ShowModelResponse, error) + + // Health checks if the Ollama service is available + Health(ctx context.Context) error + + // NewChatBuilder creates a new chat builder for conversational interactions + NewChatBuilder() ChatBuilder +} + +// HTTPClient defines the interface for HTTP operations +type HTTPClient interface { + // Post performs a POST request + Post(ctx context.Context, url string, body interface{}, result interface{}) error + + // Get performs a GET request + Get(ctx context.Context, url string, result interface{}) error + + // Delete performs a DELETE request + Delete(ctx context.Context, url string, result interface{}) error +} + +// Logger defines the interface for logging operations +type Logger interface { + Debug(msg string, fields map[string]interface{}) + Info(msg string, fields map[string]interface{}) + Warn(msg string, fields map[string]interface{}) + Error(msg string, fields map[string]interface{}) + Fatal(msg string, fields map[string]interface{}) + WithFields(fields map[string]interface{}) Logger + Sync() error +} + +// ChatBuilder provides a fluent interface for building chat requests +type ChatBuilder interface { + // SetModel sets the model to use + SetModel(model string) ChatBuilder + + // SetSystem sets the system prompt + SetSystem(system string) ChatBuilder + + // AddMessage adds a message to the conversation + AddMessage(role MessageRole, content string) ChatBuilder + + // SetTemperature sets the temperature for generation + SetTemperature(temperature float64) ChatBuilder + + // SetTopP sets the top_p parameter + SetTopP(topP float64) ChatBuilder + + // SetMaxTokens sets the maximum number of tokens to generate + SetMaxTokens(maxTokens int) ChatBuilder + + // SetStream sets whether to stream the response + SetStream(stream bool) ChatBuilder + + // SetFormat sets the response format + SetFormat(format string) ChatBuilder + + // SetOptions sets custom options + SetOptions(options map[string]interface{}) ChatBuilder + + // Build creates the chat request + Build() (*ChatRequest, error) + + // Send builds and sends the chat request + Send(ctx context.Context) (*ChatResponse, error) + + // Continue continues the conversation with the current context + Continue(ctx context.Context, userMessage string) (*ChatResponse, error) +} + +// MessageRole represents the role of a message in a chat +type MessageRole string + +const ( + MessageRoleSystem MessageRole = "system" + MessageRoleUser MessageRole = "user" + MessageRoleAssistant MessageRole = "assistant" +) + +// GenerateOptions represents additional options for text generation +type GenerateOptions struct { + Temperature *float64 `json:"temperature,omitempty"` + TopP *float64 `json:"top_p,omitempty"` + TopK *int `json:"top_k,omitempty"` + RepeatPenalty *float64 `json:"repeat_penalty,omitempty"` + Seed *int `json:"seed,omitempty"` + Stop []string `json:"stop,omitempty"` + Format string `json:"format,omitempty"` + Options map[string]interface{} `json:"options,omitempty"` +} diff --git a/pkg/ollama/sdk.go b/pkg/ollama/sdk.go new file mode 100644 index 0000000..eab9b0c --- /dev/null +++ b/pkg/ollama/sdk.go @@ -0,0 +1,232 @@ +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 +}