From db7712b6d253dd947eee083c78c48183c49e6b98 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?= <41315874+fumiama@users.noreply.github.com> Date: Sun, 30 Mar 2025 01:48:52 +0900 Subject: [PATCH] feat: add Google GenAI API --- README.md | 2 +- api.go | 8 ++-- model.go | 1 + model/api.go | 7 +++ model/genai.go | 120 ++++++++++++++++++++++++++++++++++++++++++++++++ model/ollama.go | 10 ++++ model/openai.go | 10 ++++ 7 files changed, 153 insertions(+), 5 deletions(-) create mode 100644 model/genai.go diff --git a/README.md b/README.md index 9343010..b8e9e8c 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@ Call OpenAI compatible APIs, originally designed for DeepInfra. ## Quick Start ```go -api := NewAPI(APIDeepInfra, "PUT YOUR API KEY HERE") +api := NewAPI(OpenAIDeepInfra, "PUT YOUR API KEY HERE") txt, err := api.Request(model.NewOpenAI(model.ModelDeepDeek, model.SeparatorThink, 0.7, 0.9, 1024). System("Be a good assistant.").User("Hello"), ) diff --git a/api.go b/api.go index 2543afe..7f0a11a 100644 --- a/api.go +++ b/api.go @@ -8,7 +8,8 @@ import ( ) const ( - APIDeepInfra = "https://api.deepinfra.com/v1/openai/chat/completions" + OpenAIDeepInfra = "https://api.deepinfra.com/v1/openai/chat/completions" + GenAIGoogle = "https://generativelanguage.googleapis.com/v1beta" ) type API struct { @@ -21,12 +22,11 @@ func NewAPI(api, key string) API { } func (api *API) Request(model Model) (string, error) { - req, err := http.NewRequest("POST", api.api, model.Body()) + req, err := http.NewRequest("POST", model.API(api.api, api.key), model.Body()) if err != nil { return "", err } - req.Header.Add("Content-Type", "application/json") - req.Header.Add("Authorization", "Bearer "+api.key) + model.Header(api.key, req.Header) resp, err := http.DefaultClient.Do(req) if err != nil { return "", err diff --git a/model.go b/model.go index 35d13f8..17bfdb6 100644 --- a/model.go +++ b/model.go @@ -7,4 +7,5 @@ import ( type Model interface { model.Inputer model.Outputer + model.Requester } diff --git a/model/api.go b/model/api.go index 6ea493a..a73af94 100644 --- a/model/api.go +++ b/model/api.go @@ -3,6 +3,7 @@ package model import ( "bytes" "io" + "net/http" ) type Inputer interface { @@ -15,6 +16,11 @@ type Outputer interface { OutputRaw() string } +type Requester interface { + API(api, key string) string // API decorator + Header(key string, h http.Header) // Header decorator +} + type MessageBuilder[T any] interface { System(prompt string) T User(prompt string) T @@ -24,5 +30,6 @@ type MessageBuilder[T any] interface { type Protocol interface { Inputer Outputer + Requester MessageBuilder[Protocol] } diff --git a/model/genai.go b/model/genai.go new file mode 100644 index 0000000..b2b66b9 --- /dev/null +++ b/model/genai.go @@ -0,0 +1,120 @@ +package model + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" +) + +const ( + ModelGemini15Flash = "models/gemini-1.5-flash" +) + +type Text struct { + Text string `json:"text"` +} + +type Content struct { + Parts []Text `json:"parts"` + Role string `json:"role,omitempty"` +} + +func (c *Content) String() string { + sb := strings.Builder{} + for _, p := range c.Parts { + sb.WriteString(p.Text) + } + return sb.String() +} + +type Candidate struct { + Content Content `json:"content"` + FinishReason string `json:"finishReason"` + Index int `json:"index"` +} + +// GenAI is Goole API format +type GenAI struct { + model string `json:"-"` + Protocol `json:"-"` + // request only + Contents []Content `json:"contents,omitempty"` + SystemInstruction *Content `json:"systemInstruction,omitempty"` + GenerationConfig struct { + Temperature float32 `json:"temperature,omitempty"` + ResponseMimeType string `json:"responseMimeType,omitempty"` + TopP float32 `json:"topP,omitempty"` + MaxOutputTokens int `json:"maxOutputTokens,omitempty"` + } `json:"generationConfig"` + // callback only + Candidates []Candidate `json:"candidates,omitempty"` +} + +// NewGenAI use temp 0.7, topp 0.9, maxn 4096 if you don't know the meaning. +func NewGenAI(model string, temp, topp float32, maxn uint) *GenAI { + opai := new(GenAI) + opai.model = model + opai.GenerationConfig.Temperature = temp + opai.GenerationConfig.ResponseMimeType = "text/plain" + opai.GenerationConfig.TopP = topp + opai.GenerationConfig.MaxOutputTokens = int(maxn) + return opai +} + +func (opai *GenAI) API(api, key string) string { + return fmt.Sprintf("%s/%s:generateContent?key=%s", api, opai.model, key) +} + +func (*GenAI) Header(_ string, h http.Header) { + h.Add("Content-Type", "application/json") +} + +func (opai *GenAI) Body() *bytes.Buffer { + w := bytes.NewBuffer(make([]byte, 0, 8192)) + err := json.NewEncoder(w).Encode(opai) + if err != nil { + panic(err) + } + return w +} + +func (opai *GenAI) Parse(body io.Reader) error { + return json.NewDecoder(body).Decode(&opai) +} + +func (opai *GenAI) Output() string { + if len(opai.Candidates) == 0 { + return "" + } + return opai.Candidates[0].Content.String() +} + +func (opai *GenAI) OutputRaw() string { + return opai.Output() +} + +func (opai *GenAI) System(prompt string) Protocol { + opai.SystemInstruction = &Content{ + Parts: []Text{{prompt}}, + } + return opai +} + +func (opai *GenAI) User(prompt string) Protocol { + opai.Contents = append(opai.Contents, Content{ + Parts: []Text{{prompt}}, + Role: "user", + }) + return opai +} + +func (opai *GenAI) Assistant(prompt string) Protocol { + opai.Contents = append(opai.Contents, Content{ + Parts: []Text{{prompt}}, + Role: "model", + }) + return opai +} diff --git a/model/ollama.go b/model/ollama.go index c3615f9..a45779b 100644 --- a/model/ollama.go +++ b/model/ollama.go @@ -4,6 +4,7 @@ import ( "bytes" "encoding/json" "io" + "net/http" ) // OLLaMA as an specified example. @@ -35,6 +36,15 @@ func NewOLLaMA(model, sep string, temp, topp float32, maxn uint) *OLLaMA { return opai } +func (*OLLaMA) API(api, _ string) string { + return api +} + +func (*OLLaMA) Header(key string, h http.Header) { + h.Add("Content-Type", "application/json") + h.Add("Authorization", "Bearer "+key) +} + func (ollm *OLLaMA) Body() *bytes.Buffer { w := bytes.NewBuffer(make([]byte, 0, 8192)) err := json.NewEncoder(w).Encode(ollm) diff --git a/model/openai.go b/model/openai.go index 5458f89..2f59a76 100644 --- a/model/openai.go +++ b/model/openai.go @@ -4,6 +4,7 @@ import ( "bytes" "encoding/json" "io" + "net/http" ) const ( @@ -49,6 +50,15 @@ func NewOpenAI(model, sep string, temp, topp float32, maxn uint) *OpenAI { return opai } +func (*OpenAI) API(api, _ string) string { + return api +} + +func (*OpenAI) Header(key string, h http.Header) { + h.Add("Content-Type", "application/json") + h.Add("Authorization", "Bearer "+key) +} + func (opai *OpenAI) Body() *bytes.Buffer { w := bytes.NewBuffer(make([]byte, 0, 8192)) err := json.NewEncoder(w).Encode(opai)