go-micro.dev/v5@v5.12.0/genai/gemini/gemini.go (about)

     1  package gemini
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"encoding/json"
     7  	"fmt"
     8  	"net/http"
     9  	"os"
    10  
    11  	"go-micro.dev/v5/genai"
    12  )
    13  
    14  // gemini implements the GenAI interface using Google Gemini 2.5 API.
    15  type gemini struct {
    16  	options genai.Options
    17  }
    18  
    19  func New(opts ...genai.Option) genai.GenAI {
    20  	var options genai.Options
    21  	for _, o := range opts {
    22  		o(&options)
    23  	}
    24  	if options.APIKey == "" {
    25  		options.APIKey = os.Getenv("GEMINI_API_KEY")
    26  	}
    27  	return &gemini{options: options}
    28  }
    29  
    30  func (g *gemini) Generate(prompt string, opts ...genai.Option) (*genai.Result, error) {
    31  	options := g.options
    32  	for _, o := range opts {
    33  		o(&options)
    34  	}
    35  	ctx := context.Background()
    36  
    37  	res := &genai.Result{Prompt: prompt, Type: options.Type}
    38  
    39  	endpoint := options.Endpoint
    40  	if endpoint == "" {
    41  		endpoint = "https://generativelanguage.googleapis.com/v1beta/models/"
    42  	}
    43  
    44  	var url string
    45  	var body map[string]interface{}
    46  
    47  	// Determine model to use
    48  	var model string
    49  	switch options.Type {
    50  	case "image":
    51  		if options.Model != "" {
    52  			model = options.Model
    53  		} else {
    54  			model = "gemini-2.5-pro-vision"
    55  		}
    56  		url = endpoint + model + ":generateContent?key=" + options.APIKey
    57  		body = map[string]interface{}{
    58  			"contents": []map[string]interface{}{
    59  				{"parts": []map[string]string{{"text": prompt}}},
    60  			},
    61  		}
    62  	case "audio":
    63  		if options.Model != "" {
    64  			model = options.Model
    65  		} else {
    66  			model = "gemini-2.5-pro"
    67  		}
    68  		url = endpoint + model + ":generateContent?key=" + options.APIKey
    69  		body = map[string]interface{}{
    70  			"contents": []map[string]interface{}{
    71  				{"parts": []map[string]string{{"text": prompt}}},
    72  			},
    73  			"response_mime_type": "audio/wav",
    74  		}
    75  	case "text":
    76  		fallthrough
    77  	default:
    78  		if options.Model != "" {
    79  			model = options.Model
    80  		} else {
    81  			model = "gemini-2.5-pro"
    82  		}
    83  		url = endpoint + model + ":generateContent?key=" + options.APIKey
    84  		body = map[string]interface{}{
    85  			"contents": []map[string]interface{}{
    86  				{"parts": []map[string]string{{"text": prompt}}},
    87  			},
    88  		}
    89  	}
    90  
    91  	b, _ := json.Marshal(body)
    92  	req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(b))
    93  	if err != nil {
    94  		return nil, err
    95  	}
    96  	req.Header.Set("Content-Type", "application/json")
    97  
    98  	resp, err := http.DefaultClient.Do(req)
    99  	if err != nil {
   100  		return nil, err
   101  	}
   102  	defer resp.Body.Close()
   103  
   104  	if options.Type == "audio" {
   105  		var result struct {
   106  			Candidates []struct {
   107  				Content struct {
   108  					Parts []struct {
   109  						InlineData struct {
   110  							Data []byte `json:"data"`
   111  						} `json:"inline_data"`
   112  					} `json:"parts"`
   113  				} `json:"content"`
   114  			} `json:"candidates"`
   115  		}
   116  		if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
   117  			return nil, err
   118  		}
   119  		if len(result.Candidates) == 0 || len(result.Candidates[0].Content.Parts) == 0 {
   120  			return nil, fmt.Errorf("no audio returned")
   121  		}
   122  		res.Data = result.Candidates[0].Content.Parts[0].InlineData.Data
   123  		return res, nil
   124  	}
   125  
   126  	var result struct {
   127  		Candidates []struct {
   128  			Content struct {
   129  				Parts []struct {
   130  					Text string `json:"text"`
   131  				} `json:"parts"`
   132  			} `json:"content"`
   133  		} `json:"candidates"`
   134  	}
   135  	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
   136  		return nil, err
   137  	}
   138  	if len(result.Candidates) == 0 || len(result.Candidates[0].Content.Parts) == 0 {
   139  		return nil, fmt.Errorf("no candidates returned")
   140  	}
   141  	res.Text = result.Candidates[0].Content.Parts[0].Text
   142  	return res, nil
   143  }
   144  
   145  func (g *gemini) Stream(prompt string, opts ...genai.Option) (*genai.Stream, error) {
   146  	results := make(chan *genai.Result)
   147  	go func() {
   148  		defer close(results)
   149  		res, err := g.Generate(prompt, opts...)
   150  		if err != nil {
   151  			// Send error via Stream.Err, not channel
   152  			return
   153  		}
   154  		results <- res
   155  	}()
   156  	return &genai.Stream{Results: results}, nil
   157  }
   158  
   159  func (g *gemini) String() string {
   160  	return "gemini"
   161  }
   162  
   163  func init() {
   164  	genai.Register("gemini", New())
   165  }