mazyar/tm_back
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'feature/ai-translator' into develop

Browse Source
This commit is contained in:
n.nakhostin
2025-10-18 09:04:32 +03:30
parent 952986331a 8e5308b70b
commit 69d201479f
27 changed files with 3354 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/config/config.go
+15
View File
@@ -74,6 +74,21 @@ type NotificationConfig struct {
UserAgent string `env:"NOTIFICATION_USER_AGENT"`
}
// OllamaConfig represents the configuration for the Ollama SDK
type OllamaConfig struct {
BaseURL string `env:"OLLAMA_BASE_URL" default:"http://localhost:11434"`
Timeout time.Duration `env:"OLLAMA_TIMEOUT" default:"300s"`
RetryAttempts int `env:"OLLAMA_RETRY_ATTEMPTS" default:"3"`
RetryDelay time.Duration `env:"OLLAMA_RETRY_DELAY" default:"2s"`
EnableLogging bool `env:"OLLAMA_ENABLE_LOGGING" default:"true"`
UserAgent string `env:"OLLAMA_USER_AGENT" default:"TenderManagement-OllamaSDK/1.0"`
DefaultModel string `env:"OLLAMA_DEFAULT_MODEL" default:"llama2"`
DefaultTemperature float64 `env:"OLLAMA_DEFAULT_TEMPERATURE" default:"0.7"`
DefaultTopP float64 `env:"OLLAMA_DEFAULT_TOP_P" default:"0.9"`
DefaultMaxTokens int `env:"OLLAMA_DEFAULT_MAX_TOKENS" default:"1000"`
EnableStreaming bool `env:"OLLAMA_ENABLE_STREAMING" default:"false"`
}
// LoadConfig loads configuration with priority: OS environment variables > .env file
// OS environment variables take precedence over .env file values
func LoadConfig[T any](path string, config T) (T, error) {
pkg/mongo/repository.go
+29
View File
@@ -47,6 +47,9 @@ type Repository[T any] interface {
// Create inserts a new document
Create(ctx context.Context, model *T) error
// CreateMany inserts multiple documents
CreateMany(ctx context.Context, models []T) error
// Update updates an existing document
Update(ctx context.Context, model *T) error
@@ -239,6 +242,32 @@ func (r *repository[T]) Create(ctx context.Context, model *T) error {
return nil
}
// CreateMany inserts multiple documents
func (r *repository[T]) CreateMany(ctx context.Context, models []T) error {
// Set timestamps if the model implements Timestampable
for _, v := range models {
if timestampable, ok := any(v).(Timestampable); ok {
now := time.Now().Unix()
timestampable.SetCreatedAt(now)
timestampable.SetUpdatedAt(now)
}
}
result, err := r.collection.InsertMany(ctx, models)
if err != nil {
r.logger.Error("Failed to create multiple documents", map[string]interface{}{
"error": err.Error(),
})
return fmt.Errorf("failed to create multiple documents: %w", err)
}
r.logger.Info("Documents created successfully", map[string]interface{}{
"ids": result.InsertedIDs,
})
return nil
}
// Update updates an existing document
func (r *repository[T]) Update(ctx context.Context, model *T) error {
// Set updated timestamp if the model implements Timestampable
pkg/ollama/README.md
+394
View File
@@ -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{"</s>", "\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.
pkg/ollama/USAGE.md
+263
View File
@@ -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?")
```
pkg/ollama/client.go
+290
View File
@@ -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
}
pkg/ollama/config.go
+91
View File
@@ -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
}
pkg/ollama/entities.go
+182
View File
@@ -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"`
}
pkg/ollama/errors.go
+146
View File
@@ -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)
}
pkg/ollama/interfaces.go
+121
View File
@@ -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"`
}
pkg/ollama/sdk.go
+232
View File
@@ -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
}
pkg/ollama/service.go
+444
View File
@@ -0,0 +1,444 @@
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
}
pkg/schedule/cron.go
+4 -4
View File
@@ -2,7 +2,7 @@ package schedule
import (
"time"
"tm/internal/tender"
"tm/internal/notice"
"tm/pkg/logger"
"tm/pkg/mongo"
@@ -13,7 +13,7 @@ import (
type CronScheduler struct {
cron *cron.Cron
mongoManager *mongo.ConnectionManager
tenderRepo tender.TenderRepository
noticeRepo notice.Repository
logger logger.Logger
}
@@ -31,14 +31,14 @@ type Job struct {
}
// NewCronScheduler creates a new cron scheduler instance
func NewCronScheduler(logger logger.Logger, mongoManager *mongo.ConnectionManager, tenderRepo tender.TenderRepository) *CronScheduler {
func NewCronScheduler(logger logger.Logger, mongoManager *mongo.ConnectionManager, noticeRepo notice.Repository) *CronScheduler {
// Create cron with timezone support and seconds precision
c := cron.New(cron.WithLocation(time.UTC), cron.WithSeconds())
return &CronScheduler{
cron: c,
mongoManager: mongoManager,
tenderRepo: tenderRepo,
noticeRepo: noticeRepo,
logger: logger,
}
}