opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Enhance Neo API assistant with function support and improved message handling

Browse source 
- Introduced a new Function type to define assistant functions, enhancing the assistant's capabilities.
- Updated the Assistant struct to include a Functions field, allowing for dynamic function management.
- Refactored the streamChat method to handle tool calls, improving interaction with functions during chat streaming.
- Enhanced message handling in the NewOpenAI function to support tool call responses, ensuring proper message parsing.
- Implemented loading and saving of functions from a JSON file, streamlining function management within the assistant.

These changes improve the robustness and maintainability of the Neo API, paving the way for more flexible and powerful assistant functionalities.
This commit is contained in:
Max 2025-01-15 11:20:55 +08:00
parent 860bfd1d0d
commit e7185aad10
8 changed files with 220 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

52
neo/assistant/api.go
View file

@ -163,11 +163,11 @@ func (next *NextAction) Execute(c *gin.Context, ctx chatctx.Context) error {
func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, messages []message.Message, options map[string]interface{}) error {
clientBreak := make(chan bool, 1)
done := make(chan bool, 1)
content := []byte{}
content := message.NewContent("text")
// Chat with AI in background
go func() {
err := ast.streamChat(c, ctx, messages, options, clientBreak, done, &content)
err := ast.streamChat(c, ctx, messages, options, clientBreak, done, content)
if err != nil {
chatMessage.New().Error(err).Done().Write(c.Writer)
}
@ -187,8 +187,14 @@ func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, mess
}
// streamChat handles the streaming chat interaction
func (ast *Assistant) streamChat(c *gin.Context, ctx chatctx.Context, messages []message.Message, options map[string]interface{},
clientBreak chan bool, done chan bool, content *[]byte) error {
func (ast *Assistant) streamChat(
c *gin.Context,
ctx chatctx.Context,
messages []message.Message,
options map[string]interface{},
clientBreak chan bool,
done chan bool,
content *message.Content) error {
return ast.Chat(c.Request.Context(), messages, options, func(data []byte) int {
select {
@ -204,7 +210,7 @@ func (ast *Assistant) streamChat(c *gin.Context, ctx chatctx.Context, messages [
// Handle error
if msg.Type == "error" {
value := msg.String()
res, hookErr := ast.HookFail(c, ctx, messages, string(*content), fmt.Errorf("%s", value))
res, hookErr := ast.HookFail(c, ctx, messages, content.String(), fmt.Errorf("%s", value))
if hookErr == nil && res != nil && (res.Output != "" || res.Error != "") {
value = res.Output
if res.Error != "" {
@ -215,12 +221,26 @@ func (ast *Assistant) streamChat(c *gin.Context, ctx chatctx.Context, messages [
return 0 // break
}
// Handle tool call
if msg.Type == "tool_calls" {
content.SetType("function") // Set type to function
// Set id
if id, ok := msg.Props["id"].(string); ok && id != "" {
content.SetID(id)
}
// Set name
if name, ok := msg.Props["name"].(string); ok && name != "" {
content.SetName(name)
}
}
// Append content and send message
*content = msg.Append(*content)
value := msg.String()
content.Append(value)
if value != "" {
// Handle stream
res, err := ast.HookStream(c, ctx, messages, string(*content))
res, err := ast.HookStream(c, ctx, messages, content.String(), msg.Type == "tool_calls")
if err == nil && res != nil {
if res.Output != "" {
value = res.Output
@ -256,7 +276,8 @@ func (ast *Assistant) streamChat(c *gin.Context, ctx chatctx.Context, messages [
// }
// Call HookDone
res, hookErr := ast.HookDone(c, ctx, messages, string(*content))
content.SetStatus(message.ContentStatusDone)
res, hookErr := ast.HookDone(c, ctx, messages, content.String(), msg.Type == "tool_calls")
if hookErr == nil && res != nil {
if res.Output != "" {
chatMessage.New().
@ -295,13 +316,13 @@ func (ast *Assistant) streamChat(c *gin.Context, ctx chatctx.Context, messages [
}
// saveChatHistory saves the chat history if storage is available
func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []message.Message, content []byte) {
if len(content) > 0 && ctx.Sid != "" && len(messages) > 0 {
func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []message.Message, content *message.Content) {
if len(content.Bytes) > 0 && ctx.Sid != "" && len(messages) > 0 {
storage.SaveHistory(
ctx.Sid,
[]map[string]interface{}{
{"role": "user", "content": messages[len(messages)-1].Content(), "name": ctx.Sid},
{"role": "assistant", "content": string(content), "name": ctx.Sid},
{"role": "assistant", "content": content.String(), "name": ctx.Sid},
},
ctx.ChatID,
nil,
@ -319,6 +340,15 @@ func (ast *Assistant) withOptions(options map[string]interface{}) map[string]int
options[key] = value
}
}
// Add functions
if ast.Functions != nil {
options["tools"] = ast.Functions
if options["tool_choice"] == nil {
options["tool_choice"] = "auto"
}
}
return options
}

1
neo/assistant/assistant.go
View file

@ -132,6 +132,7 @@ func (ast *Assistant) Map() map[string]interface{} {
"description": ast.Description,
"options": ast.Options,
"prompts": ast.Prompts,
"functions": ast.Functions,
"tags": ast.Tags,
"mentionable": ast.Mentionable,
"automated": ast.Automated,

8
neo/assistant/hooks.go
View file

@ -57,13 +57,13 @@ func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []
}
// HookStream Handle streaming response from LLM
func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input []message.Message, output string) (*ResHookStream, error) {
func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input []message.Message, output string, toolcall bool) (*ResHookStream, error) {
// Create timeout context
ctx, cancel := ast.createTimeoutContext(c)
defer cancel()
v, err := ast.call(ctx, "Stream", context, input, output, c.Writer)
v, err := ast.call(ctx, "Stream", context, input, output, toolcall, c.Writer)
if err != nil {
if err.Error() == HookErrorMethodNotFound {
return nil, nil
@ -100,12 +100,12 @@ func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input
}
// HookDone Handle completion of assistant response
func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []message.Message, output string) (*ResHookDone, error) {
func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []message.Message, output string, toolcall bool) (*ResHookDone, error) {
// Create timeout context
ctx, cancel := ast.createTimeoutContext(c)
defer cancel()
v, err := ast.call(ctx, "Done", context, input, output, c.Writer)
v, err := ast.call(ctx, "Done", context, input, output, toolcall, c.Writer)
if err != nil {
if err.Error() == HookErrorMethodNotFound {
return nil, nil

55
neo/assistant/load.go
View file

@ -246,6 +246,16 @@ func LoadPath(path string) (*Assistant, error) {
}
// load functions
functionsfile := filepath.Join(path, "functions.json")
if has, _ := app.Exists(functionsfile); has {
functions, ts, err := loadFunctions(functionsfile)
if err != nil {
return nil, err
}
data["functions"] = functions
updatedAt = max(updatedAt, ts)
data["updated_at"] = updatedAt
}
// load flow
@ -340,6 +350,25 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
assistant.Prompts = prompts
}
// functions
if funcs, has := data["functions"]; has {
switch vv := funcs.(type) {
case []Function:
assistant.Functions = vv
default:
raw, err := jsoniter.Marshal(vv)
if err != nil {
return nil, err
}
var functions []Function
err = jsoniter.Unmarshal(raw, &functions)
if err != nil {
return nil, err
}
assistant.Functions = functions
}
}
// script
if data["script"] != nil {
switch v := data["script"].(type) {
@ -382,6 +411,32 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
return assistant, nil
}
func loadFunctions(file string) ([]Function, int64, error) {
app, err := fs.Get("app")
if err != nil {
return nil, 0, err
}
ts, err := app.ModTime(file)
if err != nil {
return nil, 0, err
}
raw, err := app.ReadFile(file)
if err != nil {
return nil, 0, err
}
var functions []Function
err = jsoniter.Unmarshal(raw, &functions)
if err != nil {
return nil, 0, err
}
return functions, ts.UnixNano(), nil
}
func loadPrompts(file string, root string) (string, int64, error) {
app, err := fs.Get("app")

11
neo/assistant/types.go
View file

@ -85,6 +85,16 @@ type Prompt struct {
Name string `json:"name,omitempty"`
}
// Function a function
type Function struct {
Type string `json:"type"`
Function struct {
Name string `json:"name"`
Description string `json:"description"`
Parameters map[string]interface{} `json:"parameters"`
} `json:"function"`
}
// QueryParam the assistant query param
type QueryParam struct {
Limit uint `json:"limit"`
@ -110,6 +120,7 @@ type Assistant struct {
Automated bool `json:"automated,omitempty"` // Whether this assistant is automated
Options map[string]interface{} `json:"options,omitempty"` // AI Options
Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts
Functions []Function `json:"functions,omitempty"` // Assistant Functions
Flows []map[string]interface{} `json:"flows,omitempty"` // Assistant Flows
Script *v8.Script `json:"-" yaml:"-"` // Assistant Script
CreatedAt int64 `json:"created_at"` // Creation timestamp

67
neo/message/content.go Normal file
View file

@ -0,0 +1,67 @@
package message
import "fmt"
const (
// ContentStatusPending the content status pending
ContentStatusPending = iota
// ContentStatusDone the content status done
ContentStatusDone
// ContentStatusError the content status error
ContentStatusError
)
// Content the content
type Content struct {
ID string `json:"id"`
Name string `json:"name"`
Bytes []byte `json:"bytes"`
Type string `json:"type"` // text, function, error
Status uint8 `json:"status"` // 0: pending, 1: done
}
// NewContent create a new content
func NewContent(typ string) *Content {
if typ == "" {
typ = "text"
}
return &Content{
Bytes: []byte{},
Type: typ,
Status: ContentStatusPending,
}
}
// String the content string
func (c *Content) String() string {
if c.Type == "function" {
return fmt.Sprintf(`{"id":"%s","type": "function", "function": {"name": "%s", "arguments": "%s"}}`, c.ID, c.Name, c.Bytes)
}
return string(c.Bytes)
}
// SetID set the content id
func (c *Content) SetID(id string) {
c.ID = id
}
// SetName set the content name
func (c *Content) SetName(name string) {
c.Name = name
}
// SetType set the content type
func (c *Content) SetType(typ string) {
c.Type = typ
}
// Append append the content
func (c *Content) Append(data string) {
c.Bytes = append(c.Bytes, []byte(data)...)
}
// SetStatus set the content status
func (c *Content) SetStatus(status uint8) {
c.Status = status
}

20
neo/message/message.go
View file

@ -49,7 +49,7 @@ type Action struct {
// New create a new message
func New() *Message {
return &Message{Actions: []Action{}}
return &Message{Actions: []Action{}, Props: map[string]interface{}{}}
}
// NewString create a new message from string
@ -75,6 +75,21 @@ func NewOpenAI(data []byte) *Message {
data = []byte(strings.TrimPrefix(text, "data: "))
switch {
case strings.Contains(text, `"delta":{`) && strings.Contains(text, `"tool_calls"`):
var toolCalls openai.ToolCalls
if err := jsoniter.Unmarshal(data, &toolCalls); err != nil {
msg.Text = err.Error() + "\n" + string(data)
return msg
}
msg.Type = "tool_calls"
if len(toolCalls.Choices) > 0 && len(toolCalls.Choices[0].Delta.ToolCalls) > 0 {
msg.Props["id"] = toolCalls.Choices[0].Delta.ToolCalls[0].ID
msg.Props["name"] = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Name
msg.Text = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Arguments
}
case strings.Contains(text, `"delta":{`) && strings.Contains(text, `"content":`):
var message openai.Message
if err := jsoniter.Unmarshal(data, &message); err != nil {
@ -92,6 +107,9 @@ func NewOpenAI(data []byte) *Message {
case strings.Contains(text, `"finish_reason":"stop"`):
msg.IsDone = true
case strings.Contains(text, `"finish_reason":"tool_calls"`):
msg.IsDone = true
default:
str := strings.TrimPrefix(strings.Trim(string(data), "\""), "data: ")
msg.Type = "error"

22
openai/types.go
View file

@ -16,6 +16,28 @@ type Message struct {
} `json:"choices,omitempty"`
}
// ToolCalls is the response from OpenAI
type ToolCalls struct {
ID string `json:"id,omitempty"`
Object string `json:"object,omitempty"`
Created int64 `json:"created,omitempty"`
Model string `json:"model,omitempty"`
Choices []struct {
Delta struct {
ToolCalls []struct {
ID string `json:"id,omitempty"`
Type string `json:"type,omitempty"`
Function struct {
Name string `json:"name,omitempty"`
Arguments string `json:"arguments,omitempty"`
} `json:"function,omitempty"`
} `json:"tool_calls,omitempty"`
} `json:"delta,omitempty"`
Index int `json:"index,omitempty"`
FinishReason string `json:"finish_reason,omitempty"`
} `json:"choices,omitempty"`
}
// ErrorMessage is the error response from OpenAI
type ErrorMessage struct {
Error Error `json:"error,omitempty"`