chore: import upstream snapshot with attribution

ai/gemini/gemini.go

@@ -0,0 +1,358 @@
+// Package gemini implements the Google Gemini model provider.
+//
+// Usage:
+//
+//	import _ "go-micro.dev/v6/ai/gemini"
+//
+//	m := ai.New("gemini",
+//	    ai.WithAPIKey("your-api-key"),
+//	)
+package gemini
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+
+	"go-micro.dev/v6/ai"
+)
+
+func init() {
+	ai.Register("gemini", func(opts ...ai.Option) ai.Model {
+		return NewProvider(opts...)
+	})
+	ai.RegisterStream("gemini")
+}
+
+// Provider implements the ai.Model interface for Google Gemini.
+type Provider struct {
+	opts ai.Options
+}
+
+// NewProvider creates a new Gemini provider.
+func NewProvider(opts ...ai.Option) *Provider {
+	options := ai.NewOptions(opts...)
+
+	if options.Model == "" {
+		options.Model = "gemini-2.5-flash"
+	}
+	if options.BaseURL == "" {
+		options.BaseURL = "https://generativelanguage.googleapis.com"
+	}
+
+	return &Provider{opts: options}
+}
+
+func (p *Provider) Init(opts ...ai.Option) error {
+	for _, o := range opts {
+		o(&p.opts)
+	}
+	return nil
+}
+
+func (p *Provider) Options() ai.Options { return p.opts }
+func (p *Provider) String() string      { return "gemini" }
+
+func (p *Provider) Generate(ctx context.Context, req *ai.Request, opts ...ai.GenerateOption) (*ai.Response, error) {
+	var tools []map[string]any
+	for _, t := range req.Tools {
+		tools = append(tools, map[string]any{
+			"name":        t.Name,
+			"description": t.Description,
+			"parameters": map[string]any{
+				"type":       "object",
+				"properties": t.Properties,
+			},
+		})
+	}
+
+	contents := geminiContents(req)
+
+	apiReq := map[string]any{
+		"contents": contents,
+	}
+
+	if req.SystemPrompt != "" {
+		apiReq["system_instruction"] = map[string]any{
+			"parts": []map[string]any{{"text": req.SystemPrompt}},
+		}
+	}
+
+	if len(tools) > 0 {
+		apiReq["tools"] = []map[string]any{
+			{"functionDeclarations": tools},
+		}
+	}
+
+	resp, rawParts, err := p.callAPI(ctx, apiReq)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(resp.ToolCalls) == 0 {
+		return resp, nil
+	}
+
+	if p.opts.ToolHandler != nil {
+		var resultParts []map[string]any
+		for _, tc := range resp.ToolCalls {
+			result := p.opts.ToolHandler(ctx, tc).Value
+			resultParts = append(resultParts, map[string]any{
+				"functionResponse": map[string]any{
+					"name":     tc.Name,
+					"id":       tc.ID,
+					"response": result,
+				},
+			})
+		}
+
+		followUpContents := append(contents,
+			map[string]any{"role": "model", "parts": rawParts},
+			map[string]any{"role": "user", "parts": resultParts},
+		)
+
+		followUpReq := map[string]any{
+			"contents": followUpContents,
+		}
+		if req.SystemPrompt != "" {
+			followUpReq["system_instruction"] = map[string]any{
+				"parts": []map[string]any{{"text": req.SystemPrompt}},
+			}
+		}
+
+		followUpResp, _, err := p.callAPI(ctx, followUpReq)
+		if err == nil && followUpResp.Reply != "" {
+			resp.Answer = followUpResp.Reply
+		}
+	}
+
+	return resp, nil
+}
+
+func (p *Provider) Stream(ctx context.Context, req *ai.Request, opts ...ai.GenerateOption) (ai.Stream, error) {
+	apiReq := map[string]any{
+		"contents": geminiContents(req),
+	}
+	if req.SystemPrompt != "" {
+		apiReq["system_instruction"] = map[string]any{
+			"parts": []map[string]any{{"text": req.SystemPrompt}},
+		}
+	}
+	if p.opts.MaxTokens > 0 {
+		apiReq["generationConfig"] = map[string]any{"maxOutputTokens": p.opts.MaxTokens}
+	}
+
+	reqBody, err := json.Marshal(apiReq)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal stream request: %w", err)
+	}
+
+	apiURL := strings.TrimRight(p.opts.BaseURL, "/") +
+		"/v1beta/models/" + p.opts.Model + ":streamGenerateContent?alt=sse"
+	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(reqBody))
+	if err != nil {
+		return nil, fmt.Errorf("failed to create stream request: %w", err)
+	}
+	httpReq.Header.Set("Content-Type", "application/json")
+	httpReq.Header.Set("Accept", "text/event-stream")
+	httpReq.Header.Set("x-goog-api-key", p.opts.APIKey)
+
+	httpResp, err := http.DefaultClient.Do(httpReq)
+	if err != nil {
+		return nil, fmt.Errorf("stream API request failed: %w", err)
+	}
+	if httpResp.StatusCode != http.StatusOK {
+		defer httpResp.Body.Close()
+		respBody, _ := io.ReadAll(httpResp.Body)
+		return nil, ai.NewHTTPError(httpResp, respBody)
+	}
+	return &streamReader{body: httpResp.Body, scanner: bufio.NewScanner(httpResp.Body)}, nil
+}
+
+type streamReader struct {
+	body    io.ReadCloser
+	scanner *bufio.Scanner
+	closed  bool
+}
+
+func (s *streamReader) Recv() (*ai.Response, error) {
+	for s.scanner.Scan() {
+		line := strings.TrimSpace(s.scanner.Text())
+		if line == "" || strings.HasPrefix(line, ":") || strings.HasPrefix(line, "event:") {
+			continue
+		}
+		if !strings.HasPrefix(line, "data:") {
+			continue
+		}
+		data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
+		if data == "[DONE]" {
+			return nil, io.EOF
+		}
+
+		var chunk struct {
+			Error *struct {
+				Code    int    `json:"code"`
+				Message string `json:"message"`
+				Status  string `json:"status"`
+			} `json:"error"`
+			Candidates []struct {
+				Content struct {
+					Parts []struct {
+						Text string `json:"text"`
+					} `json:"parts"`
+				} `json:"content"`
+			} `json:"candidates"`
+			UsageMetadata *struct {
+				PromptTokenCount     int `json:"promptTokenCount"`
+				CandidatesTokenCount int `json:"candidatesTokenCount"`
+				TotalTokenCount      int `json:"totalTokenCount"`
+			} `json:"usageMetadata"`
+		}
+		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
+			return nil, fmt.Errorf("failed to parse stream chunk: %w", err)
+		}
+		if chunk.Error != nil {
+			return nil, fmt.Errorf("gemini stream error (%s): %s", chunk.Error.Status, chunk.Error.Message)
+		}
+		for _, candidate := range chunk.Candidates {
+			var parts []string
+			for _, part := range candidate.Content.Parts {
+				if part.Text != "" {
+					parts = append(parts, part.Text)
+				}
+			}
+			if len(parts) > 0 {
+				return &ai.Response{Reply: strings.Join(parts, "")}, nil
+			}
+		}
+		if chunk.UsageMetadata != nil {
+			return &ai.Response{Usage: ai.Usage{
+				InputTokens:  chunk.UsageMetadata.PromptTokenCount,
+				OutputTokens: chunk.UsageMetadata.CandidatesTokenCount,
+				TotalTokens:  chunk.UsageMetadata.TotalTokenCount,
+			}}, nil
+		}
+	}
+	if err := s.scanner.Err(); err != nil {
+		return nil, err
+	}
+	return nil, io.EOF
+}
+
+func (s *streamReader) Close() error {
+	if s.closed {
+		return nil
+	}
+	s.closed = true
+	return s.body.Close()
+}
+
+func (p *Provider) callAPI(ctx context.Context, req map[string]any) (*ai.Response, []map[string]any, error) {
+	reqBody, err := json.Marshal(req)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to marshal request: %w", err)
+	}
+
+	apiURL := strings.TrimRight(p.opts.BaseURL, "/") +
+		"/v1beta/models/" + p.opts.Model + ":generateContent"
+
+	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(reqBody))
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to create request: %w", err)
+	}
+
+	httpReq.Header.Set("Content-Type", "application/json")
+	httpReq.Header.Set("x-goog-api-key", p.opts.APIKey)
+
+	httpResp, err := http.DefaultClient.Do(httpReq)
+	if err != nil {
+		return nil, nil, fmt.Errorf("API request failed: %w", err)
+	}
+	defer httpResp.Body.Close()
+
+	respBody, _ := io.ReadAll(httpResp.Body)
+	if httpResp.StatusCode != http.StatusOK {
+		return nil, nil, ai.NewHTTPError(httpResp, respBody)
+	}
+
+	var geminiResp struct {
+		Candidates []struct {
+			Content struct {
+				Parts []struct {
+					Text         string          `json:"text"`
+					FunctionCall *functionCallPB `json:"functionCall"`
+				} `json:"parts"`
+			} `json:"content"`
+		} `json:"candidates"`
+	}
+
+	if err := json.Unmarshal(respBody, &geminiResp); err != nil {
+		return nil, nil, fmt.Errorf("failed to parse response: %w", err)
+	}
+
+	if len(geminiResp.Candidates) == 0 {
+		return nil, nil, fmt.Errorf("no response from API")
+	}
+
+	parts := geminiResp.Candidates[0].Content.Parts
+	response := &ai.Response{}
+
+	var replyParts []string
+	var rawParts []map[string]any
+
+	for _, part := range parts {
+		if part.Text != "" {
+			replyParts = append(replyParts, part.Text)
+			rawParts = append(rawParts, map[string]any{"text": part.Text})
+		}
+		if part.FunctionCall != nil {
+			response.ToolCalls = append(response.ToolCalls, ai.ToolCall{
+				ID:    part.FunctionCall.ID,
+				Name:  part.FunctionCall.Name,
+				Input: part.FunctionCall.Args,
+			})
+			rawParts = append(rawParts, map[string]any{
+				"functionCall": map[string]any{
+					"id":   part.FunctionCall.ID,
+					"name": part.FunctionCall.Name,
+					"args": part.FunctionCall.Args,
+				},
+			})
+		}
+	}
+
+	if len(replyParts) > 0 {
+		response.Reply = strings.Join(replyParts, "\n")
+	}
+
+	return response, rawParts, nil
+}
+
+type functionCallPB struct {
+	ID   string         `json:"id"`
+	Name string         `json:"name"`
+	Args map[string]any `json:"args"`
+}
+
+func geminiContents(req *ai.Request) []map[string]any {
+	contents := make([]map[string]any, 0, len(req.Messages)+1)
+	for _, m := range req.Messages {
+		role := m.Role
+		if role == "assistant" {
+			role = "model"
+		}
+		if role == "system" || role == "" {
+			continue
+		}
+		contents = append(contents, map[string]any{"role": role, "parts": []map[string]any{{"text": fmt.Sprint(m.Content)}}})
+	}
+	if req.Prompt != "" {
+		contents = append(contents, map[string]any{"role": "user", "parts": []map[string]any{{"text": req.Prompt}}})
+	}
+	return contents
+}

