loveuer
98 lines
2.1 KiB
Go
98 lines
2.1 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
)
|
|
|
|
const (
|
|
RoleSystem = "system"
|
|
RoleUser = "user"
|
|
RoleAssistant = "assistant"
|
|
RoleTool = "tool"
|
|
)
|
|
|
|
type Provider interface {
|
|
ChatStream(ctx context.Context, req ChatRequest, emit func(StreamEvent) error) error
|
|
}
|
|
|
|
type ChatRequest struct {
|
|
Model string `json:"model"`
|
|
Messages []Message `json:"messages"`
|
|
Tools []Tool `json:"tools,omitempty"`
|
|
ToolChoice string `json:"tool_choice,omitempty"`
|
|
Stream bool `json:"stream"`
|
|
Extra map[string]any `json:"-"`
|
|
}
|
|
|
|
func (r ChatRequest) MarshalJSON() ([]byte, error) {
|
|
payload := map[string]any{
|
|
"model": r.Model,
|
|
"messages": r.Messages,
|
|
"stream": r.Stream,
|
|
}
|
|
if len(r.Tools) > 0 {
|
|
payload["tools"] = r.Tools
|
|
}
|
|
if r.ToolChoice != "" {
|
|
payload["tool_choice"] = r.ToolChoice
|
|
}
|
|
for key, value := range r.Extra {
|
|
if _, exists := payload[key]; exists {
|
|
continue
|
|
}
|
|
payload[key] = value
|
|
}
|
|
return json.Marshal(payload)
|
|
}
|
|
|
|
type Message struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content,omitempty"`
|
|
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
|
ToolCallID string `json:"tool_call_id,omitempty"`
|
|
}
|
|
|
|
type ToolCall struct {
|
|
ID string `json:"id,omitempty"`
|
|
Type string `json:"type,omitempty"`
|
|
Function FunctionCall `json:"function,omitempty"`
|
|
}
|
|
|
|
type FunctionCall struct {
|
|
Name string `json:"name,omitempty"`
|
|
Arguments string `json:"arguments,omitempty"`
|
|
}
|
|
|
|
type Tool struct {
|
|
Type string `json:"type"`
|
|
Function ToolFunction `json:"function"`
|
|
}
|
|
|
|
type ToolFunction struct {
|
|
Name string `json:"name"`
|
|
Description string `json:"description"`
|
|
Parameters json.RawMessage `json:"parameters"`
|
|
}
|
|
|
|
type StreamEvent struct {
|
|
Content string
|
|
ToolCalls []ToolCallDelta
|
|
FinishReason string
|
|
Usage *Usage
|
|
}
|
|
|
|
type Usage struct {
|
|
PromptTokens int `json:"prompt_tokens"`
|
|
CompletionTokens int `json:"completion_tokens"`
|
|
TotalTokens int `json:"total_tokens"`
|
|
}
|
|
|
|
type ToolCallDelta struct {
|
|
Index int
|
|
ID string
|
|
Type string
|
|
Name string
|
|
Arguments string
|
|
}
|