vertex.go

package iteragent

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

type VertexConfig struct {
	ProjectID   string
	Location    string
	Model       string
	Credentials string
	MaxTokens   int
	Temperature float32
}

type VertexProvider struct {
	config VertexConfig
	client *http.Client
}

func NewVertex(config VertexConfig) *VertexProvider {
	return &VertexProvider{
		config: config,
		client: &http.Client{Timeout: 120 * time.Second},
	}
}

func (p *VertexProvider) Name() string {
	return fmt.Sprintf("vertex(%s)", p.config.Model)
}

func (p *VertexProvider) getAccessToken(ctx context.Context) (string, error) {
	tokenSrc := os.Getenv("GOOGLE_ACCESS_TOKEN")
	if tokenSrc != "" {
		return tokenSrc, nil
	}

	credFile := p.config.Credentials
	if credFile == "" {
		credFile = os.Getenv("GOOGLE_APPLICATION_CREDENTIALS")
	}

	if credFile != "" {
		return "", fmt.Errorf("service account credentials file found but JWT signing is not implemented; use gcloud auth or set GOOGLE_ACCESS_TOKEN instead")
	}

	return "", fmt.Errorf("no credentials found for Vertex AI")
}

func (p *VertexProvider) Complete(ctx context.Context, messages []Message, opts ...CompletionOptions) (string, error) {
	location := p.config.Location
	if location == "" {
		location = "us-central1"
	}

	url := fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/google/models/%s:generateContent",
		location, p.config.ProjectID, location, p.config.Model)

	token, err := p.getAccessToken(ctx)
	if err != nil {
		return "", err
	}

	var system string
	var contents []map[string]interface{}
	for _, m := range messages {
		if m.Role == "system" {
			system = m.Content
		} else {
			role := "user"
			if m.Role == "assistant" {
				role = "model"
			}
			contents = append(contents, map[string]interface{}{
				"role": role,
				"parts": []map[string]string{
					{"text": m.Content},
				},
			})
		}
	}

	body := map[string]interface{}{
		"contents": contents,
	}
	if system != "" {
		body["systemInstruction"] = map[string]interface{}{
			"parts": []map[string]string{
				{"text": system},
			},
		}
	}
	if p.config.MaxTokens > 0 {
		body["maxOutputTokens"] = p.config.MaxTokens
	}
	if p.config.Temperature > 0 {
		body["temperature"] = p.config.Temperature
	}

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

	req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(jsonBody))
	if err != nil {
		return "", err
	}

	req.Header.Set("Authorization", "Bearer "+token)
	req.Header.Set("Content-Type", "application/json")

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

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("read response body: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("Vertex AI error (%d): %s", resp.StatusCode, string(respBody))
	}

	var response struct {
		Candidates []struct {
			Content struct {
				Parts []struct {
					Text string `json:"text"`
				} `json:"parts"`
			} `json:"content"`
		} `json:"candidates"`
	}

	if err := json.Unmarshal(respBody, &response); err != nil {
		return "", fmt.Errorf("parse response: %w", err)
	}

	if len(response.Candidates) == 0 {
		return "", fmt.Errorf("no response candidates")
	}

	texts := []string{}
	for _, part := range response.Candidates[0].Content.Parts {
		texts = append(texts, part.Text)
	}

	return strings.Join(texts, ""), nil
}

// CompleteStream implements Provider for Vertex AI using the streamGenerateContent SSE endpoint.
func (p *VertexProvider) CompleteStream(ctx context.Context, messages []Message, opt CompletionOptions, onToken func(string)) (string, error) {
	location := p.config.Location
	if location == "" {
		location = "us-central1"
	}
	streamURL := fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/google/models/%s:streamGenerateContent",
		location, p.config.ProjectID, location, p.config.Model)

	accessToken, err := p.getAccessToken(ctx)
	if err != nil {
		return "", err
	}

	var system string
	var contents []map[string]interface{}
	for _, m := range messages {
		if m.Role == "system" {
			system = m.Content
		} else {
			role := "user"
			if m.Role == "assistant" {
				role = "model"
			}
			contents = append(contents, map[string]interface{}{
				"role":  role,
				"parts": []map[string]string{{"text": m.Content}},
			})
		}
	}
	body := map[string]interface{}{"contents": contents}
	if system != "" {
		body["systemInstruction"] = map[string]interface{}{
			"parts": []map[string]string{{"text": system}},
		}
	}
	if opt.MaxTokens > 0 {
		body["maxOutputTokens"] = opt.MaxTokens
	}
	if opt.Temperature > 0 {
		body["temperature"] = opt.Temperature
	}
	jsonBody, err := json.Marshal(body)
	if err != nil {
		return "", fmt.Errorf("marshal request: %w", err)
	}

	var full strings.Builder
	sseClient := NewSSEClient()
	err = sseClient.Stream(ctx, streamURL, map[string]string{"Authorization": "Bearer " + accessToken}, jsonBody, func(e SSEEvent) {
		if tok, ok := ParseGeminiSSE(e.Data); ok && tok != "" {
			full.WriteString(tok)
			if onToken != nil {
				onToken(tok)
			}
		}
	})
	if err != nil {
		return "", fmt.Errorf("vertex stream: %w", err)
	}
	result := full.String()
	if result == "" {
		return "", fmt.Errorf("empty streaming response from vertex")
	}
	return result, nil
}