petal/internal/llm/vllm.go

package llm

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

// VLLMClient talks to an OpenAI-compatible /v1/chat/completions endpoint
// (vLLM's API server). It is the default backend.
type VLLMClient struct {
	backend
}

// vllmRequest is the OpenAI-compatible request body. repetition_penalty is a
// vLLM extension to the OpenAI schema, which vLLM accepts.
type vllmRequest struct {
	Model             string    `json:"model"`
	Messages          []Message `json:"messages"`
	MaxTokens         int       `json:"max_tokens"`
	Temperature       float64   `json:"temperature"`
	RepetitionPenalty float64   `json:"repetition_penalty"`
	TopP              float64   `json:"top_p"`
	Stop              []string  `json:"stop,omitempty"`
	Stream            bool      `json:"stream"`
}

func (c *VLLMClient) body(req CompletionRequest) vllmRequest {
	return vllmRequest{
		Model:             c.model(req),
		Messages:          req.Messages,
		MaxTokens:         req.MaxTokens,
		Temperature:       req.Temperature,
		RepetitionPenalty: req.RepetitionPenalty,
		TopP:              req.TopP,
		Stop:              req.Stop,
		Stream:            req.Stream,
	}
}

// Complete sends a non-streaming request and returns the full message content.
func (c *VLLMClient) Complete(ctx context.Context, req CompletionRequest) (string, error) {
	req.Stream = false
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()
	resp, err := c.post(ctx, c.body(req))
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return "", httpError("vllm", resp)
	}

	var out struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return "", fmt.Errorf("vllm: decode response: %w", err)
	}
	if len(out.Choices) == 0 {
		return "", fmt.Errorf("vllm: empty choices in response")
	}
	return out.Choices[0].Message.Content, nil
}

// Stream sends a streaming request and emits delta content chunks on the
// returned channel, which closes when the stream ends ([DONE]) or ctx is done.
func (c *VLLMClient) Stream(ctx context.Context, req CompletionRequest) (<-chan string, error) {
	req.Stream = true
	resp, err := c.post(ctx, c.body(req))
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != http.StatusOK {
		defer resp.Body.Close()
		return nil, httpError("vllm", resp)
	}

	ch := make(chan string)
	go func() {
		defer close(ch)
		defer resp.Body.Close()
		sc := bufio.NewScanner(resp.Body)
		sc.Buffer(make([]byte, 0, 64*1024), 1024*1024)
		for sc.Scan() {
			line := strings.TrimSpace(sc.Text())
			if !strings.HasPrefix(line, "data:") {
				continue
			}
			data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
			if data == "[DONE]" {
				return
			}
			var chunk struct {
				Choices []struct {
					Delta struct {
						Content string `json:"content"`
					} `json:"delta"`
				} `json:"choices"`
			}
			if err := json.Unmarshal([]byte(data), &chunk); err != nil {
				continue // skip keep-alives / malformed frames
			}
			if len(chunk.Choices) == 0 {
				continue
			}
			if text := chunk.Choices[0].Delta.Content; text != "" {
				select {
				case ch <- text:
				case <-ctx.Done():
					return
				}
			}
		}
	}()
	return ch, nil
}

func (c *VLLMClient) post(ctx context.Context, body any) (*http.Response, error) {
	return postJSON(ctx, c.endpoint+"/v1/chat/completions", body)
}

// --- shared HTTP helpers (used by both backends) ---------------------------

// llmHTTP has no client-level timeout: Complete bounds itself with a context
// deadline (LLM_TIMEOUT), and streaming requests must stay open as long as the
// model is generating. Cancellation flows through the request context.
var llmHTTP = &http.Client{}

func postJSON(ctx context.Context, url string, body any) (*http.Response, error) {
	buf, err := json.Marshal(body)
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(buf))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	return llmHTTP.Do(req)
}

func httpError(backend string, resp *http.Response) error {
	b, _ := io.ReadAll(io.LimitReader(resp.Body, 2048))
	return fmt.Errorf("%s: %s: %s", backend, resp.Status, strings.TrimSpace(string(b)))
}