470 lines
12 KiB
Go
470 lines
12 KiB
Go
// This is free software for the public good of a permacomputer hosted at
|
|
// permacomputer.com, an always-on computer by the people, for the people.
|
|
// One which is durable, easy to repair, & distributed like tap water
|
|
// for machine learning intelligence.
|
|
//
|
|
// The permacomputer is community-owned infrastructure optimized around
|
|
// four values:
|
|
//
|
|
// TRUTH First principles, math & science, open source code freely distributed
|
|
// FREEDOM Voluntary partnerships, freedom from tyranny & corporate control
|
|
// HARMONY Minimal waste, self-renewing systems with diverse thriving connections
|
|
// LOVE Be yourself without hurting others, cooperation through natural law
|
|
//
|
|
// This software contributes to that vision by making machine learning
|
|
// accessible to everyone through a free, open, embeddable chat interface.
|
|
// Code is seeds to sprout on any abandoned technology.
|
|
|
|
package uncloseai
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
// Error types
|
|
type Error string
|
|
|
|
const (
|
|
ErrConnection Error = "connection error"
|
|
ErrModelNotFound Error = "model not found"
|
|
ErrStreaming Error = "streaming error"
|
|
ErrNoModels Error = "no models available"
|
|
ErrNoTTSEndpoints Error = "no TTS endpoints available"
|
|
ErrInvalidResponse Error = "invalid response"
|
|
)
|
|
|
|
func (e Error) Error() string {
|
|
return string(e)
|
|
}
|
|
|
|
// ModelInfo contains metadata about a discovered model
|
|
type ModelInfo struct {
|
|
ID string `json:"id"`
|
|
Endpoint string `json:"endpoint"`
|
|
MaxTokens int `json:"max_tokens"`
|
|
}
|
|
|
|
// Message represents a chat message
|
|
type Message struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
}
|
|
|
|
// ChatResponse is the response from a non-streaming chat completion
|
|
type ChatResponse struct {
|
|
ID string `json:"id"`
|
|
Model string `json:"model"`
|
|
Choices []ChatChoice `json:"choices"`
|
|
}
|
|
|
|
// ChatChoice represents a single choice in the response
|
|
type ChatChoice struct {
|
|
Index int `json:"index"`
|
|
Message Message `json:"message"`
|
|
FinishReason string `json:"finish_reason"`
|
|
}
|
|
|
|
// StreamChunk represents a chunk from streaming chat
|
|
type StreamChunk struct {
|
|
ID string `json:"id"`
|
|
Model string `json:"model"`
|
|
Choices []StreamChunkChoice `json:"choices"`
|
|
}
|
|
|
|
// StreamChunkChoice represents a single choice in a streaming chunk
|
|
type StreamChunkChoice struct {
|
|
Index int `json:"index"`
|
|
Delta StreamDelta `json:"delta"`
|
|
}
|
|
|
|
// StreamDelta represents the delta content in a streaming chunk
|
|
type StreamDelta struct {
|
|
Role string `json:"role,omitempty"`
|
|
Content string `json:"content,omitempty"`
|
|
}
|
|
|
|
// Config holds configuration options for the UncloseAI client
|
|
type Config struct {
|
|
// Endpoints for model discovery (nil = auto-discover from environment)
|
|
Endpoints []string
|
|
// TTS endpoints (nil = auto-discover from environment)
|
|
TTSEndpoints []string
|
|
// API key for authentication (empty = no authentication)
|
|
APIKey string
|
|
// HTTP client timeout (0 = 30 seconds default)
|
|
Timeout time.Duration
|
|
// Enable debug logging
|
|
Debug bool
|
|
}
|
|
|
|
// ChatOptions holds options for chat completions
|
|
type ChatOptions struct {
|
|
// Model ID (empty = auto-select first available)
|
|
Model string
|
|
// Maximum tokens to generate (0 = no limit)
|
|
MaxTokens int
|
|
// Sampling temperature (0.0 - 2.0)
|
|
Temperature float64
|
|
// Nucleus sampling (0.0 - 1.0)
|
|
TopP float64
|
|
}
|
|
|
|
// Client is the main UncloseAI client
|
|
type Client struct {
|
|
models []ModelInfo
|
|
ttsEndpoints []string
|
|
apiKey string
|
|
timeout time.Duration
|
|
httpClient *http.Client
|
|
debug bool
|
|
}
|
|
|
|
// New creates a new UncloseAI client with optional configuration.
|
|
// If config is nil, defaults are used and endpoints are auto-discovered from environment variables.
|
|
func New(config *Config) (*Client, error) {
|
|
if config == nil {
|
|
config = &Config{}
|
|
}
|
|
|
|
timeout := config.Timeout
|
|
if timeout == 0 {
|
|
timeout = 30 * time.Second
|
|
}
|
|
|
|
client := &Client{
|
|
models: make([]ModelInfo, 0),
|
|
ttsEndpoints: make([]string, 0),
|
|
apiKey: config.APIKey,
|
|
timeout: timeout,
|
|
httpClient: &http.Client{Timeout: timeout},
|
|
debug: config.Debug,
|
|
}
|
|
|
|
// Discover endpoints from environment if not provided
|
|
endpoints := config.Endpoints
|
|
if endpoints == nil {
|
|
endpoints = discoverEnvEndpoints("MODEL_ENDPOINT")
|
|
}
|
|
|
|
ttsEndpoints := config.TTSEndpoints
|
|
if ttsEndpoints == nil {
|
|
ttsEndpoints = discoverEnvEndpoints("TTS_ENDPOINT")
|
|
}
|
|
|
|
if client.debug {
|
|
fmt.Printf("[DEBUG] Initialized with %d model endpoint(s) and %d TTS endpoint(s)\n",
|
|
len(endpoints), len(ttsEndpoints))
|
|
}
|
|
|
|
// Discover models from each endpoint
|
|
for _, endpoint := range endpoints {
|
|
if err := client.discoverModelsFromEndpoint(endpoint); err != nil {
|
|
if client.debug {
|
|
fmt.Printf("[DEBUG] Failed to discover models from %s: %v\n", endpoint, err)
|
|
}
|
|
}
|
|
}
|
|
|
|
client.ttsEndpoints = ttsEndpoints
|
|
|
|
return client, nil
|
|
}
|
|
|
|
// discoverEnvEndpoints finds endpoints from environment variables like PREFIX_1, PREFIX_2, ...
|
|
func discoverEnvEndpoints(prefix string) []string {
|
|
endpoints := make([]string, 0)
|
|
for i := 1; i <= 9999; i++ {
|
|
endpoint := os.Getenv(fmt.Sprintf("%s_%d", prefix, i))
|
|
if endpoint == "" {
|
|
break
|
|
}
|
|
endpoints = append(endpoints, endpoint)
|
|
}
|
|
return endpoints
|
|
}
|
|
|
|
// discoverModelsFromEndpoint discovers available models from an endpoint
|
|
func (c *Client) discoverModelsFromEndpoint(endpoint string) error {
|
|
if c.debug {
|
|
fmt.Printf("[DEBUG] Discovering models from: %s\n", endpoint)
|
|
}
|
|
|
|
req, err := http.NewRequest("GET", endpoint+"/models", nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if c.apiKey != "" {
|
|
req.Header.Set("Authorization", "Bearer "+c.apiKey)
|
|
}
|
|
|
|
resp, err := c.httpClient.Do(req)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return fmt.Errorf("unexpected status: %d", resp.StatusCode)
|
|
}
|
|
|
|
var result struct {
|
|
Data []struct {
|
|
ID string `json:"id"`
|
|
MaxModelLen int `json:"max_model_len"`
|
|
} `json:"data"`
|
|
}
|
|
|
|
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
|
return err
|
|
}
|
|
|
|
for _, model := range result.Data {
|
|
maxTokens := model.MaxModelLen
|
|
if maxTokens == 0 {
|
|
maxTokens = 8192
|
|
}
|
|
c.models = append(c.models, ModelInfo{
|
|
ID: model.ID,
|
|
Endpoint: endpoint,
|
|
MaxTokens: maxTokens,
|
|
})
|
|
|
|
if c.debug {
|
|
fmt.Printf("[DEBUG] Discovered: %s\n", model.ID)
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// ListModels returns all discovered models
|
|
func (c *Client) ListModels() []ModelInfo {
|
|
return c.models
|
|
}
|
|
|
|
// Chat performs a non-streaming chat completion
|
|
func (c *Client) Chat(ctx context.Context, messages []Message, options *ChatOptions) (*ChatResponse, error) {
|
|
if options == nil {
|
|
options = &ChatOptions{Temperature: 0.7, TopP: 1.0}
|
|
}
|
|
|
|
modelInfo, err := c.getModelInfo(options.Model)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
payload := map[string]interface{}{
|
|
"model": modelInfo.ID,
|
|
"messages": messages,
|
|
"temperature": options.Temperature,
|
|
"top_p": options.TopP,
|
|
"stream": false,
|
|
}
|
|
|
|
if options.MaxTokens > 0 {
|
|
payload["max_tokens"] = options.MaxTokens
|
|
}
|
|
|
|
body, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, "POST", modelInfo.Endpoint+"/chat/completions", bytes.NewBuffer(body))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
req.Header.Set("Content-Type", "application/json")
|
|
if c.apiKey != "" {
|
|
req.Header.Set("Authorization", "Bearer "+c.apiKey)
|
|
}
|
|
|
|
resp, err := c.httpClient.Do(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("%w: %v", ErrConnection, err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
bodyBytes, _ := io.ReadAll(resp.Body)
|
|
return nil, fmt.Errorf("%w: status %d: %s", ErrInvalidResponse, resp.StatusCode, string(bodyBytes))
|
|
}
|
|
|
|
var chatResp ChatResponse
|
|
if err := json.NewDecoder(resp.Body).Decode(&chatResp); err != nil {
|
|
return nil, fmt.Errorf("%w: %v", ErrInvalidResponse, err)
|
|
}
|
|
|
|
return &chatResp, nil
|
|
}
|
|
|
|
// ChatStream performs a streaming chat completion, returning channels for chunks and errors
|
|
func (c *Client) ChatStream(ctx context.Context, messages []Message, options *ChatOptions) (<-chan StreamChunk, <-chan error) {
|
|
chunkChan := make(chan StreamChunk)
|
|
errChan := make(chan error, 1)
|
|
|
|
go func() {
|
|
defer close(chunkChan)
|
|
defer close(errChan)
|
|
|
|
if options == nil {
|
|
options = &ChatOptions{Temperature: 0.7, TopP: 1.0}
|
|
}
|
|
|
|
modelInfo, err := c.getModelInfo(options.Model)
|
|
if err != nil {
|
|
errChan <- err
|
|
return
|
|
}
|
|
|
|
payload := map[string]interface{}{
|
|
"model": modelInfo.ID,
|
|
"messages": messages,
|
|
"temperature": options.Temperature,
|
|
"top_p": options.TopP,
|
|
"stream": true,
|
|
}
|
|
|
|
if options.MaxTokens > 0 {
|
|
payload["max_tokens"] = options.MaxTokens
|
|
}
|
|
|
|
body, err := json.Marshal(payload)
|
|
if err != nil {
|
|
errChan <- err
|
|
return
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, "POST", modelInfo.Endpoint+"/chat/completions", bytes.NewBuffer(body))
|
|
if err != nil {
|
|
errChan <- err
|
|
return
|
|
}
|
|
|
|
req.Header.Set("Content-Type", "application/json")
|
|
if c.apiKey != "" {
|
|
req.Header.Set("Authorization", "Bearer "+c.apiKey)
|
|
}
|
|
|
|
resp, err := c.httpClient.Do(req)
|
|
if err != nil {
|
|
errChan <- fmt.Errorf("%w: %v", ErrConnection, err)
|
|
return
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
bodyBytes, _ := io.ReadAll(resp.Body)
|
|
errChan <- fmt.Errorf("%w: status %d: %s", ErrStreaming, resp.StatusCode, string(bodyBytes))
|
|
return
|
|
}
|
|
|
|
// Parse SSE stream
|
|
scanner := bufio.NewScanner(resp.Body)
|
|
for scanner.Scan() {
|
|
line := scanner.Text()
|
|
|
|
// SSE format: "data: {...}"
|
|
if strings.HasPrefix(line, "data: ") {
|
|
data := strings.TrimPrefix(line, "data: ")
|
|
|
|
// Check for stream termination
|
|
if strings.TrimSpace(data) == "[DONE]" {
|
|
break
|
|
}
|
|
|
|
var chunk StreamChunk
|
|
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
|
|
if c.debug {
|
|
fmt.Printf("[DEBUG] Failed to parse chunk: %s\n", data)
|
|
}
|
|
continue // Skip malformed chunks
|
|
}
|
|
|
|
select {
|
|
case chunkChan <- chunk:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
if err := scanner.Err(); err != nil {
|
|
errChan <- fmt.Errorf("%w: %v", ErrStreaming, err)
|
|
}
|
|
}()
|
|
|
|
return chunkChan, errChan
|
|
}
|
|
|
|
// TTS generates speech from text
|
|
func (c *Client) TTS(ctx context.Context, text, voice, model string) ([]byte, error) {
|
|
if len(c.ttsEndpoints) == 0 {
|
|
return nil, ErrNoTTSEndpoints
|
|
}
|
|
|
|
endpoint := c.ttsEndpoints[0]
|
|
|
|
payload := map[string]interface{}{
|
|
"model": model,
|
|
"voice": voice,
|
|
"input": text,
|
|
}
|
|
|
|
body, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, "POST", endpoint+"/audio/speech", bytes.NewBuffer(body))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
req.Header.Set("Content-Type", "application/json")
|
|
if c.apiKey != "" {
|
|
req.Header.Set("Authorization", "Bearer "+c.apiKey)
|
|
}
|
|
|
|
resp, err := c.httpClient.Do(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("%w: %v", ErrConnection, err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return nil, fmt.Errorf("%w: status %d", ErrInvalidResponse, resp.StatusCode)
|
|
}
|
|
|
|
return io.ReadAll(resp.Body)
|
|
}
|
|
|
|
// getModelInfo retrieves model info by ID or returns first available
|
|
func (c *Client) getModelInfo(modelID string) (*ModelInfo, error) {
|
|
if len(c.models) == 0 {
|
|
return nil, ErrNoModels
|
|
}
|
|
|
|
if modelID == "" {
|
|
return &c.models[0], nil
|
|
}
|
|
|
|
for i := range c.models {
|
|
if c.models[i].ID == modelID {
|
|
return &c.models[i], nil
|
|
}
|
|
}
|
|
|
|
return nil, fmt.Errorf("%w: '%s'", ErrModelNotFound, modelID)
|
|
}
|