557 lines
14 KiB
Go
557 lines
14 KiB
Go
// 织忆 CLI — 命令行工具（零外部依赖，纯 flag 实现）
|
||
package main
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/json"
|
||
"flag"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"os"
|
||
"sort"
|
||
"strings"
|
||
"unicode/utf8"
|
||
)
|
||
|
||
const (
|
||
apiKeyDefault = "zhiyi-dev-key-2026"
|
||
apiURLDefault = "http://localhost:7821"
|
||
namespaceDefault = "hermes-main"
|
||
)
|
||
|
||
// ─── 全局参数 ──────────────────────────────────────────────
|
||
|
||
var (
|
||
apiURL = flag.String("url", getEnv("ZHIYI_API_URL", apiURLDefault), "织忆 API 地址")
|
||
apiKey = flag.String("key", getEnv("ZHIYI_API_KEY", apiKeyDefault), "API Key")
|
||
ns = flag.String("n", getEnv("ZHIYI_NAMESPACE", namespaceDefault), "命名空间")
|
||
)
|
||
|
||
// ─── 入口 ─────────────────────────────────────────────────
|
||
|
||
func main() {
|
||
flag.Usage = usage
|
||
flag.Parse()
|
||
|
||
if flag.NArg() == 0 {
|
||
flag.Usage()
|
||
os.Exit(1)
|
||
}
|
||
|
||
cmd := flag.Arg(0)
|
||
args := flag.Args()[1:]
|
||
|
||
var err error
|
||
switch cmd {
|
||
case "tree": err = runTree(args)
|
||
case "graph": err = runGraph(args)
|
||
case "recall": err = runRecall(args)
|
||
case "stats": err = runStats()
|
||
case "entity": err = runEntity(args)
|
||
case "help", "--help", "-h":
|
||
flag.Usage()
|
||
os.Exit(0)
|
||
default:
|
||
fmt.Fprintf(os.Stderr, "未知命令: %s\n", cmd)
|
||
flag.Usage()
|
||
os.Exit(1)
|
||
}
|
||
if err != nil {
|
||
fmt.Fprintln(os.Stderr, "Error:", err)
|
||
os.Exit(1)
|
||
}
|
||
}
|
||
|
||
func getEnv(key, fallback string) string {
|
||
if v := os.Getenv(key); v != "" {
|
||
return v
|
||
}
|
||
return fallback
|
||
}
|
||
|
||
func usage() {
|
||
fmt.Fprint(os.Stderr, `织忆 CLI — 记忆系统命令行工具
|
||
|
||
用法:
|
||
zhiyi [全局选项] <命令> [命令参数]
|
||
|
||
全局选项:
|
||
-url <地址> 织忆 API 地址（默认 http://localhost:7821）
|
||
-key <key> API Key（默认 ZHIYI_API_KEY 环境变量）
|
||
-n <ns> 命名空间（默认 hermes-main）
|
||
|
||
命令:
|
||
tree 树形展示记忆结构（按 category 分组）
|
||
graph [实体] ASCII 渲染 ego-network 图谱（省略实体自动取 Top-1）
|
||
recall <query> 语义搜索，返回 top-10 结果
|
||
stats 显示系统统计（记忆数、蒸馏状态等）
|
||
entity <name> 查询实体详情（出现次数、关联记忆）
|
||
|
||
示例:
|
||
zhiyi tree
|
||
zhiyi graph 牧尘
|
||
zhiyi recall 牧尘的偏好
|
||
zhiyi stats
|
||
zhiyi entity 织忆
|
||
`)
|
||
}
|
||
|
||
// ─── API 调用 ──────────────────────────────────────────────
|
||
|
||
func apiGet(path string, v interface{}) error {
|
||
return getJSON(*apiURL+path, *apiKey, v)
|
||
}
|
||
|
||
func apiPost(path string, body, v interface{}) error {
|
||
return postJSON(*apiURL+path, *apiKey, body, v)
|
||
}
|
||
|
||
func getJSON(url, apiKey string, v interface{}) error {
|
||
req, err := http.NewRequest("GET", url, nil)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
req.Header.Set("X-API-Key", apiKey)
|
||
resp, err := http.DefaultClient.Do(req)
|
||
if err != nil {
|
||
return fmt.Errorf("请求失败: %w", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
if resp.StatusCode != 200 {
|
||
body, _ := io.ReadAll(resp.Body)
|
||
return fmt.Errorf("API %d: %s", resp.StatusCode, string(body))
|
||
}
|
||
return json.NewDecoder(resp.Body).Decode(v)
|
||
}
|
||
|
||
func postJSON(url, apiKey string, body interface{}, v interface{}) error {
|
||
bodyBytes, _ := json.Marshal(body)
|
||
req, err := http.NewRequest("POST", url, bytes.NewReader(bodyBytes))
|
||
if err != nil {
|
||
return err
|
||
}
|
||
req.Header.Set("X-API-Key", apiKey)
|
||
req.Header.Set("Content-Type", "application/json")
|
||
resp, err := http.DefaultClient.Do(req)
|
||
if err != nil {
|
||
return fmt.Errorf("请求失败: %w", err)
|
||
}
|
||
defer resp.Body.Close()
|
||
if resp.StatusCode != 200 {
|
||
body, _ := io.ReadAll(resp.Body)
|
||
return fmt.Errorf("API %d: %s", resp.StatusCode, string(body))
|
||
}
|
||
return json.NewDecoder(resp.Body).Decode(v)
|
||
}
|
||
|
||
// ─── 命令实现 ──────────────────────────────────────────────
|
||
|
||
func runTree(args []string) error {
|
||
// 从图谱 export 获取节点（带 category），按 category 分组
|
||
var r struct {
|
||
Nodes []struct {
|
||
ID string `json:"id"`
|
||
Name string `json:"name"`
|
||
Type string `json:"type,omitempty"`
|
||
Category string `json:"category,omitempty"`
|
||
Weight float64 `json:"weight,omitempty"`
|
||
} `json:"nodes"`
|
||
Edges []struct {
|
||
Source string `json:"source"`
|
||
Target string `json:"target"`
|
||
Rel string `json:"relation"`
|
||
} `json:"edges"`
|
||
Count struct {
|
||
Nodes int `json:"nodes"`
|
||
Edges int `json:"edges"`
|
||
} `json:"count"`
|
||
}
|
||
if err := apiGet(fmt.Sprintf("/api/v1/graph/export?namespace=%s&limit=500", *ns), &r); err != nil {
|
||
return err
|
||
}
|
||
|
||
// 按 type（实体类型）分组，忽略 episode 节点
|
||
groups := make(map[string][]string)
|
||
for _, n := range r.Nodes {
|
||
// 过滤 episode 节点（以 ep_ 开头的是 episode ID）
|
||
if strings.HasPrefix(n.ID, "ep_") {
|
||
continue
|
||
}
|
||
cat := n.Category
|
||
if cat == "" {
|
||
cat = n.Type
|
||
}
|
||
if cat == "" || cat == "entity" {
|
||
cat = "概念"
|
||
}
|
||
groups[cat] = append(groups[cat], n.Name)
|
||
}
|
||
|
||
cats := make([]string, 0, len(groups))
|
||
for k := range groups {
|
||
cats = append(cats, k)
|
||
}
|
||
sort.Strings(cats)
|
||
|
||
fmt.Printf("🧠 织忆记忆树 [%s] (%d 节点 %d 边)\n", *ns, r.Count.Nodes, r.Count.Edges)
|
||
fmt.Println(strings.Repeat("─", 50))
|
||
|
||
catIcons := map[string]string{
|
||
"人物": "👤",
|
||
"项目": "📦",
|
||
"事件": "📅",
|
||
"概念": "💡",
|
||
"位置": "📍",
|
||
"组织": "🏢",
|
||
"distilled": "🔄",
|
||
}
|
||
for _, cat := range cats {
|
||
items := groups[cat]
|
||
icon := catIcons[cat]
|
||
if icon == "" {
|
||
icon = "📄"
|
||
}
|
||
fmt.Printf("\n%s %s (%d)\n", icon, cat, len(items))
|
||
for i, label := range items {
|
||
if i >= 20 {
|
||
fmt.Printf(" … 还有 %d 个实体\n", len(items)-20)
|
||
break
|
||
}
|
||
fmt.Printf(" %2d. %s\n", i+1, label)
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func runGraph(args []string) error {
|
||
var entity string
|
||
if len(args) == 0 {
|
||
// 从 pagerank map 取 score 最高的非 episode 实体
|
||
var pr struct {
|
||
Pagerank map[string]float64 `json:"pagerank"`
|
||
Count int `json:"count"`
|
||
}
|
||
if err := apiGet("/api/v1/graph/pagerank", &pr); err != nil {
|
||
return fmt.Errorf("获取 Top 实体失败: %w", err)
|
||
}
|
||
var topEntity string
|
||
var topScore float64
|
||
for e, s := range pr.Pagerank {
|
||
if !strings.HasPrefix(e, "ep_") && s > topScore {
|
||
topScore = s
|
||
topEntity = e
|
||
}
|
||
}
|
||
if topEntity == "" {
|
||
return fmt.Errorf("图谱为空，无实体")
|
||
}
|
||
entity = stripPrefix(topEntity)
|
||
} else {
|
||
entity = args[0]
|
||
}
|
||
|
||
type navigateResp struct {
|
||
Entity string `json:"entity"`
|
||
Count int `json:"count"`
|
||
Paths []struct {
|
||
From string `json:"from"`
|
||
To string `json:"to"`
|
||
Relation string `json:"relation"`
|
||
Hop int `json:"hop"`
|
||
Weight float64 `json:"weight"`
|
||
} `json:"paths"`
|
||
}
|
||
var resp navigateResp
|
||
if err := apiPost("/api/v1/graph/navigate", map[string]interface{}{
|
||
"entity": entity, "max_hops": 1, "namespace": *ns,
|
||
}, &resp); err != nil {
|
||
return fmt.Errorf("获取邻居失败: %w", err)
|
||
}
|
||
|
||
// 统计关系和邻居
|
||
neighbors := make([]string, 0)
|
||
seen := make(map[string]bool)
|
||
for _, p := range resp.Paths {
|
||
if p.From == entity || p.From == "n_"+entity {
|
||
neighbor := stripPrefix(p.To)
|
||
if !seen[neighbor] {
|
||
neighbors = append(neighbors, neighbor)
|
||
seen[neighbor] = true
|
||
}
|
||
} else if p.To == entity || p.To == "n_"+entity {
|
||
neighbor := stripPrefix(p.From)
|
||
if !seen[neighbor] {
|
||
neighbors = append(neighbors, neighbor)
|
||
seen[neighbor] = true
|
||
}
|
||
}
|
||
}
|
||
|
||
relMap := make(map[string]string)
|
||
for _, p := range resp.Paths {
|
||
if p.From == entity || p.From == "n_"+entity {
|
||
relMap[stripPrefix(p.To)] = p.Relation
|
||
} else if p.To == entity || p.To == "n_"+entity {
|
||
relMap[stripPrefix(p.From)] = p.Relation
|
||
}
|
||
}
|
||
|
||
printASCIIGraph(entity, neighbors, relMap)
|
||
return nil
|
||
}
|
||
|
||
func printASCIIGraph(center string, neighbors []string, rels map[string]string) {
|
||
fmt.Printf("📐 织忆图谱 — %s\n", center)
|
||
fmt.Println(strings.Repeat("─", 50))
|
||
|
||
if len(neighbors) == 0 {
|
||
fmt.Println(" (无邻居)")
|
||
return
|
||
}
|
||
|
||
// 中心节点（未使用渲染，留空用于后续扩展）
|
||
_ = fmt.Sprintf(" %s %s", box("center", center), color("dim", "[中心节点]"))
|
||
|
||
// 分两列最多 6 个邻居
|
||
sort.Strings(neighbors)
|
||
_ = neighbors[:len(neighbors)/2]
|
||
|
||
// 计算分支
|
||
fmt.Println("")
|
||
fmt.Printf(" ┌─── %s%s ───┐\n", color("green", "◉"), color("bright", center))
|
||
for i, n := range neighbors {
|
||
rel := rels[n]
|
||
relStr := ""
|
||
if rel != "" {
|
||
relStr = color("dim", "("+rel+")")
|
||
}
|
||
prefix := " │"
|
||
if i < len(neighbors)-1 {
|
||
fmt.Printf("%s ○ %s %s\n", prefix, color("cyan", truncate(n, 12)), relStr)
|
||
} else {
|
||
fmt.Printf("%s ○ %s %s\n", prefix, color("cyan", truncate(n, 12)), relStr)
|
||
}
|
||
}
|
||
fmt.Printf(" └%s (%d 个邻居)\n", strings.Repeat("─", 20), len(neighbors))
|
||
}
|
||
|
||
func runRecall(args []string) error {
|
||
if len(args) == 0 {
|
||
return fmt.Errorf("用法: zhiyi recall <query>")
|
||
}
|
||
query := args[0]
|
||
|
||
var resp struct {
|
||
Results []struct {
|
||
ID string `json:"id"`
|
||
Content string `json:"content"`
|
||
Category string `json:"category"`
|
||
Score float64 `json:"score"`
|
||
} `json:"results"`
|
||
Count int `json:"count"`
|
||
}
|
||
if err := apiPost("/api/v1/recall", map[string]interface{}{
|
||
"query": query, "namespace": *ns, "top_k": 10,
|
||
}, &resp); err != nil {
|
||
return err
|
||
}
|
||
|
||
fmt.Printf("🔍 搜索: %s (%d 结果)\n", query, resp.Count)
|
||
fmt.Println(strings.Repeat("─", 50))
|
||
|
||
for i, m := range resp.Results {
|
||
_ = scoreBar(m.Score)
|
||
fmt.Printf("\n[%d] %s %.3f %s\n", i+1, color("green", "●"), m.Score, color("dim", m.Category))
|
||
fmt.Printf(" %s\n", truncate(m.Content, 100))
|
||
fmt.Printf(" %s\n", color("faint", m.ID))
|
||
}
|
||
if len(resp.Results) == 0 {
|
||
fmt.Println(" (无结果)")
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func runStats() error {
|
||
var stats struct {
|
||
TotalMemories int `json:"total_memories"`
|
||
TotalEpisodes int `json:"total_episodes"`
|
||
Backend string `json:"backend"`
|
||
TombstoneCount int `json:"tombstone_count"`
|
||
}
|
||
var quota struct {
|
||
Remaining int `json:"remaining"`
|
||
Used int `json:"used"`
|
||
Limit int `json:"limit"`
|
||
Status string `json:"status"`
|
||
}
|
||
var distStatus struct {
|
||
QueueLen int `json:"queue_len"`
|
||
DailyUsed int `json:"daily_used"`
|
||
DailyLimit int `json:"daily_limit"`
|
||
BatchSize int `json:"batch_size"`
|
||
}
|
||
|
||
apiGet("/api/v1/stats", &stats)
|
||
apiGet("/api/v1/distill/quota", "a)
|
||
apiGet("/api/v1/distill/status", &distStatus)
|
||
|
||
barLen := 40
|
||
filled := 0
|
||
if quota.Limit > 0 {
|
||
filled = int(float64(barLen) * float64(quota.Used) / float64(quota.Limit))
|
||
if filled > barLen {
|
||
filled = barLen
|
||
}
|
||
}
|
||
|
||
fmt.Println("🧠 织忆系统状态")
|
||
fmt.Println(strings.Repeat("─", 50))
|
||
fmt.Printf(" 📊 记忆总数: %d (episodes: %d, backend: %s)\n", stats.TotalMemories, stats.TotalEpisodes, stats.Backend)
|
||
fmt.Printf(" 🗑️ 墓碑: %d\n", stats.TombstoneCount)
|
||
fmt.Printf(" ⚙️ 蒸馏队列: %d 条\n", distStatus.QueueLen)
|
||
fmt.Println("")
|
||
fmt.Printf(" 📈 蒸馏配额 [%-*s] %d/%d (%s)\n",
|
||
barLen, strings.Repeat("█", filled)+strings.Repeat("░", barLen-filled),
|
||
quota.Used, quota.Limit, quota.Status)
|
||
|
||
return nil
|
||
}
|
||
|
||
func runEntity(args []string) error {
|
||
if len(args) == 0 {
|
||
return fmt.Errorf("用法: zhiyi entity <name>")
|
||
}
|
||
entity := args[0]
|
||
|
||
// 从 navigate 获取证据数量（并发）
|
||
type navResp struct {
|
||
Entity string `json:"entity"`
|
||
Count int `json:"count"`
|
||
}
|
||
var (
|
||
evCh chan navResp
|
||
mrCh chan struct {
|
||
Results []struct {
|
||
ID string `json:"id"`
|
||
Content string `json:"content"`
|
||
} `json:"results"`
|
||
}
|
||
)
|
||
|
||
evCh = make(chan navResp, 1)
|
||
mrCh = make(chan struct {
|
||
Results []struct {
|
||
ID string `json:"id"`
|
||
Content string `json:"content"`
|
||
} `json:"results"`
|
||
}, 1)
|
||
|
||
// 并发请求（navigate 和 recall 都用 POST）
|
||
go func() {
|
||
var nav navResp
|
||
type navPostReq struct {
|
||
Entity string `json:"entity"`
|
||
MaxHops int `json:"max_hops"`
|
||
Namespace string `json:"namespace"`
|
||
}
|
||
if err := postJSON(*apiURL+"/api/v1/graph/navigate", *apiKey, navPostReq{entity, 1, *ns}, &nav); err == nil {
|
||
evCh <- nav
|
||
} else {
|
||
evCh <- navResp{}
|
||
}
|
||
}()
|
||
go func() {
|
||
var recallResp struct {
|
||
Results []struct {
|
||
ID string `json:"id"`
|
||
Content string `json:"content"`
|
||
} `json:"results"`
|
||
}
|
||
if err := apiPost("/api/v1/recall", map[string]interface{}{
|
||
"query": entity, "namespace": *ns, "top_k": 10,
|
||
}, &recallResp); err == nil {
|
||
mrCh <- recallResp
|
||
} else {
|
||
mrCh <- struct {
|
||
Results []struct {
|
||
ID string `json:"id"`
|
||
Content string `json:"content"`
|
||
} `json:"results"`
|
||
}{}
|
||
}
|
||
}()
|
||
|
||
nav := <-evCh
|
||
recallData := <-mrCh
|
||
|
||
fmt.Printf("📌 实体: %s\n", entity)
|
||
fmt.Println(strings.Repeat("─", 50))
|
||
fmt.Printf(" 🔢 邻居数量: %d\n", nav.Count)
|
||
fmt.Printf(" 📄 关联记忆: %d 条\n", len(recallData.Results))
|
||
if len(recallData.Results) > 0 {
|
||
fmt.Println("")
|
||
for i, m := range recallData.Results {
|
||
if i >= 10 {
|
||
fmt.Printf(" … 还有 %d 条\n", len(recallData.Results)-10)
|
||
break
|
||
}
|
||
fmt.Printf(" [%d] %s\n", i+1, truncate(m.Content, 80))
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// ─── 辅助函数 ──────────────────────────────────────────────
|
||
|
||
func truncate(s string, max int) string {
|
||
if utf8.RuneCountInString(s) <= max {
|
||
return s
|
||
}
|
||
r := []rune(s)
|
||
return string(r[:max-1]) + "…"
|
||
}
|
||
|
||
// stripPrefix removes the "n_" prefix from entity names if present
|
||
func stripPrefix(s string) string {
|
||
if strings.HasPrefix(s, "n_") {
|
||
return s[2:]
|
||
}
|
||
return s
|
||
}
|
||
|
||
func scoreBar(score float64) string {
|
||
n := int(score * 10)
|
||
if n > 10 {
|
||
n = 10
|
||
}
|
||
return strings.Repeat("█", n) + strings.Repeat("░", 10-n)
|
||
}
|
||
|
||
// ANSI 颜色
|
||
func color(c, s string) string {
|
||
m := map[string]string{
|
||
"green": "\033[32m",
|
||
"cyan": "\033[36m",
|
||
"dim": "\033[2m",
|
||
"bright": "\033[1m",
|
||
"faint": "\033[2m",
|
||
}
|
||
magenta := "\033[35m"
|
||
reset := "\033[0m"
|
||
if col, ok := m[c]; ok {
|
||
return col + s + reset
|
||
}
|
||
if c == "magenta" {
|
||
return magenta + s + reset
|
||
}
|
||
return s
|
||
}
|
||
|
||
func box(style, s string) string {
|
||
switch style {
|
||
case "center":
|
||
return "◉ " + s
|
||
}
|
||
return s
|
||
} |