git-sync/crawlab
1
0
Fork 0
mirror of https://github.com/crawlab-team/crawlab.git synced 2026-01-21 17:21:09 +01:00
Code Issues Packages Projects Releases Wiki Activity

feat: introduce LLMResponseUsage struct and update ChatMessage models

Browse Source 
- Added LLMResponseUsage struct to encapsulate token usage details for language model responses.
- Updated ChatMessageContent and ChatMessage models to include a Usage field of type *entity.LLMResponseUsage, enhancing message handling capabilities.
- Improved structure and clarity of models by ensuring consistent field definitions and descriptions.
This commit is contained in:
Marvin Zhang
2025-04-11 22:29:02 +08:00
parent 631f6f96b2
commit 2549abdc61
3 changed files with 30 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
core/entity/llm.go Normal file
View File

@@ -0,0 +1,7 @@
package entity
type LLMResponseUsage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
TotalTokens int `json:"total_tokens"`
}

22
core/models/models/chat_message.go
View File

@@ -1,23 +1,23 @@
package models
import (
"github.com/crawlab-team/crawlab/core/entity"
"go.mongodb.org/mongo-driver/bson/primitive"
)
type ChatMessage struct {
any `collection:"chat_messages"`
BaseModel `bson:",inline"`
ConversationId primitive.ObjectID `json:"conversation_id" bson:"conversation_id" description:"Conversation ID"`
Role string `json:"role" bson:"role" description:"Message role (user/assistant/system)"`
Content string `json:"content,omitempty" bson:"content,omitempty" description:"Message content for user/system only"`
IsAgent bool `json:"is_agent,omitempty" bson:"is_agent,omitempty" description:"Is agent"`
ContentIds []primitive.ObjectID `json:"content_ids,omitempty" bson:"content_ids,omitempty" description:"Content IDs"`
Contents []ChatMessageContent `json:"contents,omitempty" bson:"-" description:"Contents"`
Tokens int `json:"tokens" bson:"tokens" description:"Number of tokens in the message"`
Model string `json:"model" bson:"model" description:"AI model used"`
Metadata map[string]any `json:"metadata,omitempty" bson:"metadata,omitempty" description:"Additional metadata"`
Status string `json:"status" bson:"status" description:"Message status (pending/completed/failed)"`
Error string `json:"error,omitempty" bson:"error,omitempty" description:"Error message if failed"`
ConversationId primitive.ObjectID `json:"conversation_id" bson:"conversation_id" description:"Conversation ID"`
Role string `json:"role" bson:"role" description:"Message role (user/assistant/system)"`
Content string `json:"content,omitempty" bson:"content,omitempty" description:"Message content for user/system only"`
IsAgent bool `json:"is_agent,omitempty" bson:"is_agent,omitempty" description:"Is agent"`
ContentIds []primitive.ObjectID `json:"content_ids,omitempty" bson:"content_ids,omitempty" description:"Content IDs"`
Contents []ChatMessageContent `json:"contents,omitempty" bson:"-" description:"Contents"`
Model string `json:"model" bson:"model" description:"AI model used"`
Status string `json:"status" bson:"status" description:"Message status (pending/completed/failed)"`
Error string `json:"error,omitempty" bson:"error,omitempty" description:"Error message if failed"`
Usage *entity.LLMResponseUsage `json:"usage,omitempty" bson:"usage,omitempty" description:"Usage"`
}
func (m *ChatMessage) GetContent() string {

20
core/models/models/chat_message_content.go
View File

@@ -19,16 +19,20 @@
package models
import "go.mongodb.org/mongo-driver/bson/primitive"
import (
"github.com/crawlab-team/crawlab/core/entity"
"go.mongodb.org/mongo-driver/bson/primitive"
)
type ChatMessageContent struct {
any `collection:"chat_message_contents"`
BaseModel `bson:",inline"`
MessageId primitive.ObjectID `json:"message_id" bson:"message_id" description:"Message ID"`
Key string `json:"key" bson:"key" description:"Message content key"`
Content string `json:"content" bson:"content" description:"Message content"`
Type string `json:"type" bson:"type" description:"Message type (text/action)"`
Action string `json:"action,omitempty" bson:"action,omitempty" description:"Action name"`
ActionStatus string `json:"action_status,omitempty" bson:"action_status,omitempty" description:"Action status"`
Hidden bool `json:"hidden,omitempty" bson:"hidden,omitempty" description:"Hidden"`
MessageId primitive.ObjectID `json:"message_id" bson:"message_id" description:"Message ID"`
Key string `json:"key" bson:"key" description:"Message content key"`
Content string `json:"content" bson:"content" description:"Message content"`
Type string `json:"type" bson:"type" description:"Message type (text/action)"`
Action string `json:"action,omitempty" bson:"action,omitempty" description:"Action name"`
ActionStatus string `json:"action_status,omitempty" bson:"action_status,omitempty" description:"Action status"`
Hidden bool `json:"hidden,omitempty" bson:"hidden,omitempty" description:"Hidden"`
Usage *entity.LLMResponseUsage `json:"usage,omitempty" bson:"usage,omitempty" description:"Usage"`
}