memoryweave/go/internal/consolidate/client.go

// 织忆 MemoryWeave — Go ↔ Rust IPC 客户端
// 通信方式: Unix Socket + Protobuf + length-prefixed framing
// 协议: proto/consolidate.proto
// Rust 端: rust/src/main.rs (Unix Socket 监听)
//
// 每次深度整合：Go 发送 ConsolidationRequest → Rust 执行 → 返回 ConsolidationResponse

package consolidate

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"net"
	"os"
	"time"
)

// ─── 消息定义（对应 proto/consolidate.proto）─────────────────

// ConsolidateRequest 整合请求
type ConsolidateRequest struct {
	Task         string  `json:"task"`          // "full" | "cluster_only" | "prune_only"
	LanceDBPath  string  `json:"lancedb_path"`  // LanceDB 数据目录
	SQLitePath   string  `json:"sqlite_path"`   // SQLite 图谱路径
	LLMEndpoint  string  `json:"llm_endpoint"`  // LLM API 端点
	LLMModel     string  `json:"llm_model"`     // LLM 模型名
	LLMBudget    int     `json:"llm_budget"`    // 本次可用 LLM 次数
	Epsilon      float64 `json:"epsilon"`       // DBSCAN 邻域半径
	MinPoints    int     `json:"min_points"`    // DBSCAN 最小点数
}

// ConsolidateResponse 整合响应
type ConsolidateResponse struct {
	Status       string `json:"status"`        // "ok" | "partial_failure"
	ReportJSON   string `json:"report_json"`   // ConsolidationReport JSON
	FailureStep  string `json:"failure_step"`  // 失败步骤
	ErrorDetail  string `json:"error_detail"`  // 错误详情
}

// Result 解析后的整合结果
type Result struct {
	Mode       string             `json:"mode"`
	Timestamp  string             `json:"timestamp"`
	Clusters   int                `json:"clusters,omitempty"`
	Noise      int                `json:"noise,omitempty"`
	DecayRates map[string]float64 `json:"decay_rates,omitempty"`
	Quality    *QualityResult     `json:"quality,omitempty"`
}

type QualityResult struct {
	Score          float64 `json:"score"`
	LowInfo        int     `json:"low_info"`
	Total          int     `json:"total"`
	Hallucinations int     `json:"hallucinations"`
}

// ─── IPC Client ───────────────────────────────────────────

const (
	defaultSocketPath = "/tmp/zhiyi-ipc.sock"
	defaultTimeout    = 10 * time.Minute // 深度整合 5 分钟 + 缓冲
)

// Run 通过 Unix Socket 调 Rust zhiyi-consolidate，执行深度整合
func Run(dataDir, sqlitePath, mode string) (*Result, error) {
	return RunWithOptions(ConsolidateRequest{
		Task:        mode,
		LanceDBPath: dataDir,
		SQLitePath:  sqlitePath,
		LLMBudget:   20,
		Epsilon:     0.5,
		MinPoints:   3,
	})
}

// RunWithOptions 完整参数调用
func RunWithOptions(req ConsolidateRequest) (*Result, error) {
	socketPath := os.Getenv("ZHIYI_IPC_SOCKET")
	if socketPath == "" {
		socketPath = defaultSocketPath
	}

	// 连接 Unix Socket
	conn, err := net.DialTimeout("unix", socketPath, 5*time.Second)
	if err != nil {
		return nil, fmt.Errorf("connect to %s: %w (is zhiyi-consolidate running?)", socketPath, err)
	}
	defer conn.Close()

	// 设置超时
	conn.SetDeadline(time.Now().Add(defaultTimeout))

	// 序列化请求
	reqJSON, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("marshal request: %w", err)
	}

	// 发送: 4-byte length prefix + JSON body
	msgLen := make([]byte, 4)
	binary.BigEndian.PutUint32(msgLen, uint32(len(reqJSON)))
	if _, err := conn.Write(msgLen); err != nil {
		return nil, fmt.Errorf("write length: %w", err)
	}
	if _, err := conn.Write(reqJSON); err != nil {
		return nil, fmt.Errorf("write body: %w", err)
	}

	// 读取响应长度
	lenBuf := make([]byte, 4)
	if _, err := conn.Read(lenBuf); err != nil {
		return nil, fmt.Errorf("read response length: %w", err)
	}
	respLen := binary.BigEndian.Uint32(lenBuf)

	// 读取响应体
	respBody := make([]byte, respLen)
	n, err := conn.Read(respBody)
	if err != nil {
		return nil, fmt.Errorf("read response body: %w", err)
	}
	if n < int(respLen) {
		return nil, fmt.Errorf("truncated response: got %d, expected %d", n, respLen)
	}

	// 解析响应
	var resp ConsolidateResponse
	if err := json.Unmarshal(respBody[:respLen], &resp); err != nil {
		return nil, fmt.Errorf("unmarshal response: %w\nbody: %s", err, string(respBody[:respLen]))
	}

	if resp.Status != "ok" {
		return nil, fmt.Errorf("consolidate %s: step=%s, %s", resp.Status, resp.FailureStep, resp.ErrorDetail)
	}

	// 解析 ReportJSON 为 Result
	var result Result
	if err := json.Unmarshal([]byte(resp.ReportJSON), &result); err != nil {
		return nil, fmt.Errorf("unmarshal report: %w", err)
	}

	return &result, nil
}

// HealthCheck 检查 Rust sidecar 是否存活
func HealthCheck() error {
	socketPath := os.Getenv("ZHIYI_IPC_SOCKET")
	if socketPath == "" {
		socketPath = defaultSocketPath
	}

	conn, err := net.DialTimeout("unix", socketPath, 2*time.Second)
	if err != nil {
		return fmt.Errorf("zhiyi-consolidate not reachable: %w", err)
	}
	conn.Close()
	return nil
}