Add Ollama SDK with Core Functionality and Documentation
- Introduced the Ollama SDK, providing a comprehensive interface for interacting with Ollama AI models, including text generation, chat conversations, embeddings, and model management. - Implemented core components such as Client, Config, and various request/response entities to facilitate API interactions. - Added structured error handling with specific error types for improved clarity and debugging. - Included a README and usage guide to assist developers in integrating the SDK into their applications, ensuring a smooth onboarding experience. - Established a fluent API for chat interactions and streaming responses, enhancing usability and flexibility for developers.
This commit is contained in:
n.nakhostin
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user