package proto

import (
	"github.com/ollama/ollama/api"
	"gorm.io/gorm"
)

type AIQueueMessage struct {
	Type string               `json:"type"` //声明不同消息类型
	ID   uint                 `json:"id"`   //消息id
	Msg  api.GenerateResponse `json:"msg"`  //原始消息
}

type WSMessage struct {
	Type      string         `json:"type"`       //接收及发送消息类型
	Msg       AIQueueMessage `json:"msg"`        //消息内容,只进行转发，不做处理
	SessionID int            `json:"session_id"` //应用层会话id
	ToID      int            `json:"to_id"`      //接收者id
}

type WSMessageReq struct {
	Type      string `json:"type"`       //接收及发送消息类型
	Function  string `json:"function"`   //功能名称
	Msg       string `json:"msg"`        //消息内容,只进行转发，不做处理
	SessionID int    `json:"session_id"` //应用层会话id
	ToID      int    `json:"to_id"`      //接收者id
}

type ModelParam struct {
	Model     string `json:"model"`     //模型名称
	Domain    string `json:"domain"`    //domain参数，spark需要
	Url       string `json:"url"`       //模型地址
	APPID     string `json:"appid"`     //应用id
	APISecret string `json:"apiSecret"` //应用密钥
	APIKey    string `json:"apiKey"`    //应用key
	System    string `json:"system"`    //
	//其他参数
	Temperature float64 `json:"temperature"` //温度,默认0.5,核采样阈值。取值越高随机性越强，即相同的问题得到的不同答案的可能性越大
	MaxTokens   int     `json:"maxTokens"`   //最大生成长度,默认4096
	TopK        float64 `json:"topK"`        //取概率最大的前k个词
	TopP        float64 `json:"topP"`        //取
}

func (m *ModelParam) SetDefaultParams() {
	if m.Temperature == 0 {
		m.Temperature = DefaultTemperature
	}
	if m.MaxTokens == 0 {
		m.MaxTokens = DefaultMaxTokens
	}
	if m.TopK == 0 {
		m.TopK = DefaultTopK
	}
	if m.TopP == 0 {
		m.TopP = DefaultTopP
	}
}

type SessionResponse struct {
	gorm.Model
	UserID     int    `gorm:"column:user_id"` //只能由用户创建
	Name       string `gorm:"column:name"`
	Context    []int  `gorm:"column:context;type:json"` //会话上下文
	MsgCount   int    `gorm:"column:msg_count"`         //消息数量
	TokenUsage uint   `gorm:"column:token_usage"`       //token使用数量
}

type IMParamContext struct {
	UserID     int    `json:"user_id"`     //用户id
	SessionID  int    `json:"session_id"`  //会话id
	FunctionID int    `json:"function_id"` //功能id
	ModelID    int    `json:"model_id"`    //模型id
	ModelType  string `json:"model_type"`  //模型类型
	Question   string `json:"question"`    //问题
	Channel    string `json:"channel"`     //消息队列
}