1
0
mirror of https://github.com/fumiama/deepinfra.git synced 2026-06-13 06:40:25 +08:00

feat: add custom

This commit is contained in:
源文雨
2025-02-14 16:29:37 +09:00
parent 507e407127
commit 12ba460588
6 changed files with 155 additions and 26 deletions

33
model/api.go Normal file
View File

@@ -0,0 +1,33 @@
package model
import (
"bytes"
"io"
)
type Inputer interface {
Body() *bytes.Buffer
Parse(io.Reader) error
}
type Outputer interface {
Output() string
OutputRaw() string
}
type MessageBuilder[T any] interface {
System(prompt string) T
User(prompt string) T
Assistant(prompt string) T
}
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"`
}

89
model/custom.go Normal file
View File

@@ -0,0 +1,89 @@
package model
import (
"bytes"
"encoding/json"
"io"
)
// Custom as an compatible example.
type Custom struct {
Inputer
Outputer
MessageBuilder[*DeepSeek]
sep string
// 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 16384
}
func NewCustom(model, sep string, temp, topp float32, maxn uint) *Custom {
c := new(Custom)
c.sep = sep
c.Model = model
c.Temperature = temp
c.TopP = topp
c.MaxTokens = int(maxn)
return c
}
func (c *Custom) Parse(body io.Reader) error {
return json.NewDecoder(body).Decode(&c)
}
func (c *Custom) Output() string {
if len(c.Choices) == 0 {
return ""
}
return CutLast(c.Choices[len(c.Choices)-1].Message.Content, c.sep)
}
func (c *Custom) OutputRaw() string {
if len(c.Choices) == 0 {
return ""
}
return c.Choices[len(c.Choices)-1].Message.Content
}
func (ds *Custom) System(prompt string) *Custom {
ds.Messages = make([]Message, 1, 8)
ds.Messages[0] = Message{
Role: "system",
Content: prompt,
}
return ds
}
func (ds *Custom) User(prompt string) *Custom {
ds.Messages = append(ds.Messages, Message{
Role: "user",
Content: prompt,
})
return ds
}
func (ds *Custom) Assistant(prompt string) *Custom {
ds.Messages = append(ds.Messages, Message{
Role: "assistant",
Content: prompt,
})
return ds
}
func (ds *Custom) Body() *bytes.Buffer {
w := bytes.NewBuffer(make([]byte, 0, 16384))
err := json.NewEncoder(w).Encode(ds)
if err != nil {
panic(err)
}
return w
}

View File

@@ -10,8 +10,11 @@ const (
modelDeepDeek = "deepseek-ai/DeepSeek-R1"
)
// DeepSeek as an example.
// DeepSeek as an specified example.
type DeepSeek struct {
Inputer
Outputer
MessageBuilder[*DeepSeek]
// callback only
ID string `json:"id,omitempty"`
Object string `json:"object,omitempty"`
@@ -26,17 +29,6 @@ type DeepSeek struct {
}
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"`
}
// NewDeepSeek 0.7, 0.9
func NewDeepSeek(temp, topp float32, maxn uint) *DeepSeek {
ds := new(DeepSeek)
@@ -47,6 +39,15 @@ func NewDeepSeek(temp, topp float32, maxn uint) *DeepSeek {
return ds
}
func (ds *DeepSeek) Body() *bytes.Buffer {
w := bytes.NewBuffer(make([]byte, 0, 16384))
err := json.NewEncoder(w).Encode(ds)
if err != nil {
panic(err)
}
return w
}
func (ds *DeepSeek) Parse(body io.Reader) error {
return json.NewDecoder(body).Decode(&ds)
}
@@ -89,12 +90,3 @@ func (ds *DeepSeek) Assistant(prompt string) *DeepSeek {
})
return ds
}
func (ds *DeepSeek) Body() *bytes.Buffer {
w := bytes.NewBuffer(make([]byte, 0, 16384))
err := json.NewEncoder(w).Encode(ds)
if err != nil {
panic(err)
}
return w
}

View File

@@ -7,6 +7,9 @@ const (
)
func CutLast(txt, sep string) string {
if sep == "" { // no need to cut
return txt
}
a := strings.LastIndex(txt, sep)
if a < 0 {
return ""