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

Refactor Neo API context handling and improve context management

Browse source 
- Replaced the custom context creation functions with a new context package, enhancing consistency and maintainability across the API.
- Updated multiple methods in the DSL to utilize the new context package, ensuring a unified approach to context management.
- Removed the obsolete context.go file, streamlining the codebase and reducing redundancy.

These changes improve the overall structure and clarity of the Neo API, paving the way for future enhancements in context handling and assistant functionalities.
This commit is contained in:
Max 2025-01-06 14:20:39 +08:00
parent ca78c40293
commit 90967f58c1
7 changed files with 109 additions and 90 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
neo/api.go
View file

@ -15,6 +15,7 @@ import (
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/process"
"github.com/yaoapp/yao/helper"
chatctx "github.com/yaoapp/yao/neo/context"
"github.com/yaoapp/yao/neo/message"
"github.com/yaoapp/yao/neo/store"
)
@ -173,7 +174,7 @@ func (neo *DSL) handleUpload(c *gin.Context) {
}
// Set the context
ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "")
ctx, cancel := chatctx.NewWithCancel(sid, c.Query("chat_id"), "")
defer cancel()
// Upload the file
@ -214,7 +215,7 @@ func (neo *DSL) handleChat(c *gin.Context) {
}
// Set the context with validated chat_id
ctx, cancel := NewContextWithCancel(sid, chatID, c.Query("context"))
ctx, cancel := chatctx.NewWithCancel(sid, chatID, c.Query("context"))
defer cancel()
neo.Answer(ctx, content, c)
@ -297,7 +298,7 @@ func (neo *DSL) handleDownload(c *gin.Context) {
}
// Set the context
ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "")
ctx, cancel := chatctx.NewWithCancel(sid, c.Query("chat_id"), "")
defer cancel()
// Download the file
@ -537,7 +538,7 @@ func (neo *DSL) handleChatUpdate(c *gin.Context) {
// If content is not empty, Generate the chat title
if body.Content != "" {
ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "")
ctx, cancel := chatctx.NewWithCancel(sid, c.Query("chat_id"), "")
defer cancel()
title, err := neo.GenerateChatTitle(ctx, body.Content, c, true)
@ -729,7 +730,7 @@ func (neo *DSL) handleGenerateTitle(c *gin.Context) {
return
}
ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "")
ctx, cancel := chatctx.NewWithCancel(resp.sid, c.Query("chat_id"), "")
defer cancel()
// Use silent mode for regular HTTP requests, streaming for SSE
@ -780,7 +781,7 @@ func (neo *DSL) handleGeneratePrompts(c *gin.Context) {
return
}
ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "")
ctx, cancel := chatctx.NewWithCancel(resp.sid, c.Query("chat_id"), "")
defer cancel()
// Use silent mode for regular HTTP requests, streaming for SSE
@ -831,7 +832,7 @@ func (neo *DSL) handleGenerateCustom(c *gin.Context) {
return
}
ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "")
ctx, cancel := chatctx.NewWithCancel(resp.sid, c.Query("chat_id"), "")
defer cancel()
// Use silent mode for regular HTTP requests, streaming for SSE

1
neo/assistant/hooks.go Normal file
View file

@ -0,0 +1 @@
package assistant

49
neo/context.go
View file

@ -1,49 +0,0 @@
package neo
import (
"context"
"time"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/kun/log"
)
// NewContext create a new context
func NewContext(sid, cid, payload string) Context {
ctx := Context{Context: context.Background(), Sid: sid, ChatID: cid}
if payload == "" {
return ctx
}
err := jsoniter.Unmarshal([]byte(payload), &ctx)
if err != nil {
log.Error("%s", err.Error())
}
return ctx
}
// NewContextWithCancel create a new context with cancel
func NewContextWithCancel(sid, cid, payload string) (Context, context.CancelFunc) {
ctx := NewContext(sid, cid, payload)
return ContextWithCancel(ctx)
}
// NewContextWithTimeout create a new context with timeout
func NewContextWithTimeout(sid, cid, payload string, timeout time.Duration) (Context, context.CancelFunc) {
ctx := NewContext(sid, cid, payload)
return ContextWithTimeout(ctx, timeout)
}
// ContextWithCancel create a new context
func ContextWithCancel(parent Context) (Context, context.CancelFunc) {
new, cancel := context.WithCancel(parent.Context)
parent.Context = new
return parent, cancel
}
// ContextWithTimeout create a new context
func ContextWithTimeout(parent Context, timeout time.Duration) (Context, context.CancelFunc) {
new, cancel := context.WithTimeout(parent.Context, timeout)
parent.Context = new
return parent, cancel
}

82
neo/context/context.go Normal file
View file

