From 239a22899b8835472874b5f285888cfc29ece930 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 3 Jul 2026 20:46:04 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E6=A8=A1=E5=9E=8B=E6=80=9D=E8=80=83?= =?UTF-8?q?=E6=A8=A1=E5=BC=8F=E9=85=8D=E7=BD=AE=20+=20Unicode=20=E6=88=AA?= =?UTF-8?q?=E6=96=AD=20+=20=E5=AE=A1=E8=AE=A1=E4=BF=AE=E5=A4=8D=20(13=20fi?= =?UTF-8?q?les)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 模型模式: - 新增 LLMConfig/Source.ThinkingEnabled 配置，通过 ExtraBody 控制 DeepSeek thinking mode，默认关闭 - SeedDefaults/ToConfig 读写 core.llm.thinking_enabled - deepseek.lua 移除硬编码 temperature=0 Unicode 截断: - truncateStr 改按 rune 计数，修复中文截断乱码 审计修复 (Critical): - graph.go: defer rows.Close 在 for 循环 → 显式 Close (连接池泄漏) - cli/openclaw/plugin.go: bare type assertion → comma-ok (panic) - channel.go: payload["type"].(string) → comma-ok (panic) - webui/handler.go: .(string) → fmt.Sprint (panic) - agent.go: 添加 nil provider 错误返回 审计修复 (High): - events/bus.go: copy handler slice under RLock (data race) - webui/handler.go: SSE 通过 channel 串行化写入 (data race) - timer/plugin.go: time.Sleep → select with stopCh (Stop 阻塞) - provider.go: stream ch <- 添加 select ctx.Done (goroutine 泄漏) - main.go: outputCh goroutine 添加 ctx.Done 退出路径 --- cmd/homed/main.go | 63 +++++++++++++++++------------ internal/agent/api/provider.go | 20 +++++++-- internal/agent/core/agent.go | 41 +++++++++++++++---- internal/agent/io/channel.go | 3 +- internal/config/registry.go | 34 ++++++++++------ internal/events/bus.go | 6 ++- internal/lua/adapters/deepseek.lua | 3 +- internal/memory/graph.go | 13 +++--- internal/plugins/cli/plugin.go | 6 ++- internal/plugins/openclaw/plugin.go | 6 ++- internal/plugins/timer/plugin.go | 20 +++++---- internal/plugins/webui/handler.go | 26 +++++++++--- pkg/types/types.go | 30 +++++++------- 13 files changed, 182 insertions(+), 89 deletions(-) diff --git a/cmd/homed/main.go b/cmd/homed/main.go index 64b34d8..fdb5ee4 100644 --- a/cmd/homed/main.go +++ b/cmd/homed/main.go @@ -1,6 +1,7 @@ package main import ( + "context" "flag" "log" "os" @@ -162,34 +163,45 @@ func main() { log.Printf("[homed] text memory active at %s", filepath.Join(cfg.Daemon.DataDir, "memory", "text")) } + ctx, stop := context.WithCancel(context.Background()) + defer stop() + go func() { - for evt := range iom.OutputChan() { - if evt.Target == "memory" && evt.Type == "memory_candidate" { - source, _ := evt.Payload["source"].(string) - input, _ := evt.Payload["input"].(string) - response, _ := evt.Payload["response"].(string) - toolsUsed, _ := evt.Payload["tools_used"].([]string) - agentID, _ := evt.Payload["agent_id"].(string) - - if input != "" && textMem != nil { - te := text.Event{ - Timestamp: time.Now().Unix(), - Source: source, - Input: input, - Response: response, - ToolsUsed: toolsUsed, - AgentID: agentID, - } - if err := textMem.Append(te); err != nil { - log.Printf("[homed] text memory append: %v", err) - } + for { + select { + case <-ctx.Done(): + return + case evt, ok := <-iom.OutputChan(): + if !ok { + return } + if evt.Target == "memory" && evt.Type == "memory_candidate" { + source, _ := evt.Payload["source"].(string) + input, _ := evt.Payload["input"].(string) + response, _ := evt.Payload["response"].(string) + toolsUsed, _ := evt.Payload["tools_used"].([]string) + agentID, _ := evt.Payload["agent_id"].(string) - if input != "" { - distiller.Append("agent", "user", input) - } - if response != "" { - distiller.Append("agent", "assistant", response) + if input != "" && textMem != nil { + te := text.Event{ + Timestamp: time.Now().Unix(), + Source: source, + Input: input, + Response: response, + ToolsUsed: toolsUsed, + AgentID: agentID, + } + if err := textMem.Append(te); err != nil { + log.Printf("[homed] text memory append: %v", err) + } + } + + if input != "" { + distiller.Append("agent", "user", input) + } + if response != "" { + distiller.Append("agent", "assistant", response) + } } } } @@ -345,6 +357,7 @@ func main() { ContextSavePath: filepath.Join(cfg.Daemon.DataDir, "memory", "context.json"), StageHost: stageHost, EventBus: evBus, + ThinkingEnabled: cfg.LLM.ThinkingEnabled, }) agent.Start() defer agent.Stop() diff --git a/internal/agent/api/provider.go b/internal/agent/api/provider.go index 4d764b3..b410561 100644 --- a/internal/agent/api/provider.go +++ b/internal/agent/api/provider.go @@ -265,13 +265,17 @@ func (p *OpenAIProvider) ChatStream(ctx context.Context, req *CompletionRequest) } if err := decoder.Decode(&line); err != nil { - break + return } if len(line.Choices) > 0 { - ch <- StreamChunk{ + select { + case ch <- StreamChunk{ Content: line.Choices[0].Delta.Content, Done: line.Choices[0].FinishReason != nil, + }: + case <-ctx.Done(): + return } } } @@ -498,9 +502,13 @@ func (p *LuaAdaptedProvider) ChatStream(ctx context.Context, req *CompletionRequ continue } if len(raw.Choices) > 0 { - ch <- StreamChunk{ + select { + case ch <- StreamChunk{ Content: raw.Choices[0].Delta.Content, Done: raw.Choices[0].FinishReason != nil, + }: + case <-ctx.Done(): + return } } continue @@ -509,7 +517,11 @@ func (p *LuaAdaptedProvider) ChatStream(ctx context.Context, req *CompletionRequ // Lua 返回了变换后的统一格式 var chunk StreamChunk if err := json.Unmarshal([]byte(unified), &chunk); err == nil { - ch <- chunk + select { + case ch <- chunk: + case <-ctx.Done(): + return + } } } }() diff --git a/internal/agent/core/agent.go b/internal/agent/core/agent.go index 5eb546a..cf1f32b 100644 --- a/internal/agent/core/agent.go +++ b/internal/agent/core/agent.go @@ -7,6 +7,7 @@ import ( "strings" "sync" "time" + "unicode/utf8" agentAPI "gitcode.com/JianFeeeee/HomeAgent/internal/agent/api" agentPkg "gitcode.com/JianFeeeee/HomeAgent/internal/agent" @@ -89,6 +90,9 @@ type Agent struct { // 进行中的 LLM 请求取消函数，interceptLoop 可调用以在请求中打断 cancelLLM context.CancelFunc llmMu sync.Mutex + + // 模型思考模式（thinking/reasoning） + thinkingEnabled bool } type AgentConfig struct { @@ -115,6 +119,7 @@ type AgentConfig struct { ContextSavePath string // 上下文持久化路径，空则不持久化 StageHost *StageHost EventBus *events.Bus + ThinkingEnabled bool } func New(cfg AgentConfig) *Agent { @@ -156,6 +161,7 @@ func New(cfg AgentConfig) *Agent { selfInputCh: make(chan string, 64), childResults: make(map[string]string), interceptCh: make(chan string, 64), + thinkingEnabled: cfg.ThinkingEnabled, } } @@ -383,6 +389,10 @@ func (a *Agent) process(input string, stageCtx *sdk.StageContext) (response stri a.mu.Lock() defer a.mu.Unlock() + if a.provider == nil { + return "", nil, fmt.Errorf("agent: no LLM provider configured") + } + memContext := a.buildMemoryContext(input) sysPrompt := a.buildSystemPrompt(memContext, input) tools := a.buildToolDefs() @@ -418,14 +428,16 @@ func (a *Agent) process(input string, stageCtx *sdk.StageContext) (response stri log.Printf("[agent] interrupt injected before LLM call (turn %d)", turn) } + eb := map[string]interface{}{} + if !a.thinkingEnabled { + eb["thinking"] = map[string]interface{}{"type": "disabled"} + } req := &agentAPI.CompletionRequest{ Messages: msgs, MaxTokens: 4096, Tools: tools, ToolChoice: "auto", - ExtraBody: map[string]interface{}{ - "thinking": map[string]interface{}{"type": "disabled"}, - }, + ExtraBody: eb, } // 可取消的 LLM 调用：interceptLoop 通过 cancelLLM 打断进行中的请求 @@ -1839,6 +1851,10 @@ func (a *Agent) executeSpawnChild(tc agentAPI.ToolCall) string { // runChildTask 后台运行子 Agent 任务，完成后将结果存储并通过 selfInputCh 通知主 Agent func (a *Agent) runChildTask(taskID, task string) { + if a.provider == nil { + log.Printf("[child] %s failed: no LLM provider configured", taskID) + return + } log.Printf("[child] %s started: %s", taskID, truncateStr(task, 80)) sysPrompt := fmt.Sprintf(`你是 HomeAgent 的子任务助手。 @@ -1871,14 +1887,16 @@ func (a *Agent) runChildTask(taskID, task string) { var finalResult string for turn := 0; turn < 5; turn++ { + eb := map[string]interface{}{} + if !a.thinkingEnabled { + eb["thinking"] = map[string]interface{}{"type": "disabled"} + } req := &agentAPI.CompletionRequest{ Messages: msgs, MaxTokens: 4096, Tools: childTools, ToolChoice: "auto", - ExtraBody: map[string]interface{}{ - "thinking": map[string]interface{}{"type": "disabled"}, - }, + ExtraBody: eb, } resp, err := a.provider.Chat(a.ctx, req) @@ -2038,10 +2056,17 @@ func getFloat(m map[string]interface{}, key string) float64 { } func truncateStr(s string, max int) string { - if len(s) <= max { + if utf8.RuneCountInString(s) <= max { return s } - return s[:max] + "..." + var truncated int + for i := range s { + if truncated >= max { + return s[:i] + "..." + } + truncated++ + } + return s } // drainInterrupt 非阻塞读取 interceptCh 中的一条打断消息。 diff --git a/internal/agent/io/channel.go b/internal/agent/io/channel.go index b31ee0d..7afcd2f 100644 --- a/internal/agent/io/channel.go +++ b/internal/agent/io/channel.go @@ -246,10 +246,11 @@ func (m *IOManager) InjectInterrupt(source, channel string, payload map[string]i if payload == nil { payload = map[string]interface{}{} } + evtType, _ := payload["type"].(string) m.interruptCh <- &InputEvent{ RequestID: m.nextRequestID(), Source: source, - Type: payload["type"].(string), + Type: evtType, Payload: payload, OutputChannel: channel, } diff --git a/internal/config/registry.go b/internal/config/registry.go index 178db4b..1926523 100644 --- a/internal/config/registry.go +++ b/internal/config/registry.go @@ -186,25 +186,29 @@ func (r *ConfigRegistry) SeedDefaults(dataDir string) { set("core.llm.provider", "deepseek") set("core.llm.model", "deepseek-v4-flash") set("core.llm.base_url", "https://api.deepseek.com") + set("core.llm.api_key", "") set("core.llm.adapter", "deepseek") set("core.llm.temperature", "0.7") set("core.llm.max_tokens", "4096") + set("core.llm.thinking_enabled", "false") // llm sources sources := map[string]map[string]string{ - "deepseek": {"base_url": "https://api.deepseek.com", "model": "deepseek-v4-flash", "adapter": "deepseek", "adapter_path": "adapters/deepseek.lua"}, - "openai": {"base_url": "https://api.openai.com/v1", "model": "gpt-4o", "adapter": "openai", "adapter_path": "adapters/openai.lua"}, - "anthropic": {"base_url": "https://api.anthropic.com", "model": "claude-sonnet-4-20250514", "adapter": "anthropic", "adapter_path": "adapters/anthropic.lua"}, - "gemini": {"base_url": "https://generativelanguage.googleapis.com", "model": "gemini-2.0-flash", "adapter": "gemini", "adapter_path": "adapters/gemini.lua"}, - "mistral": {"base_url": "https://api.mistral.ai", "model": "mistral-large-latest", "adapter": "mistral", "adapter_path": "adapters/mistral.lua"}, - "groq": {"base_url": "https://api.groq.com", "model": "llama3-70b-8192", "adapter": "groq", "adapter_path": "adapters/groq.lua"}, - "github": {"base_url": "https://models.inference.ai.azure.com", "model": "gpt-4o", "adapter": "github", "adapter_path": "adapters/github.lua"}, - "ollama": {"base_url": "http://localhost:11434", "model": "llama3", "adapter": "ollama", "adapter_path": "adapters/ollama.lua"}, + "deepseek": {"base_url": "https://api.deepseek.com", "model": "deepseek-v4-flash", "api_key": "", "thinking_enabled": "false", "adapter": "deepseek", "adapter_path": "adapters/deepseek.lua"}, + "openai": {"base_url": "https://api.openai.com/v1", "model": "gpt-4o", "api_key": "", "thinking_enabled": "false", "adapter": "openai", "adapter_path": "adapters/openai.lua"}, + "anthropic": {"base_url": "https://api.anthropic.com", "model": "claude-sonnet-4-20250514", "api_key": "", "thinking_enabled": "false", "adapter": "anthropic", "adapter_path": "adapters/anthropic.lua"}, + "gemini": {"base_url": "https://generativelanguage.googleapis.com", "model": "gemini-2.0-flash", "api_key": "", "thinking_enabled": "false", "adapter": "gemini", "adapter_path": "adapters/gemini.lua"}, + "mistral": {"base_url": "https://api.mistral.ai", "model": "mistral-large-latest", "api_key": "", "thinking_enabled": "false", "adapter": "mistral", "adapter_path": "adapters/mistral.lua"}, + "groq": {"base_url": "https://api.groq.com", "model": "llama3-70b-8192", "api_key": "", "thinking_enabled": "false", "adapter": "groq", "adapter_path": "adapters/groq.lua"}, + "github": {"base_url": "https://models.inference.ai.azure.com", "model": "gpt-4o", "api_key": "", "thinking_enabled": "false", "adapter": "github", "adapter_path": "adapters/github.lua"}, + "ollama": {"base_url": "http://localhost:11434", "model": "llama3", "api_key": "", "thinking_enabled": "false", "adapter": "ollama", "adapter_path": "adapters/ollama.lua"}, } for name, props := range sources { p := "core.llm.sources." + name set(p+".base_url", props["base_url"]) set(p+".model", props["model"]) + set(p+".api_key", props["api_key"]) + set(p+".thinking_enabled", props["thinking_enabled"]) set(p+".adapter", props["adapter"]) set(p+".adapter_path", props["adapter_path"]) } @@ -345,9 +349,11 @@ func (r *ConfigRegistry) ToConfig() *types.Config { cfg.LLM.Provider = read("core.llm.provider", cfg.LLM.Provider) cfg.LLM.Model = read("core.llm.model", cfg.LLM.Model) cfg.LLM.BaseURL = read("core.llm.base_url", cfg.LLM.BaseURL) + cfg.LLM.APIKey = read("core.llm.api_key", cfg.LLM.APIKey) cfg.LLM.Adapter = read("core.llm.adapter", cfg.LLM.Adapter) cfg.LLM.Temperature = float64(readInt("core.llm.temperature", int(cfg.LLM.Temperature*100))) / 100 cfg.LLM.MaxTokens = readInt("core.llm.max_tokens", cfg.LLM.MaxTokens) + cfg.LLM.ThinkingEnabled = readBool("core.llm.thinking_enabled", cfg.LLM.ThinkingEnabled) // 重建 sources —— 从 DB 中按前缀扫描，按名称排序保证确定性 sourceNames := make([]string, 0) @@ -362,11 +368,13 @@ func (r *ConfigRegistry) ToConfig() *types.Config { for _, name := range sourceNames { p := "core.llm.sources." + name cfg.LLM.Sources = append(cfg.LLM.Sources, types.LLMSource{ - Name: name, - BaseURL: read(p+".base_url", ""), - Model: read(p+".model", ""), - Adapter: read(p+".adapter", ""), - AdapterPath: read(p+".adapter_path", ""), + Name: name, + BaseURL: read(p+".base_url", ""), + Model: read(p+".model", ""), + APIKey: read(p+".api_key", ""), + Adapter: read(p+".adapter", ""), + AdapterPath: read(p+".adapter_path", ""), + ThinkingEnabled: readBool(p+".thinking_enabled", false), }) } diff --git a/internal/events/bus.go b/internal/events/bus.go index 635cf5e..6b87b4b 100644 --- a/internal/events/bus.go +++ b/internal/events/bus.go @@ -38,8 +38,10 @@ func NewBus() *Bus { func (b *Bus) Publish(evt *Event) { b.mu.RLock() - allHandlers := b.subs[EventAll] - typeHandlers := b.subs[evt.Type] + allHandlers := make([]Handler, len(b.subs[EventAll])) + copy(allHandlers, b.subs[EventAll]) + typeHandlers := make([]Handler, len(b.subs[evt.Type])) + copy(typeHandlers, b.subs[evt.Type]) b.mu.RUnlock() for _, h := range allHandlers { diff --git a/internal/lua/adapters/deepseek.lua b/internal/lua/adapters/deepseek.lua index 0675dd1..81476af 100644 --- a/internal/lua/adapters/deepseek.lua +++ b/internal/lua/adapters/deepseek.lua @@ -5,12 +5,11 @@ adapter.version = "2.0.0" adapter.endpoint = "/chat/completions" adapter.headers = {} --- DeepSeek 格式与 OpenAI 兼容，只需要强制 temperature=0（禁用 thinking） +-- DeepSeek 格式与 OpenAI 兼容，thinking 模式由 Go 端 ExtraBody 控制 function adapter.transform_request(raw_body) local ok, req = pcall(json.decode, raw_body) if not ok then return raw_body end req.model = req.model or "deepseek-chat" - req.temperature = 0.0 req.stream = req.stream or false return json.encode(req) end diff --git a/internal/memory/graph.go b/internal/memory/graph.go index 6ede294..2b0b4f0 100644 --- a/internal/memory/graph.go +++ b/internal/memory/graph.go @@ -272,11 +272,11 @@ func (g *GraphDB) Recall(keywords []string, seedEntities []string, depth int, se if err != nil { return nil, err } - defer rows.Close() for rows.Next() { var e Entity if err := rows.Scan(&e.ID, &e.Name, &e.Type, &e.MentionCount, &e.CreatedAt, &e.UpdatedAt); err != nil { + rows.Close() return nil, err } if !entityIDs[e.ID] { @@ -284,6 +284,7 @@ func (g *GraphDB) Recall(keywords []string, seedEntities []string, depth int, se result.Entities = append(result.Entities, e) } } + rows.Close() } for _, se := range seedEntities { @@ -336,7 +337,6 @@ func (g *GraphDB) Recall(keywords []string, seedEntities []string, depth int, se if err != nil { return nil, err } - defer relRows.Close() newIDs := make(map[int64]bool) for relRows.Next() { @@ -345,6 +345,7 @@ func (g *GraphDB) Recall(keywords []string, seedEntities []string, depth int, se &rel.SourceName, &rel.TargetName, &rel.RelationType, &rel.Confidence, &rel.Status, &rel.SessionID, &rel.TurnID, &rel.CreatedAt, &rel.DateBucket); err != nil { + relRows.Close() return nil, err } result.Relations = append(result.Relations, rel) @@ -356,6 +357,7 @@ func (g *GraphDB) Recall(keywords []string, seedEntities []string, depth int, se newIDs[rel.TargetID] = true } } + relRows.Close() if len(newIDs) == 0 { break @@ -375,11 +377,11 @@ func (g *GraphDB) Recall(keywords []string, seedEntities []string, depth int, se if err != nil { return nil, err } - defer eRows.Close() for eRows.Next() { var e Entity if err := eRows.Scan(&e.ID, &e.Name, &e.Type, &e.MentionCount, &e.CreatedAt, &e.UpdatedAt); err != nil { + eRows.Close() return nil, err } if !entityIDs[e.ID] { @@ -387,6 +389,7 @@ func (g *GraphDB) Recall(keywords []string, seedEntities []string, depth int, se result.Entities = append(result.Entities, e) } } + eRows.Close() for id := range newIDs { entityIDs[id] = true @@ -408,13 +411,13 @@ func (g *GraphDB) Purge(criteria map[string]string, mode string) (int, error) { if err != nil { return 0, err } - defer rows.Close() var ids []interface{} for rows.Next() { var id int64 rows.Scan(&id) ids = append(ids, id) } + rows.Close() if len(ids) > 0 { conds = append(conds, fmt.Sprintf("source_id IN (%s)", placeholders(len(ids)))) args = append(args, ids...) @@ -426,13 +429,13 @@ func (g *GraphDB) Purge(criteria map[string]string, mode string) (int, error) { if err != nil { return 0, err } - defer rows.Close() var ids []interface{} for rows.Next() { var id int64 rows.Scan(&id) ids = append(ids, id) } + rows.Close() if len(ids) > 0 { conds = append(conds, fmt.Sprintf("target_id IN (%s)", placeholders(len(ids)))) args = append(args, ids...) diff --git a/internal/plugins/cli/plugin.go b/internal/plugins/cli/plugin.go index 3fcad85..09f7472 100644 --- a/internal/plugins/cli/plugin.go +++ b/internal/plugins/cli/plugin.go @@ -22,7 +22,11 @@ func init() { plugin.RegisterFactory("cli", func(name string, config map[string]interface{}) (sdk.Plugin, error) { sock := DefaultSocket if sock == "" { - sock = filepath.Join(config["data_dir"].(string), "cli.sock") + dataDir, ok := config["data_dir"].(string) + if !ok { + return nil, fmt.Errorf("cli plugin: config missing 'data_dir' or not a string") + } + sock = filepath.Join(dataDir, "cli.sock") } return New(name, sock), nil }) diff --git a/internal/plugins/openclaw/plugin.go b/internal/plugins/openclaw/plugin.go index a33d27a..7922db8 100644 --- a/internal/plugins/openclaw/plugin.go +++ b/internal/plugins/openclaw/plugin.go @@ -17,7 +17,11 @@ func init() { plugin.RegisterFactory("openclaw", func(name string, config map[string]interface{}) (sdk.Plugin, error) { dir := SkillsDir if dir == "" { - dir = filepath.Join(config["data_dir"].(string), "skills") + dataDir, ok := config["data_dir"].(string) + if !ok { + return nil, fmt.Errorf("openclaw plugin: config missing 'data_dir' or not a string") + } + dir = filepath.Join(dataDir, "skills") } return New(name, dir), nil }) diff --git a/internal/plugins/timer/plugin.go b/internal/plugins/timer/plugin.go index 92598b9..0b56faf 100644 --- a/internal/plugins/timer/plugin.go +++ b/internal/plugins/timer/plugin.go @@ -17,9 +17,10 @@ func init() { } type Plugin struct { - name string - mu sync.Mutex - wg sync.WaitGroup + name string + mu sync.Mutex + wg sync.WaitGroup + stopCh chan struct{} } type timerTask struct { @@ -31,7 +32,7 @@ type timerTask struct { } func New(name string) *Plugin { - return &Plugin{name: name} + return &Plugin{name: name, stopCh: make(chan struct{})} } func (p *Plugin) Name() string { return p.name } @@ -75,9 +76,13 @@ func (p *Plugin) Start(s *sdk.PluginSDK) error { go func() { defer p.wg.Done() - time.Sleep(dur) - log.Printf("[timer] firing: %s (%s later)", message, dur) - s.InjectInterruptText("timer", "timer", fmt.Sprintf("timer: %s", message)) + select { + case <-time.After(dur): + log.Printf("[timer] firing: %s (%s later)", message, dur) + s.InjectInterruptText("timer", "timer", fmt.Sprintf("timer: %s", message)) + case <-p.stopCh: + log.Printf("[timer] cancelled: %s", message) + } }() doneAt := time.Now().Add(dur) @@ -93,6 +98,7 @@ func (p *Plugin) Start(s *sdk.PluginSDK) error { } func (p *Plugin) Stop() error { + close(p.stopCh) p.wg.Wait() return nil } diff --git a/internal/plugins/webui/handler.go b/internal/plugins/webui/handler.go index 6822b3d..0d8da77 100644 --- a/internal/plugins/webui/handler.go +++ b/internal/plugins/webui/handler.go @@ -558,10 +558,22 @@ func (h *Handler) handleChatEvents(w http.ResponseWriter, r *http.Request) { ticker := time.NewTicker(15 * time.Second) defer ticker.Stop() + writeCh := make(chan string, 64) + defer close(writeCh) + + go func() { + for line := range writeCh { + fmt.Fprintf(w, "%s\n", line) + flusher.Flush() + } + }() + unsub := h.eventBus.Subscribe(events.EventAll, func(evt *events.Event) { data, _ := json.Marshal(evt) - fmt.Fprintf(w, "event: %s\ndata: %s\n\n", evt.Type, data) - flusher.Flush() + select { + case writeCh <- fmt.Sprintf("event: %s\ndata: %s", evt.Type, string(data)): + default: + } }) defer unsub() @@ -570,8 +582,10 @@ func (h *Handler) handleChatEvents(w http.ResponseWriter, r *http.Request) { case <-done: return case <-ticker.C: - fmt.Fprintf(w, ": heartbeat\n\n") - flusher.Flush() + select { + case writeCh <- ": heartbeat": + default: + } } } } @@ -693,8 +707,8 @@ func (h *Handler) handleOpenAICompletions(w http.ResponseWriter, r *http.Request }, "usage": map[string]interface{}{ "prompt_tokens": len(lastMsg.Content) / 2, - "completion_tokens": len(response.Payload["content"].(string)) / 2, - "total_tokens": (len(lastMsg.Content) + len(response.Payload["content"].(string))) / 2, + "completion_tokens": len(fmt.Sprint(response.Payload["content"])) / 2, + "total_tokens": (len(lastMsg.Content) + len(fmt.Sprint(response.Payload["content"]))) / 2, }, } diff --git a/pkg/types/types.go b/pkg/types/types.go index fc751e5..77ce133 100644 --- a/pkg/types/types.go +++ b/pkg/types/types.go @@ -94,23 +94,25 @@ type OperationLog struct { } type LLMSource struct { - Name string `json:"name"` - BaseURL string `json:"base_url"` - Model string `json:"model"` - APIKey string `json:"api_key,omitempty"` - Adapter string `json:"adapter"` - AdapterPath string `json:"adapter_path,omitempty"` + Name string `json:"name"` + BaseURL string `json:"base_url"` + Model string `json:"model"` + APIKey string `json:"api_key,omitempty"` + Adapter string `json:"adapter"` + AdapterPath string `json:"adapter_path,omitempty"` + ThinkingEnabled bool `json:"thinking_enabled,omitempty"` } type LLMConfig struct { - Provider string `json:"provider"` - Model string `json:"model"` - BaseURL string `json:"base_url"` - APIKey string `json:"api_key"` - Adapter string `json:"adapter"` - Temperature float64 `json:"temperature"` - MaxTokens int `json:"max_tokens"` - Sources []LLMSource `json:"sources,omitempty"` + Provider string `json:"provider"` + Model string `json:"model"` + BaseURL string `json:"base_url"` + APIKey string `json:"api_key"` + Adapter string `json:"adapter"` + Temperature float64 `json:"temperature"` + MaxTokens int `json:"max_tokens"` + ThinkingEnabled bool `json:"thinking_enabled"` + Sources []LLMSource `json:"sources,omitempty"` } type Config struct {