mirror of
https://gitcode.com/JianFeeeee/HomeAgent.git
synced 2026-10-04 00:03:59 +00:00
fix: correct context pruning order and vector alignment
- Prune context BEFORE processing (LSTM forget gate pattern) so LLM only sees relevant context, instead of pruning after the fact - Fix ensureTrained() to recompute all event vectors after retraining vectorizer, fixing feature-space mismatch between stored vectors and query vector that made relevance scoring effectively random - Add read lock to knowledge BuildTree() (data race fix) - Log writeIndex() errors instead of discarding them - Fix TOCTOU race in document ContextToDoc() dedup - Fix healthcheck timing (measure elapsed before cleanup) - Refactor waiter CLI into separate files (state, conn, config, editor, history, builtin) for maintainability - Add CLI plugin API key authentication
This commit is contained in:
141
cmd/waiter/conn.go
Normal file
141
cmd/waiter/conn.go
Normal file
@ -0,0 +1,141 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Conn interface {
|
||||
Send(line string) error
|
||||
ReadLine() (string, error)
|
||||
Close() error
|
||||
}
|
||||
|
||||
func dial(cfg *Config) (Conn, error) {
|
||||
if cfg.Remote != "" {
|
||||
return dialRemote(cfg.Remote, cfg.APIKey)
|
||||
}
|
||||
return dialLocal(cfg.Socket, cfg.APIKey)
|
||||
}
|
||||
|
||||
func dialLocal(socket, apiKey string) (Conn, error) {
|
||||
if socket == "" {
|
||||
socket = discoverSocket("")
|
||||
}
|
||||
c, err := net.DialTimeout("unix", socket, 5*time.Second)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unix %s: %w", socket, err)
|
||||
}
|
||||
lc := &localConn{conn: c, r: bufio.NewReader(c)}
|
||||
if apiKey != "" {
|
||||
if err := lc.Send("/auth " + apiKey); err != nil {
|
||||
c.Close()
|
||||
return nil, fmt.Errorf("auth send: %w", err)
|
||||
}
|
||||
line, err := lc.ReadLine()
|
||||
if err != nil {
|
||||
c.Close()
|
||||
return nil, fmt.Errorf("auth read: %w", err)
|
||||
}
|
||||
if strings.Contains(line, "unauthorized") {
|
||||
c.Close()
|
||||
return nil, fmt.Errorf("auth rejected: %s", line)
|
||||
}
|
||||
}
|
||||
return lc, nil
|
||||
}
|
||||
|
||||
type localConn struct {
|
||||
conn net.Conn
|
||||
r *bufio.Reader
|
||||
}
|
||||
|
||||
func (c *localConn) Send(line string) error {
|
||||
_, err := fmt.Fprintf(c.conn, "%s\n", line)
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *localConn) ReadLine() (string, error) {
|
||||
s, err := c.r.ReadString('\n')
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return strings.TrimSuffix(s, "\n"), nil
|
||||
}
|
||||
|
||||
func (c *localConn) Close() error {
|
||||
return c.conn.Close()
|
||||
}
|
||||
|
||||
func dialRemote(baseURL, apiKey string) (Conn, error) {
|
||||
baseURL = strings.TrimRight(baseURL, "/")
|
||||
return &remoteConn{url: baseURL + "/api/v1/chat", apiKey: apiKey}, nil
|
||||
}
|
||||
|
||||
type remoteConn struct {
|
||||
url string
|
||||
apiKey string
|
||||
mu sync.Mutex
|
||||
buf []string
|
||||
closed bool
|
||||
}
|
||||
|
||||
func (c *remoteConn) Send(line string) error {
|
||||
req, err := http.NewRequest("POST", c.url, strings.NewReader(
|
||||
fmt.Sprintf(`{"message":%q}`, line),
|
||||
))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if c.apiKey != "" {
|
||||
req.Header.Set("X-API-Key", c.apiKey)
|
||||
}
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
var result map[string]interface{}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
return err
|
||||
}
|
||||
content, _ := result["response"].(string)
|
||||
|
||||
c.mu.Lock()
|
||||
c.buf = append(c.buf, content)
|
||||
c.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *remoteConn) ReadLine() (string, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
for len(c.buf) == 0 && !c.closed {
|
||||
c.mu.Unlock()
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
c.mu.Lock()
|
||||
}
|
||||
if c.closed && len(c.buf) == 0 {
|
||||
return "", io.EOF
|
||||
}
|
||||
s := c.buf[0]
|
||||
c.buf = c.buf[1:]
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func (c *remoteConn) Close() error {
|
||||
c.mu.Lock()
|
||||
c.closed = true
|
||||
c.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user