@ -0,0 +1,82 @@
package context
import (
"context"
"time"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/kun/log"
)
// Context the context
type Context struct {
context.Context
Sid string `json:"sid" yaml:"-"` // Session ID
ChatID string `json:"chat_id,omitempty"` // Chat ID, use to select chat
AssistantID string `json:"assistant_id,omitempty"` // Assistant ID, use to select assistant
Stack string `json:"stack,omitempty"`
Path string `json:"pathname,omitempty"`
FormData map[string]interface{} `json:"formdata,omitempty"`
Field *Field `json:"field,omitempty"`
Namespace string `json:"namespace,omitempty"`
Config map[string]interface{} `json:"config,omitempty"`
Signal interface{} `json:"signal,omitempty"`
Upload *FileUpload `json:"upload,omitempty"`
}
// Field the context field
type Field struct {
Name string `json:"name,omitempty"`
Type string `json:"type,omitempty"`
Bind string `json:"bind,omitempty"`
Props map[string]interface{} `json:"props,omitempty"`
Children []interface{} `json:"children,omitempty"`
}
// FileUpload the file upload
type FileUpload struct {
Name string `json:"name,omitempty"`
Type string `json:"type,omitempty"`
Size int64 `json:"size,omitempty"`
TempFile string `json:"temp_file,omitempty"`
}
// New create a new context
func New(sid, cid, payload string) Context {
ctx := Context{Context: context.Background(), Sid: sid, ChatID: cid}
if payload == "" {
return ctx
}
err := jsoniter.Unmarshal([]byte(payload), &ctx)
if err != nil {
log.Error("%s", err.Error())
}
return ctx
}
// NewWithCancel create a new context with cancel
func NewWithCancel(sid, cid, payload string) (Context, context.CancelFunc) {
ctx := New(sid, cid, payload)
return WithCancel(ctx)
}
// NewWithTimeout create a new context with timeout
func NewWithTimeout(sid, cid, payload string, timeout time.Duration) (Context, context.CancelFunc) {
ctx := New(sid, cid, payload)
return WithTimeout(ctx, timeout)
}
// WithCancel create a new context
func WithCancel(parent Context) (Context, context.CancelFunc) {
new, cancel := context.WithCancel(parent.Context)
parent.Context = new
return parent, cancel
}
// WithTimeout create a new context
func WithTimeout(parent Context, timeout time.Duration) (Context, context.CancelFunc) {
new, cancel := context.WithTimeout(parent.Context, timeout)
parent.Context = new
return parent, cancel
}

7
neo/hooks.go
View file

