llama-link/internal/llama/proxy.go

package llama

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"strings"
	"time"
)

type Proxy struct {
	manager *Manager
	client  *http.Client
}

func NewProxy(manager *Manager) *Proxy {
	return &Proxy{
		manager: manager,
		client: &http.Client{
			Timeout: 60 * time.Second,
		},
	}
}

func (p *Proxy) Manager() *Manager {
	return p.manager
}

type ChatMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
	Name    string `json:"name,omitempty"`
}

type ChatCompletionRequest struct {
	Model    string        `json:"model"`
	Messages []ChatMessage `json:"messages"`
	Stream   bool          `json:"stream,omitempty"`
	MaxTokens int          `json:"max_tokens,omitempty"`
	Temperature float64    `json:"temperature,omitempty"`
	TopP      float64      `json:"top_p,omitempty"`
}

type ChatCompletionResponse struct {
	ID      string   `json:"id"`
	Object  string   `json:"object"`
	Created int64    `json:"created"`
	Model   string   `json:"model"`
	Choices []Choice `json:"choices"`
	Usage   Usage    `json:"usage"`
}

type Choice struct {
	Index        int       `json:"index"`
	Message      ChatMessage `json:"message"`
	FinishReason string    `json:"finish_reason"`
}

type Usage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens     int `json:"total_tokens"`
}

type StreamChoice struct {
	Index        int           `json:"index"`
	Delta        ChatMessage   `json:"delta"`
	FinishReason string        `json:"finish_reason,omitempty"`
}

type StreamResponse struct {
	ID      string          `json:"id"`
	Object string          `json:"object"`
	Created int64          `json:"created"`
	Model  string          `json:"model"`
	Choices []StreamChoice `json:"choices"`
}

// ChatCompletion calls llama-server and returns the response
func (p *Proxy) ChatCompletion(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionResponse, error) {
	if !p.manager.IsReady() {
		return nil, ErrServerNotRunning
	}

	// Convert to llama-server format
	llamaReq := map[string]interface{}{
		"model":    req.Model,
		"messages": req.Messages,
		"stream":   false,
	}

	body, err := json.Marshal(llamaReq)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request: %w", err)
	}

	httpReq, err := http.NewRequestWithContext(ctx, "POST", p.manager.url+"/v1/chat/completions", bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	httpReq.Header.Set("Content-Type", "application/json")

	resp, err := p.client.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("llama-server request failed: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		bodyBytes, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("llama-server returned %d: %s", resp.StatusCode, string(bodyBytes))
	}

	var result ChatCompletionResponse
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return nil, fmt.Errorf("failed to decode response: %w", err)
	}

	return &result, nil
}

// ChatCompletionStream returns a channel of streaming responses
func (p *Proxy) ChatCompletionStream(ctx context.Context, req ChatCompletionRequest) (<-chan *StreamResponse, <-chan error) {
	stream := make(chan *StreamResponse, 100)
	errCh := make(chan error, 1)

	if !p.manager.IsReady() {
		errCh <- ErrServerNotRunning
		close(stream)
		return stream, errCh
	}

	go func() {
		defer close(stream)
		defer close(errCh)

		llamaReq := map[string]interface{}{
			"model":    req.Model,
			"messages": req.Messages,
			"stream":   true,
		}

		body, err := json.Marshal(llamaReq)
		if err != nil {
			errCh <- err
			return
		}

		httpReq, err := http.NewRequestWithContext(ctx, "POST", p.manager.url+"/v1/chat/completions", bytes.NewReader(body))
		if err != nil {
			errCh <- err
			return
		}
		httpReq.Header.Set("Content-Type", "application/json")

		resp, err := p.client.Do(httpReq)
		if err != nil {
			errCh <- err
			return
		}
		defer resp.Body.Close()

		if resp.StatusCode != http.StatusOK {
			bodyBytes, _ := io.ReadAll(resp.Body)
			errCh <- fmt.Errorf("llama-server returned %d: %s", resp.StatusCode, string(bodyBytes))
			return
		}

		reader := bufio.NewReader(resp.Body)
		for {
			line, err := reader.ReadString('\n')
			if err != nil {
				if err != io.EOF {
					errCh <- err
				}
				break
			}

			line = strings.TrimSpace(line)
			if line == "" || !strings.HasPrefix(line, "data: ") {
				continue
			}

			if line == "data: [DONE]" {
				break
			}

			data := strings.TrimPrefix(line, "data: ")
			var streamResp StreamResponse
			if err := json.Unmarshal([]byte(data), &streamResp); err != nil {
				slog.Debug("failed to parse stream chunk", "error", err, "data", data)
				continue
			}

			select {
			case stream <- &streamResp:
			case <-ctx.Done():
				return
			}
		}
	}()

	return stream, errCh
}

// ModelsList returns available models from llama-server
func (p *Proxy) ModelsList(ctx context.Context) ([]string, error) {
	if !p.manager.IsReady() {
		return nil, ErrServerNotRunning
	}

	req, err := http.NewRequestWithContext(ctx, "GET", p.manager.url+"/v1/models", nil)
	if err != nil {
		return nil, err
	}

	resp, err := p.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("llama-server returned %d", resp.StatusCode)
	}

	var result struct {
		Data []struct {
			ID string `json:"id"`
		} `json:"data"`
	}

	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return nil, err
	}

	models := make([]string, len(result.Data))
	for i, m := range result.Data {
		models[i] = m.ID
	}
	return models, nil
}