mirror of
https://github.com/fumiama/deepinfra.git
synced 2026-06-05 16:50:35 +08:00
113 lines
2.4 KiB
Go
113 lines
2.4 KiB
Go
package model
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"io"
|
|
"net/http"
|
|
)
|
|
|
|
const (
|
|
ModelDeepDeek = "deepseek-ai/DeepSeek-R1"
|
|
)
|
|
|
|
type Message struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
}
|
|
|
|
type Choice struct {
|
|
Index int `json:"index"`
|
|
Message Message `json:"message"`
|
|
FinishReason string `json:"finish_reason"`
|
|
}
|
|
|
|
// OpenAI as an specified example.
|
|
type OpenAI struct {
|
|
sep string
|
|
Protocol `json:"-"`
|
|
// callback only
|
|
ID string `json:"id,omitempty"`
|
|
Object string `json:"object,omitempty"`
|
|
Created int `json:"created,omitempty"`
|
|
Choices []Choice `json:"choices,omitempty"`
|
|
// callback/request
|
|
Model string `json:"model"`
|
|
Messages []Message `json:"messages"`
|
|
Temperature float32 `json:"temperature"` // Temperature 0.7
|
|
TopP float32 `json:"top_p"` // TopP 0.9
|
|
MaxTokens int `json:"max_tokens"` // MaxTokens 4096
|
|
}
|
|
|
|
// NewOpenAI use temp 0.7, topp 0.9, maxn 4096 if you don't know the meaning.
|
|
func NewOpenAI(model, sep string, temp, topp float32, maxn uint) *OpenAI {
|
|
opai := new(OpenAI)
|
|
opai.sep = sep
|
|
opai.Model = model
|
|
opai.Temperature = temp
|
|
opai.TopP = topp
|
|
opai.MaxTokens = int(maxn)
|
|
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)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return w
|
|
}
|
|
|
|
func (opai *OpenAI) Parse(body io.Reader) error {
|
|
return json.NewDecoder(body).Decode(&opai)
|
|
}
|
|
|
|
func (opai *OpenAI) Output() string {
|
|
if len(opai.Choices) == 0 {
|
|
return ""
|
|
}
|
|
return CutLast(opai.Choices[len(opai.Choices)-1].Message.Content, opai.sep)
|
|
}
|
|
|
|
func (opai *OpenAI) OutputRaw() string {
|
|
if len(opai.Choices) == 0 {
|
|
return ""
|
|
}
|
|
return opai.Choices[len(opai.Choices)-1].Message.Content
|
|
}
|
|
|
|
func (opai *OpenAI) System(prompt string) Protocol {
|
|
opai.Messages = make([]Message, 1, 8)
|
|
opai.Messages[0] = Message{
|
|
Role: "system",
|
|
Content: prompt,
|
|
}
|
|
return opai
|
|
}
|
|
|
|
func (opai *OpenAI) User(prompt string) Protocol {
|
|
opai.Messages = append(opai.Messages, Message{
|
|
Role: "user",
|
|
Content: prompt,
|
|
})
|
|
return opai
|
|
}
|
|
|
|
func (opai *OpenAI) Assistant(prompt string) Protocol {
|
|
opai.Messages = append(opai.Messages, Message{
|
|
Role: "assistant",
|
|
Content: prompt,
|
|
})
|
|
return opai
|
|
}
|