1
0
mirror of https://github.com/fumiama/deepinfra.git synced 2026-06-05 00:32:46 +08:00

feat: add ollama-compatibility

This commit is contained in:
源文雨
2025-02-22 01:02:48 +09:00
parent 6b99264c6b
commit ce98a6b278
7 changed files with 209 additions and 205 deletions

View File

@@ -21,13 +21,8 @@ type MessageBuilder[T any] interface {
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"`
type Protocol[T any] interface {
Inputer
Outputer
MessageBuilder[T]
}

View File

@@ -1,89 +0,0 @@
package model
import (
"bytes"
"encoding/json"
"io"
)
// Custom as an compatible example.
type Custom struct {
Inputer `json:"-"`
Outputer `json:"-"`
MessageBuilder[*DeepSeek] `json:"-"`
sep string `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 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

@@ -1,92 +0,0 @@
package model
import (
"bytes"
"encoding/json"
"io"
)
const (
modelDeepDeek = "deepseek-ai/DeepSeek-R1"
)
// DeepSeek as an specified example.
type DeepSeek struct {
Inputer `json:"-"`
Outputer `json:"-"`
MessageBuilder[*DeepSeek] `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 16384
}
// NewDeepSeek 0.7, 0.9
func NewDeepSeek(temp, topp float32, maxn uint) *DeepSeek {
ds := new(DeepSeek)
ds.Model = modelDeepDeek
ds.Temperature = temp
ds.TopP = topp
ds.MaxTokens = int(maxn)
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)
}
func (ds *DeepSeek) Output() string {
if len(ds.Choices) == 0 {
return ""
}
return CutLast(ds.Choices[len(ds.Choices)-1].Message.Content, SeparatorThink)
}
func (ds *DeepSeek) OutputRaw() string {
if len(ds.Choices) == 0 {
return ""
}
return ds.Choices[len(ds.Choices)-1].Message.Content
}
func (ds *DeepSeek) System(prompt string) *DeepSeek {
ds.Messages = make([]Message, 1, 8)
ds.Messages[0] = Message{
Role: "system",
Content: prompt,
}
return ds
}
func (ds *DeepSeek) User(prompt string) *DeepSeek {
ds.Messages = append(ds.Messages, Message{
Role: "user",
Content: prompt,
})
return ds
}
func (ds *DeepSeek) Assistant(prompt string) *DeepSeek {
ds.Messages = append(ds.Messages, Message{
Role: "assistant",
Content: prompt,
})
return ds
}

87
model/ollama.go Normal file
View File

@@ -0,0 +1,87 @@
package model
import (
"bytes"
"encoding/json"
"io"
)
// OLLaMA as an specified example.
type OLLaMA struct {
sep string
Protocol[*OLLaMA] `json:"-"`
// callback only
ID string `json:"id,omitempty"`
Object string `json:"object,omitempty"`
Created int `json:"created,omitempty"`
Messages []Message `json:"messages"`
// callback/request
Model string `json:"model"`
Message *Message `json:"message,omitempty"`
Temperature float32 `json:"temperature"` // Temperature 0.7
TopP float32 `json:"top_p"` // TopP 0.9
MaxTokens int `json:"max_tokens"` // MaxTokens 4096
}
// NewOLLaMA use temp 0.7, topp 0.9, maxn 4096 if you don't know the meaning.
func NewOLLaMA(model, sep string, temp, topp float32, maxn uint) *OLLaMA {
opai := new(OLLaMA)
opai.sep = sep
opai.Model = model
opai.Temperature = temp
opai.TopP = topp
opai.MaxTokens = int(maxn)
return opai
}
func (ollm *OLLaMA) Body() *bytes.Buffer {
w := bytes.NewBuffer(make([]byte, 0, 8192))
err := json.NewEncoder(w).Encode(ollm)
if err != nil {
panic(err)
}
return w
}
func (ollm *OLLaMA) Parse(body io.Reader) error {
return json.NewDecoder(body).Decode(&ollm)
}
func (ollm *OLLaMA) Output() string {
if ollm.Message == nil {
return ""
}
return CutLast(ollm.Message.Content, ollm.sep)
}
func (ollm *OLLaMA) OutputRaw() string {
if ollm.Message == nil {
return ""
}
return ollm.Message.Content
}
func (ollm *OLLaMA) System(prompt string) *OLLaMA {
ollm.Messages = make([]Message, 1, 8)
ollm.Messages[0] = Message{
Role: "system",
Content: prompt,
}
return ollm
}
func (ollm *OLLaMA) User(prompt string) *OLLaMA {
ollm.Messages = append(ollm.Messages, Message{
Role: "user",
Content: prompt,
})
return ollm
}
func (ollm *OLLaMA) Assistant(prompt string) *OLLaMA {
ollm.Messages = append(ollm.Messages, Message{
Role: "assistant",
Content: prompt,
})
return ollm
}

102
model/openai.go Normal file
View File

@@ -0,0 +1,102 @@
package model
import (
"bytes"
"encoding/json"
"io"
)
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[*OpenAI] `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 (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) *OpenAI {
opai.Messages = make([]Message, 1, 8)
opai.Messages[0] = Message{
Role: "system",
Content: prompt,
}
return opai
}
func (opai *OpenAI) User(prompt string) *OpenAI {
opai.Messages = append(opai.Messages, Message{
Role: "user",
Content: prompt,
})
return opai
}
func (opai *OpenAI) Assistant(prompt string) *OpenAI {
opai.Messages = append(opai.Messages, Message{
Role: "assistant",
Content: prompt,
})
return opai
}