@ -7,10 +7,11 @@ import (
"github.com/gin-gonic/gin"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/process"
chatctx "github.com/yaoapp/yao/neo/context"
)
// HookCreate create the assistant
func (neo *DSL) HookCreate(ctx Context, messages []map[string]interface{}, c *gin.Context) (CreateResponse, error) {
func (neo *DSL) HookCreate(ctx chatctx.Context, messages []map[string]interface{}, c *gin.Context) (CreateResponse, error) {
// Default assistant
assistantID := neo.Use
@ -69,7 +70,7 @@ func (neo *DSL) HookCreate(ctx Context, messages []map[string]interface{}, c *gi
}
// HookPrepare executes the prepare hook before AI is called
func (neo *DSL) HookPrepare(ctx Context, messages []map[string]interface{}) ([]map[string]interface{}, error) {
func (neo *DSL) HookPrepare(ctx chatctx.Context, messages []map[string]interface{}) ([]map[string]interface{}, error) {
if neo.Prepare == "" {
return messages, nil
}
@ -114,7 +115,7 @@ func (neo *DSL) HookPrepare(ctx Context, messages []map[string]interface{}) ([]m
}
// HookWrite executes the write hook when response is received from AI
func (neo *DSL) HookWrite(ctx Context, messages []map[string]interface{}, response map[string]interface{}, content string, writer *gin.ResponseWriter) ([]map[string]interface{}, error) {
func (neo *DSL) HookWrite(ctx chatctx.Context, messages []map[string]interface{}, response map[string]interface{}, content string, writer *gin.ResponseWriter) ([]map[string]interface{}, error) {
if neo.Write == "" {
return []map[string]interface{}{response}, nil
}

27
neo/neo.go
View file

@ -9,6 +9,7 @@ import (
"github.com/gin-gonic/gin"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/neo/assistant"
chatctx "github.com/yaoapp/yao/neo/context"
"github.com/yaoapp/yao/neo/message"
)
@ -16,7 +17,7 @@ import (
var lock sync.Mutex = sync.Mutex{}
// Answer reply the message
func (neo *DSL) Answer(ctx Context, question string, c *gin.Context) error {
func (neo *DSL) Answer(ctx chatctx.Context, question string, c *gin.Context) error {
messages, err := neo.chatMessages(ctx, question)
if err != nil {
msg := message.New().Error(err).Done()
@ -50,7 +51,7 @@ func (neo *DSL) Select(id string) (assistant.API, error) {
}
// GeneratePrompts generate prompts for the AI assistant
func (neo *DSL) GeneratePrompts(ctx Context, input string, c *gin.Context, silent ...bool) (string, error) {
func (neo *DSL) GeneratePrompts(ctx chatctx.Context, input string, c *gin.Context, silent ...bool) (string, error) {
prompts := `
Optimize the prompts for the AI assistant
1. Optimize prompts based on the user's input
@ -68,7 +69,7 @@ func (neo *DSL) GeneratePrompts(ctx Context, input string, c *gin.Context, silen
}
// GenerateChatTitle generate the chat title
func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context, silent ...bool) (string, error) {
func (neo *DSL) GenerateChatTitle(ctx chatctx.Context, input string, c *gin.Context, silent ...bool) (string, error) {
prompts := `
Help me generate a title for the chat
1. The title should be a short and concise description of the chat.
@ -84,7 +85,7 @@ func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context, sil
}
// GenerateWithAI generate content with AI, type can be "title", "prompts", etc.
func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, systemPrompt string, c *gin.Context, silent bool) (string, error) {
func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType string, systemPrompt string, c *gin.Context, silent bool) (string, error) {
messages := []map[string]interface{}{
{"role": "system", "content": systemPrompt},
{
@ -187,7 +188,7 @@ func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, sy
}
// Upload upload a file
func (neo *DSL) Upload(ctx Context, c *gin.Context) (*assistant.File, error) {
func (neo *DSL) Upload(ctx chatctx.Context, c *gin.Context) (*assistant.File, error) {
// Get the file
tmpfile, err := c.FormFile("file")
if err != nil {
@ -212,11 +213,11 @@ func (neo *DSL) Upload(ctx Context, c *gin.Context) (*assistant.File, error) {
}
// Get file info
ctx.Upload = &FileUpload{
Bytes: int(tmpfile.Size),
Name: tmpfile.Filename,
ContentType: tmpfile.Header.Get("Content-Type"),
Option: option,
ctx.Upload = &chatctx.FileUpload{
Name: tmpfile.Filename,
Type: tmpfile.Header.Get("Content-Type"),
Size: tmpfile.Size,
TempFile: tmpfile.Filename,
}
// Default use the assistant in context
@ -235,7 +236,7 @@ func (neo *DSL) Upload(ctx Context, c *gin.Context) (*assistant.File, error) {
}
// Download downloads a file
func (neo *DSL) Download(ctx Context, c *gin.Context) (*assistant.FileResponse, error) {
func (neo *DSL) Download(ctx chatctx.Context, c *gin.Context) (*assistant.FileResponse, error) {
// Get file_id from query string
fileID := c.Query("file_id")
if fileID == "" {
@ -259,7 +260,7 @@ func (neo *DSL) Download(ctx Context, c *gin.Context) (*assistant.FileResponse,
}
// chat chat with AI
func (neo *DSL) chat(ast assistant.API, ctx Context, messages []map[string]interface{}, c *gin.Context) error {
func (neo *DSL) chat(ast assistant.API, ctx chatctx.Context, messages []map[string]interface{}, c *gin.Context) error {
if ast == nil {
msg := message.New().Error("assistant is not initialized").Done()
msg.Write(c.Writer)
@ -339,7 +340,7 @@ func (neo *DSL) chat(ast assistant.API, ctx Context, messages []map[string]inter
}
// chatMessages get the chat messages
func (neo *DSL) chatMessages(ctx Context, content ...string) ([]map[string]interface{}, error) {
func (neo *DSL) chatMessages(ctx chatctx.Context, content ...string) ([]map[string]interface{}, error) {
history, err := neo.Store.GetHistory(ctx.Sid, ctx.ChatID)
if err != nil {

18
neo/types.go
View file

@ -1,8 +1,6 @@
package neo
import (
"context"
"github.com/gin-gonic/gin"
"github.com/yaoapp/yao/neo/assistant"
"github.com/yaoapp/yao/neo/rag"
@ -48,22 +46,6 @@ type Mention struct {
Type string `json:"type,omitempty"`
}
// Context the context
type Context struct {
Sid string `json:"sid" yaml:"-"` // Session ID
ChatID string `json:"chat_id,omitempty"` // Chat ID, use to select chat
AssistantID string `json:"assistant_id,omitempty"` // Assistant ID, use to select assistant
Stack string `json:"stack,omitempty"`
Path string `json:"pathname,omitempty"`
FormData map[string]interface{} `json:"formdata,omitempty"`
Field *Field `json:"field,omitempty"`
Namespace string `json:"namespace,omitempty"`
Config map[string]interface{} `json:"config,omitempty"`
Signal interface{} `json:"signal,omitempty"`
Upload *FileUpload `json:"upload,omitempty"`
context.Context `json:"-" yaml:"-"`
}
// Field the context field
type Field struct {
Name string `json:"name,omitempty"`