ai/gemini/gemini_test.go

@@ -0,0 +1,202 @@
+package gemini
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"go-micro.dev/v6/ai"
+)
+
+func TestProvider_String(t *testing.T) {
+	p := NewProvider()
+	if p.String() != "gemini" {
+		t.Errorf("Expected provider name 'gemini', got '%s'", p.String())
+	}
+}
+
+func TestProvider_Init(t *testing.T) {
+	p := NewProvider()
+
+	err := p.Init(
+		ai.WithModel("gemini-2.0-flash"),
+		ai.WithAPIKey("test-key"),
+		ai.WithBaseURL("https://test.com"),
+	)
+
+	if err != nil {
+		t.Fatalf("Init failed: %v", err)
+	}
+
+	opts := p.Options()
+	if opts.Model != "gemini-2.0-flash" {
+		t.Errorf("Expected model 'gemini-2.0-flash', got '%s'", opts.Model)
+	}
+	if opts.APIKey != "test-key" {
+		t.Errorf("Expected API key 'test-key', got '%s'", opts.APIKey)
+	}
+	if opts.BaseURL != "https://test.com" {
+		t.Errorf("Expected base URL 'https://test.com', got '%s'", opts.BaseURL)
+	}
+}
+
+func TestProvider_Options(t *testing.T) {
+	p := NewProvider(
+		ai.WithModel("custom-model"),
+		ai.WithAPIKey("my-key"),
+	)
+
+	opts := p.Options()
+	if opts.Model != "custom-model" {
+		t.Errorf("Expected model 'custom-model', got '%s'", opts.Model)
+	}
+	if opts.APIKey != "my-key" {
+		t.Errorf("Expected API key 'my-key', got '%s'", opts.APIKey)
+	}
+}
+
+func TestProvider_Defaults(t *testing.T) {
+	p := NewProvider()
+
+	opts := p.Options()
+	if opts.Model != "gemini-2.5-flash" {
+		t.Errorf("Expected default model 'gemini-2.5-flash', got '%s'", opts.Model)
+	}
+	if opts.BaseURL != "https://generativelanguage.googleapis.com" {
+		t.Errorf("Expected default base URL 'https://generativelanguage.googleapis.com', got '%s'", opts.BaseURL)
+	}
+}
+
+func TestProvider_Generate_NoAPIKey(t *testing.T) {
+	p := NewProvider()
+
+	req := &ai.Request{
+		Prompt:       "Hello",
+		SystemPrompt: "You are helpful",
+	}
+
+	_, err := p.Generate(context.Background(), req)
+	if err == nil {
+		t.Error("Expected error when API key is missing, got nil")
+	}
+}
+
+func TestProvider_Stream(t *testing.T) {
+	var sawRequest bool
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		sawRequest = true
+		if r.URL.Path != "/v1beta/models/gemini-2.5-flash:streamGenerateContent" {
+			t.Fatalf("path = %s, want streamGenerateContent", r.URL.Path)
+		}
+		if r.URL.Query().Get("alt") != "sse" {
+			t.Fatalf("alt = %q, want sse", r.URL.Query().Get("alt"))
+		}
+		if got := r.Header.Get("Accept"); got != "text/event-stream" {
+			t.Fatalf("Accept = %q, want text/event-stream", got)
+		}
+		if got := r.Header.Get("x-goog-api-key"); got != "test-key" {
+			t.Fatalf("x-goog-api-key = %q, want test-key", got)
+		}
+
+		var body map[string]any
+		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+			t.Fatalf("decode request: %v", err)
+		}
+		contents, ok := body["contents"].([]any)
+		if !ok || len(contents) != 3 {
+			t.Fatalf("contents = %#v, want history + prompt", body["contents"])
+		}
+		second := contents[1].(map[string]any)
+		if second["role"] != "model" {
+			t.Fatalf("assistant history role = %#v, want model", second["role"])
+		}
+
+		w.Header().Set("Content-Type", "text/event-stream")
+		_, _ = w.Write([]byte("data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"hel\"}]}}]}\n\n"))
+		_, _ = w.Write([]byte("data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"lo\"}]}}],\"usageMetadata\":{\"promptTokenCount\":3,\"candidatesTokenCount\":2,\"totalTokenCount\":5}}\n\n"))
+		_, _ = w.Write([]byte("data: [DONE]\n\n"))
+	}))
+	defer ts.Close()
+
+	p := NewProvider(ai.WithAPIKey("test-key"), ai.WithBaseURL(ts.URL))
+	stream, err := p.Stream(context.Background(), &ai.Request{
+		Messages: []ai.Message{
+			{Role: "user", Content: "previous question"},
+			{Role: "assistant", Content: "previous answer"},
+		},
+		Prompt: "Hello",
+	})
+	if err != nil {
+		t.Fatalf("Stream returned error: %v", err)
+	}
+	defer stream.Close()
+	if !sawRequest {
+		t.Fatal("server did not receive stream request")
+	}
+
+	first, err := stream.Recv()
+	if err != nil || first.Reply != "hel" {
+		t.Fatalf("first chunk = %#v, %v; want hel", first, err)
+	}
+	second, err := stream.Recv()
+	if err != nil || second.Reply != "lo" {
+		t.Fatalf("second chunk = %#v, %v; want lo", second, err)
+	}
+	if _, err := stream.Recv(); !errors.Is(err, io.EOF) {
+		t.Fatalf("final error = %v, want EOF", err)
+	}
+}
+
+func TestProvider_StreamPropagatesMalformedChunk(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "text/event-stream")
+		_, _ = w.Write([]byte("data: {bad json}\n\n"))
+	}))
+	defer ts.Close()
+
+	p := NewProvider(ai.WithAPIKey("test-key"), ai.WithBaseURL(ts.URL))
+	stream, err := p.Stream(context.Background(), &ai.Request{Prompt: "Hello"})
+	if err != nil {
+		t.Fatalf("Stream returned error: %v", err)
+	}
+	defer stream.Close()
+
+	if _, err := stream.Recv(); err == nil {
+		t.Fatal("Recv returned nil error for malformed chunk")
+	}
+}
+
+func TestProvider_StreamPropagatesProviderError(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		http.Error(w, "quota exhausted", http.StatusTooManyRequests)
+	}))
+	defer ts.Close()
+
+	p := NewProvider(ai.WithAPIKey("test-key"), ai.WithBaseURL(ts.URL))
+	stream, err := p.Stream(context.Background(), &ai.Request{Prompt: "Hello"})
+	if err == nil {
+		_ = stream.Close()
+		t.Fatal("Stream returned nil error for provider failure")
+	}
+	if !strings.Contains(err.Error(), "429") || !strings.Contains(err.Error(), "quota exhausted") {
+		t.Fatalf("Stream error = %v, want provider status and body", err)
+	}
+	if strings.Contains(err.Error(), "test-key") {
+		t.Fatal("stream error leaked API key")
+	}
+}
+
+func TestProvider_Registration(t *testing.T) {
+	m := ai.New("gemini", ai.WithAPIKey("test"))
+	if m == nil {
+		t.Fatal("ai.New('gemini') returned nil — provider not registered")
+	}
+	if m.String() != "gemini" {
+		t.Errorf("Expected 'gemini', got '%s'", m.String())
+	}
+}