360 lines
9.2 KiB
Go
360 lines
9.2 KiB
Go
// 织忆 MemoryWeave — Go Client SDK
|
||
// Hermes / OpenClaw / Cron Jobs 共用
|
||
// 对接织忆 API (port 7821)，提供 Commit / Recall / Feedback / WebSocket
|
||
|
||
package client
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/gorilla/websocket"
|
||
)
|
||
|
||
// ─── Types ───────────────────────────────────────────────
|
||
|
||
// CommitResponse /commit 响应
|
||
type CommitResponse struct {
|
||
MemoryID string `json:"memory_id"`
|
||
Status string `json:"status"`
|
||
}
|
||
|
||
// CommitEntry 批量提交条目
|
||
type CommitEntry struct {
|
||
Content string `json:"content"`
|
||
Category string `json:"category"`
|
||
Namespace string `json:"namespace"`
|
||
}
|
||
|
||
// BatchResponse /batch-commit 响应
|
||
type BatchResponse struct {
|
||
MemoryIDs []string `json:"memory_ids"`
|
||
Status string `json:"status"`
|
||
}
|
||
|
||
// RecallOptions recall 可选参数
|
||
type RecallOptions struct {
|
||
Namespace string `json:"namespace"`
|
||
TopK int `json:"top_k"`
|
||
Diversity float64 `json:"diversity"`
|
||
}
|
||
|
||
// RecallResponse /recall 响应
|
||
type RecallResponse struct {
|
||
Count int `json:"count"`
|
||
Results []RecallResult `json:"results"`
|
||
}
|
||
|
||
// RecallResult 单条召回结果
|
||
type RecallResult struct {
|
||
ID string `json:"id"`
|
||
Content string `json:"content"`
|
||
Category string `json:"category"`
|
||
Score float64 `json:"score"`
|
||
Timestamp string `json:"timestamp"`
|
||
}
|
||
|
||
// RegisterResponse Agent 注册响应
|
||
type RegisterResponse struct {
|
||
AgentID string `json:"agent_id"`
|
||
APIKey string `json:"api_key"`
|
||
WSEndpoint string `json:"ws_endpoint"`
|
||
QuotaRecall int `json:"quota_recall"`
|
||
QuotaCommit int `json:"quota_commit"`
|
||
}
|
||
|
||
// WSEvent WebSocket 推送事件
|
||
type WSEvent struct {
|
||
Type string `json:"type"` // prefetch.push / gap.detected / memory.updated / conflict.detected / ...
|
||
Payload interface{} `json:"payload"`
|
||
}
|
||
|
||
// StatsResponse /api/v1/stats 响应
|
||
type StatsResponse struct {
|
||
Backend string `json:"backend"`
|
||
TotalMemories int `json:"total_memories"`
|
||
TotalEpisodes int `json:"total_episodes"`
|
||
TombstoneCount int `json:"tombstone_count"`
|
||
DataDir string `json:"data_dir"`
|
||
}
|
||
|
||
// ─── Client ──────────────────────────────────────────────
|
||
|
||
// ZhiYiClient 织忆 Go SDK（Hermes/OpenClaw/Cron Jobs 共用）
|
||
type ZhiYiClient struct {
|
||
baseURL string
|
||
apiKey string
|
||
agentID string
|
||
http *http.Client
|
||
wsConn *websocket.Conn
|
||
wsMu sync.Mutex
|
||
eventCh chan WSEvent
|
||
done chan struct{}
|
||
}
|
||
|
||
// NewZhiYiClient 创建织忆客户端
|
||
func NewZhiYiClient(baseURL, apiKey, agentID string) *ZhiYiClient {
|
||
return &ZhiYiClient{
|
||
baseURL: baseURL,
|
||
apiKey: apiKey,
|
||
agentID: agentID,
|
||
http: &http.Client{Timeout: 30 * time.Second},
|
||
eventCh: make(chan WSEvent, 100),
|
||
done: make(chan struct{}),
|
||
}
|
||
}
|
||
|
||
// ─── REST APIs ───────────────────────────────────────────
|
||
|
||
// Commit 提交单条记忆
|
||
func (c *ZhiYiClient) Commit(content, category, namespace string) (*CommitResponse, error) {
|
||
body := map[string]string{
|
||
"content": content,
|
||
"category": category,
|
||
"namespace": namespace,
|
||
}
|
||
if c.agentID != "" {
|
||
body["agent_id"] = c.agentID
|
||
}
|
||
var resp CommitResponse
|
||
if err := c.post("/api/v1/commit", body, &resp); err != nil {
|
||
return nil, err
|
||
}
|
||
return &resp, nil
|
||
}
|
||
|
||
// BatchCommit 批量提交
|
||
func (c *ZhiYiClient) BatchCommit(entries []CommitEntry) (*BatchResponse, error) {
|
||
type batchReq struct {
|
||
Entries []CommitEntry `json:"entries"`
|
||
AgentID string `json:"agent_id"`
|
||
}
|
||
var resp BatchResponse
|
||
if err := c.post("/api/v1/batch-commit", batchReq{Entries: entries, AgentID: c.agentID}, &resp); err != nil {
|
||
return nil, err
|
||
}
|
||
return &resp, nil
|
||
}
|
||
|
||
// Recall 语义搜索
|
||
func (c *ZhiYiClient) Recall(query string, opts RecallOptions) (*RecallResponse, error) {
|
||
if opts.TopK <= 0 {
|
||
opts.TopK = 10
|
||
}
|
||
if opts.Diversity <= 0 {
|
||
opts.Diversity = 0.5
|
||
}
|
||
body := map[string]interface{}{
|
||
"query": query,
|
||
"namespace": opts.Namespace,
|
||
"top_k": opts.TopK,
|
||
"diversity": opts.Diversity,
|
||
}
|
||
var resp RecallResponse
|
||
if err := c.post("/api/v1/recall", body, &resp); err != nil {
|
||
return nil, err
|
||
}
|
||
return &resp, nil
|
||
}
|
||
|
||
// FeedbackUseful 标记记忆有用
|
||
func (c *ZhiYiClient) FeedbackUseful(memoryID string) error {
|
||
return c.post("/api/v1/feedback/useful", map[string]string{"memory_id": memoryID}, nil)
|
||
}
|
||
|
||
// FeedbackNotUseful 标记记忆无用
|
||
func (c *ZhiYiClient) FeedbackNotUseful(memoryID, reason string) error {
|
||
return c.post("/api/v1/feedback/not-useful", map[string]string{
|
||
"memory_id": memoryID,
|
||
"reason": reason,
|
||
}, nil)
|
||
}
|
||
|
||
// FeedbackDeprecate 标记过时
|
||
func (c *ZhiYiClient) FeedbackDeprecate(memoryID, reason string) error {
|
||
return c.post("/api/v1/feedback/deprecate", map[string]string{
|
||
"memory_id": memoryID,
|
||
"reason": reason,
|
||
}, nil)
|
||
}
|
||
|
||
// FeedbackCorrect 提交修正
|
||
func (c *ZhiYiClient) FeedbackCorrect(memoryID, newContent, source string) error {
|
||
return c.post("/api/v1/feedback/correct", map[string]string{
|
||
"memory_id": memoryID,
|
||
"new_content": newContent,
|
||
"source": source,
|
||
}, nil)
|
||
}
|
||
|
||
// RegisterAgent 注册 Agent（获取 key + quota + WS endpoint）
|
||
func (c *ZhiYiClient) RegisterAgent(agentType, namespace string) (*RegisterResponse, error) {
|
||
body := map[string]string{
|
||
"agent_id": c.agentID,
|
||
"agent_type": agentType,
|
||
"namespace": namespace,
|
||
}
|
||
var resp RegisterResponse
|
||
if err := c.post("/api/v1/agents/register", body, &resp); err != nil {
|
||
return nil, err
|
||
}
|
||
c.agentID = resp.AgentID
|
||
c.apiKey = resp.APIKey
|
||
return &resp, nil
|
||
}
|
||
|
||
// Bootstrap 冷启动引导（~10 条核心事实）
|
||
func (c *ZhiYiClient) Bootstrap(namespace string) (*RecallResponse, error) {
|
||
url := fmt.Sprintf("%s/api/v1/bootstrap?namespace=%s", c.baseURL, namespace)
|
||
req, _ := http.NewRequest("GET", url, nil)
|
||
req.Header.Set("X-API-Key", c.apiKey)
|
||
resp, err := c.http.Do(req)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer resp.Body.Close()
|
||
if resp.StatusCode != 200 {
|
||
return nil, fmt.Errorf("bootstrap failed: %d", resp.StatusCode)
|
||
}
|
||
var r RecallResponse
|
||
json.NewDecoder(resp.Body).Decode(&r)
|
||
return &r, nil
|
||
}
|
||
|
||
// Stats 获取系统统计
|
||
func (c *ZhiYiClient) Stats() (*StatsResponse, error) {
|
||
var resp StatsResponse
|
||
if err := c.get("/api/v1/stats", &resp); err != nil {
|
||
return nil, err
|
||
}
|
||
return &resp, nil
|
||
}
|
||
|
||
// Health 健康检查（无需认证）
|
||
func (c *ZhiYiClient) Health() error {
|
||
resp, err := c.http.Get(c.baseURL + "/health")
|
||
if err != nil {
|
||
return err
|
||
}
|
||
resp.Body.Close()
|
||
if resp.StatusCode != 200 {
|
||
return fmt.Errorf("health check failed: %d", resp.StatusCode)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// ─── WebSocket ───────────────────────────────────────────
|
||
|
||
// ListenEvents 建立 WebSocket 连接，实时接收事件
|
||
// 返回只读 channel，调用 Close() 停止
|
||
func (c *ZhiYiClient) ListenEvents(ctx context.Context) (<-chan WSEvent, error) {
|
||
wsURL := fmt.Sprintf("ws://%s/api/v1/ws/%s", c.baseURL[len("http://"):], c.agentID)
|
||
wsURL = "ws" + wsURL[len("http:"):] // 替换 https→wss 简单处理
|
||
|
||
conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("ws connect: %w", err)
|
||
}
|
||
|
||
c.wsMu.Lock()
|
||
c.wsConn = conn
|
||
c.wsMu.Unlock()
|
||
|
||
go func() {
|
||
defer close(c.eventCh)
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
return
|
||
case <-c.done:
|
||
return
|
||
default:
|
||
}
|
||
|
||
_, msg, err := conn.ReadMessage()
|
||
if err != nil {
|
||
return
|
||
}
|
||
var evt WSEvent
|
||
if json.Unmarshal(msg, &evt) == nil {
|
||
c.eventCh <- evt
|
||
}
|
||
}
|
||
}()
|
||
|
||
return c.eventCh, nil
|
||
}
|
||
|
||
// Close 关闭 WebSocket 连接和 channel
|
||
func (c *ZhiYiClient) Close() {
|
||
c.wsMu.Lock()
|
||
defer c.wsMu.Unlock()
|
||
select {
|
||
case <-c.done:
|
||
return
|
||
default:
|
||
close(c.done)
|
||
}
|
||
if c.wsConn != nil {
|
||
c.wsConn.Close()
|
||
}
|
||
}
|
||
|
||
// ─── Internal Helpers ────────────────────────────────────
|
||
|
||
func (c *ZhiYiClient) post(path string, body interface{}, result interface{}) error {
|
||
data, _ := json.Marshal(body)
|
||
req, err := http.NewRequest("POST", c.baseURL+path, bytes.NewReader(data))
|
||
if err != nil {
|
||
return err
|
||
}
|
||
req.Header.Set("Content-Type", "application/json")
|
||
req.Header.Set("X-API-Key", c.apiKey)
|
||
|
||
resp, err := c.http.Do(req)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
if resp.StatusCode >= 400 {
|
||
b, _ := io.ReadAll(resp.Body)
|
||
return fmt.Errorf("zhiyi %s: %d %s", path, resp.StatusCode, string(b))
|
||
}
|
||
|
||
if result != nil {
|
||
return json.NewDecoder(resp.Body).Decode(result)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (c *ZhiYiClient) get(path string, result interface{}) error {
|
||
req, err := http.NewRequest("GET", c.baseURL+path, nil)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
req.Header.Set("X-API-Key", c.apiKey)
|
||
|
||
resp, err := c.http.Do(req)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
if resp.StatusCode >= 400 {
|
||
b, _ := io.ReadAll(resp.Body)
|
||
return fmt.Errorf("zhiyi %s: %d %s", path, resp.StatusCode, string(b))
|
||
}
|
||
|
||
if result != nil {
|
||
return json.NewDecoder(resp.Body).Decode(result)
|
||
}
|
||
return nil
|
||
}
|