465 lines
10 KiB
Go
465 lines
10 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
// 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 {
|
|
Model string `json:"model"`
|
|
Choices []struct {
|
|
Message Message `json:"message"`
|
|
} `json:"choices"`
|
|
}
|
|
|
|
// StreamChunk represents a chunk from streaming chat
|
|
type StreamChunk struct {
|
|
Model string `json:"model"`
|
|
Choices []struct {
|
|
Delta struct {
|
|
Content string `json:"content"`
|
|
} `json:"delta"`
|
|
} `json:"choices"`
|
|
}
|
|
|
|
// UncloseAI is the main client for OpenAI-compatible APIs
|
|
type UncloseAI struct {
|
|
models []ModelInfo
|
|
ttsEndpoints []string
|
|
apiKey string
|
|
timeout time.Duration
|
|
httpClient *http.Client
|
|
}
|
|
|
|
// NewUncloseAI creates a new client with auto-discovery from environment variables
|
|
func NewUncloseAI() (*UncloseAI, error) {
|
|
return NewUncloseAIWithOptions(nil, nil, "")
|
|
}
|
|
|
|
// NewUncloseAIWithOptions creates a new client with custom endpoints
|
|
func NewUncloseAIWithOptions(modelEndpoints, ttsEndpoints []string, apiKey string) (*UncloseAI, error) {
|
|
client := &UncloseAI{
|
|
models: make([]ModelInfo, 0),
|
|
ttsEndpoints: make([]string, 0),
|
|
apiKey: apiKey,
|
|
timeout: 30 * time.Second,
|
|
httpClient: &http.Client{Timeout: 30 * time.Second},
|
|
}
|
|
|
|
// Discover endpoints from environment if not provided
|
|
if modelEndpoints == nil {
|
|
modelEndpoints = discoverEnvEndpoints("MODEL_ENDPOINT")
|
|
}
|
|
if ttsEndpoints == nil {
|
|
ttsEndpoints = discoverEnvEndpoints("TTS_ENDPOINT")
|
|
}
|
|
|
|
// Discover models from each endpoint
|
|
for _, endpoint := range modelEndpoints {
|
|
if err := client.discoverModelsFromEndpoint(endpoint); err != nil {
|
|
fmt.Printf("Warning: 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 *UncloseAI) discoverModelsFromEndpoint(endpoint string) error {
|
|
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,
|
|
})
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// ListModels returns all discovered models
|
|
func (c *UncloseAI) ListModels() []ModelInfo {
|
|
return c.models
|
|
}
|
|
|
|
// Chat performs non-streaming chat completion
|
|
func (c *UncloseAI) Chat(ctx context.Context, messages []Message, model string, maxTokens int, temperature float64) (*ChatResponse, error) {
|
|
modelInfo, err := c.getModelInfo(model)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
payload := map[string]interface{}{
|
|
"model": modelInfo.ID,
|
|
"messages": messages,
|
|
"max_tokens": maxTokens,
|
|
"temperature": temperature,
|
|
"stream": false,
|
|
}
|
|
|
|
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, err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
bodyBytes, _ := io.ReadAll(resp.Body)
|
|
return nil, fmt.Errorf("unexpected status %d: %s", resp.StatusCode, string(bodyBytes))
|
|
}
|
|
|
|
var chatResp ChatResponse
|
|
if err := json.NewDecoder(resp.Body).Decode(&chatResp); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &chatResp, nil
|
|
}
|
|
|
|
// ChatStream performs streaming chat completion, returning a channel of chunks
|
|
func (c *UncloseAI) ChatStream(ctx context.Context, messages []Message, model string, maxTokens int, temperature float64) (<-chan StreamChunk, <-chan error) {
|
|
chunkChan := make(chan StreamChunk)
|
|
errChan := make(chan error, 1)
|
|
|
|
go func() {
|
|
defer close(chunkChan)
|
|
defer close(errChan)
|
|
|
|
modelInfo, err := c.getModelInfo(model)
|
|
if err != nil {
|
|
errChan <- err
|
|
return
|
|
}
|
|
|
|
payload := map[string]interface{}{
|
|
"model": modelInfo.ID,
|
|
"messages": messages,
|
|
"max_tokens": maxTokens,
|
|
"temperature": temperature,
|
|
"stream": true,
|
|
}
|
|
|
|
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 <- err
|
|
return
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
bodyBytes, _ := io.ReadAll(resp.Body)
|
|
errChan <- fmt.Errorf("unexpected status %d: %s", 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 {
|
|
continue // Skip malformed chunks
|
|
}
|
|
|
|
select {
|
|
case chunkChan <- chunk:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
if err := scanner.Err(); err != nil {
|
|
errChan <- err
|
|
}
|
|
}()
|
|
|
|
return chunkChan, errChan
|
|
}
|
|
|
|
// TTS generates speech from text
|
|
func (c *UncloseAI) TTS(text, voice, model string) ([]byte, error) {
|
|
if len(c.ttsEndpoints) == 0 {
|
|
return nil, fmt.Errorf("no TTS endpoints available")
|
|
}
|
|
|
|
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.NewRequest("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, err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return nil, fmt.Errorf("unexpected status: %d", resp.StatusCode)
|
|
}
|
|
|
|
return io.ReadAll(resp.Body)
|
|
}
|
|
|
|
// getModelInfo retrieves model info by ID or returns first available
|
|
func (c *UncloseAI) getModelInfo(modelID string) (*ModelInfo, error) {
|
|
if len(c.models) == 0 {
|
|
return nil, fmt.Errorf("no models available")
|
|
}
|
|
|
|
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("model '%s' not found", modelID)
|
|
}
|
|
|
|
// Demo usage
|
|
func main() {
|
|
fmt.Println("=== UncloseAI Go Client (with Streaming) ===\n")
|
|
|
|
// Initialize client with auto-discovery
|
|
client, err := NewUncloseAI()
|
|
if err != nil {
|
|
fmt.Printf("Error initializing client: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
models := client.ListModels()
|
|
if len(models) == 0 {
|
|
fmt.Println("ERROR: No models discovered. Set environment variables:")
|
|
fmt.Println(" MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc.")
|
|
os.Exit(1)
|
|
}
|
|
|
|
fmt.Printf("Discovered %d model(s)\n", len(models))
|
|
for _, m := range models {
|
|
fmt.Printf(" - %s (max_tokens: %d)\n", m.ID, m.MaxTokens)
|
|
}
|
|
fmt.Println()
|
|
|
|
ctx := context.Background()
|
|
|
|
// Non-streaming chat example
|
|
fmt.Println("=== Non-Streaming Chat ===")
|
|
response, err := client.Chat(
|
|
ctx,
|
|
[]Message{
|
|
{Role: "system", Content: "You are a helpful AI assistant."},
|
|
{Role: "user", Content: "Explain quantum computing in one sentence."},
|
|
},
|
|
"", // Use first available model
|
|
100,
|
|
0.7,
|
|
)
|
|
|
|
if err != nil {
|
|
fmt.Printf("Error: %v\n", err)
|
|
} else {
|
|
fmt.Printf("Model: %s\n", response.Model)
|
|
fmt.Printf("Response: %s\n\n", response.Choices[0].Message.Content)
|
|
}
|
|
|
|
// Streaming chat example
|
|
fmt.Println("=== Streaming Chat ===")
|
|
modelID := ""
|
|
if len(models) > 1 {
|
|
modelID = models[1].ID
|
|
}
|
|
|
|
fmt.Printf("Model: %s\n", modelID)
|
|
fmt.Print("Response: ")
|
|
|
|
chunkChan, errChan := client.ChatStream(
|
|
ctx,
|
|
[]Message{
|
|
{Role: "system", Content: "You are a coding assistant."},
|
|
{Role: "user", Content: "Write a Go function to check if a number is prime"},
|
|
},
|
|
modelID,
|
|
200,
|
|
0.7,
|
|
)
|
|
|
|
for {
|
|
select {
|
|
case chunk, ok := <-chunkChan:
|
|
if !ok {
|
|
goto done
|
|
}
|
|
if len(chunk.Choices) > 0 {
|
|
content := chunk.Choices[0].Delta.Content
|
|
if content != "" {
|
|
fmt.Print(content)
|
|
}
|
|
}
|
|
case err := <-errChan:
|
|
if err != nil {
|
|
fmt.Printf("\nError: %v\n", err)
|
|
}
|
|
goto done
|
|
}
|
|
}
|
|
|
|
done:
|
|
fmt.Println("\n")
|
|
|
|
// TTS example
|
|
if len(client.ttsEndpoints) > 0 {
|
|
fmt.Println("=== TTS Speech Generation ===")
|
|
audioData, err := client.TTS(
|
|
"Hello from UncloseAI Go client! This demonstrates text to speech with streaming support.",
|
|
"alloy",
|
|
"tts-1",
|
|
)
|
|
|
|
if err != nil {
|
|
fmt.Printf("Error: %v\n", err)
|
|
} else {
|
|
if err := os.WriteFile("speech.mp3", audioData, 0644); err != nil {
|
|
fmt.Printf("[ERROR] Failed to write speech file: %v\n", err)
|
|
} else {
|
|
fmt.Printf("[OK] Speech file created: speech.mp3 (%d bytes)\n\n", len(audioData))
|
|
}
|
|
}
|
|
}
|
|
|
|
fmt.Println("=== Examples Complete ===")
|
|
}
|