Memory: Multi-Session Facts and Lightweight RAG (Part 5)

// Part 5 demo: memory and lightweight RAG.
// Differences from Part 3 are in the article's "what changed" table; the core is memoryStore
// (JSON on disk), the retriever (vector search with a keyword fallback) and the four new tools.
//
// Run: go run main.go        — runs sessions A / B / C in order
//
//	go run main.go recall   — runs session B alone (proves cross-process recall, see Section 5)
package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"math"
	"net/http"
	"os"
	"sort"
	"strings"
	"time"
)

// ===========================================
// Structures matching /v1/chat/completions
// ===========================================

type Message struct {
	Role       string     `json:"role"`
	Content    string     `json:"content,omitempty"`
	ToolCalls  []ToolCall `json:"tool_calls,omitempty"`
	ToolCallID string     `json:"tool_call_id,omitempty"`
}

type ToolCall struct {
	ID       string       `json:"id"`
	Type     string       `json:"type"`
	Function FunctionCall `json:"function"`
}

type Tool struct {
	Type     string      `json:"type"`
	Function FunctionDef `json:"function"`
}

type FunctionDef struct {
	Name        string         `json:"name"`
	Description string         `json:"description"`
	Parameters  map[string]any `json:"parameters"`
}

// FunctionCall is explained in Part 2 §3 (tool_calls[].function)
type FunctionCall struct {
	Name      string `json:"name"`
	Arguments string `json:"arguments"`
}

type ChatRequest struct {
	Model       string    `json:"model"`
	Messages    []Message `json:"messages"`
	Tools       []Tool    `json:"tools,omitempty"`
	Stream      bool      `json:"stream"`
	Temperature float64   `json:"temperature"`
}

type Choice struct {
	Message      Message `json:"message"`
	FinishReason string  `json:"finish_reason"`
}

type ChatResponse struct {
	Choices []Choice `json:"choices"`
}

// ===========================================
// Memory store (cross-session, persisted to a file)
// ===========================================

type memoryStore struct {
	path string
	data map[string]string
}

func loadMemory(path string) *memoryStore {
	m := &memoryStore{path: path, data: map[string]string{}}
	if b, err := os.ReadFile(path); err == nil {
		_ = json.Unmarshal(b, &m.data)
	}
	return m
}

func (m *memoryStore) set(key, value string) error {
	m.data[key] = value
	b, _ := json.MarshalIndent(m.data, "", "  ")
	return os.WriteFile(m.path, b, 0o644)
}

// ===========================================
// Note store (standing in for "local material") and the retriever
// ===========================================

var notes = []string{
	"Ollama 默认只识别 NVIDIA(CUDA) 与 AMD(ROCm) 显卡；Intel Arc 需要在 systemd override 里设置 OLLAMA_VULKAN=true 并重启服务，才会走 Vulkan 后端。",
	"验证 Vulkan 是否生效：journalctl -u ollama --no-pager | grep 'inference compute'，看到 library=Vulkan 与 Arc A770 的描述即为成功。",
	"OLLAMA_LOAD_TIMEOUT 默认 5 分钟，控制模型加载停滞多久后放弃；驱动慢或首次编译着色器时可调大，例如 10m。",
	"OLLAMA_KEEP_ALIVE 默认 5 分钟，控制模型空闲多久后卸载；频繁调用可调大到 10m 以减少重复加载。",
	"Vulkan 后端在部分 Linux 内核 + Mesa 驱动下有显存记账失步、空闲显存被换出的已知问题，长时间运行要用 intel_gpu_top 或 journalctl 监控。",
	"llama3.1:8b 在 Intel Arc A770 上实测 33/33 层全量 offload，模型占显存约 4.4GB，生成速度约 41 tokens/s。",
}

// retriever finds the most relevant note fragments for a query.
// metric() names what score means: cosine similarity for the vector path, matched-word
// count for the keyword fallback. The two are different units, so they must not share
// the "similarity" label when printed.
type retriever interface {
	search(query string, topK int) []searchHit
	metric() string
}

type searchHit struct {
	text  string
	score float64
}

// vectorRetriever uses the vectors from /api/embed for cosine similarity.
// nomic-embed-text only performs well with the standard prefixes: documents get
// "search_document:" and queries get "search_query:" (the embedding model's own
// convention).
type vectorRetriever struct{}

func (vectorRetriever) metric() string { return "相似度" }

func (vectorRetriever) search(query string, topK int) []searchHit {
	texts := make([]string, 0, len(notes)+1)
	texts = append(texts, "search_query: "+query)
	for _, n := range notes {
		texts = append(texts, "search_document: "+n)
	}
	vecs, err := embedTexts(texts)
	if err != nil {
		fmt.Println("   ⚠️ 向量检索失败，本次降级为空结果:", err)
		return nil
	}
	if len(vecs) != len(texts) { // avoid an out-of-range panic when the count does not match the input
		fmt.Println("   ⚠️ 向量检索返回数量不符，本次降级为空结果")
		return nil
	}
	q := vecs[0]
	var hits []searchHit
	for i, n := range notes {
		s := cosine(q, vecs[i+1])
		if s >= 0.35 { // too low a threshold drags unrelated fragments in as well
			hits = append(hits, searchHit{n, s})
		}
	}
	return top(hits, topK)
}

// keywordRetriever is the fallback when there is no embedding model: score by the number of matched words.
type keywordRetriever struct{}

func (keywordRetriever) metric() string { return "命中词数" }

func (keywordRetriever) search(query string, topK int) []searchHit {
	words := tokenize(query)
	var hits []searchHit
	for _, n := range notes {
		score := 0.0
		lower := strings.ToLower(n) // case-insensitive for English (e.g. intel ↔ Intel)
		for _, w := range words {
			if strings.Contains(lower, w) {
				score++
			}
		}
		if score > 0 {
			hits = append(hits, searchHit{n, score})
		}
	}
	return top(hits, topK)
}

func tokenize(s string) []string {
	// Chinese is split into 2-character sliding windows plus English split on
	// whitespace, which is enough for a demo
	var out []string
	runes := []rune(s)
	for i := 0; i+1 < len(runes); i++ {
		if runes[i] < 128 && runes[i+1] < 128 {
			continue
		}
		out = append(out, string(runes[i:i+2]))
	}
	for _, w := range strings.FieldsFunc(s, func(r rune) bool { return !(r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9') }) {
		if len(w) > 1 {
			out = append(out, strings.ToLower(w))
		}
	}
	return out
}

func top(hits []searchHit, k int) []searchHit {
	for i := 1; i < len(hits); i++ {
		for j := i; j > 0 && hits[j].score > hits[j-1].score; j-- {
			hits[j], hits[j-1] = hits[j-1], hits[j]
		}
	}
	if len(hits) > k {
		hits = hits[:k]
	}
	return hits
}

func cosine(a, b []float64) float64 {
	if len(a) != len(b) { // different dimensions: return 0 instead of panicking on an out-of-range index
		return 0
	}
	var dot, na, nb float64
	for i := range a {
		dot += a[i] * b[i]
		na += a[i] * a[i]
		nb += b[i] * b[i]
	}
	if na == 0 || nb == 0 {
		return 0
	}
	return dot / (math.Sqrt(na) * math.Sqrt(nb))
}

// ===========================================
// Tool registration
// ===========================================

type toolEntry struct {
	description string
	parameters  map[string]any
	run         func(args json.RawMessage) (string, error)
}

var (
	mem                = loadMemory("memory.json")
	retrieve retriever = newRetriever()
)

func newRetriever() retriever {
	// Probe /api/tags for an embedding model; use vectors if present, otherwise
	// fall back to keywords. This HTTP call sits on the startup path, so it needs a
	// timeout: a stuck Ollama must not stall startup (Part 2's "give every upstream
	// request an explicit timeout budget" applies here too).
	probe := &http.Client{Timeout: 3 * time.Second}
	resp, err := probe.Get("http://localhost:11434/api/tags")
	if err == nil {
		defer resp.Body.Close()
		var tags struct {
			Models []struct {
				Name string `json:"name"`
			} `json:"models"`
		}
		if json.NewDecoder(resp.Body).Decode(&tags) == nil {
			for _, m := range tags.Models {
				if strings.Contains(m.Name, "nomic-embed-text") {
					fmt.Println("📚 检测到 nomic-embed-text，启用向量检索")
					return vectorRetriever{}
				}
			}
		}
	}
	fmt.Println("📚 未检测到 embedding 模型，降级为关键词检索（可 ollama pull nomic-embed-text 升级）")
	return keywordRetriever{}
}

var tools = map[string]*toolEntry{
	"remember": {
		description: "把一条事实写入长期记忆（跨会话保留）。key 用命名空间风格，如 用户:名字、项目:目标",
		parameters: map[string]any{"type": "object", "properties": map[string]any{
			"key":   map[string]any{"type": "string", "description": "记忆的键，如 用户:名字"},
			"value": map[string]any{"type": "string", "description": "记忆的内容"},
		}, "required": []string{"key", "value"}},
		run: func(args json.RawMessage) (string, error) {
			var p struct {
				Key   string `json:"key"`
				Value string `json:"value"`
			}
			if err := json.Unmarshal(args, &p); err != nil {
				return "", fmt.Errorf("参数解析失败：%v", err)
			}
			if err := mem.set(p.Key, p.Value); err != nil {
				return "", err
			}
			return fmt.Sprintf("已记住 %s = %s", p.Key, p.Value), nil
		},
	},
	"recall": {
		description: "从长期记忆里读取一条事实。key 与 remember 写入时一致；不确定键名时先用 list_memory 列出",
		parameters: map[string]any{"type": "object", "properties": map[string]any{
			"key": map[string]any{"type": "string", "description": "记忆的键"},
		}, "required": []string{"key"}},
		run: func(args json.RawMessage) (string, error) {
			var p struct {
				Key string `json:"key"`
			}
			if err := json.Unmarshal(args, &p); err != nil {
				return "", fmt.Errorf("参数解析失败：%v", err)
			}
			if v, ok := mem.data[p.Key]; ok {
				return v, nil
			}
			return fmt.Sprintf("长期记忆中没有 %s 的记录", p.Key), nil
		},
	},
	"list_memory": {
		description: "列出长期记忆里所有的键与值，用于不确定键名时查看",
		parameters:  map[string]any{"type": "object", "properties": map[string]any{}},
		run: func(args json.RawMessage) (string, error) {
			if len(mem.data) == 0 {
				return "长期记忆为空", nil
			}
			// Sort the output: Go iterates maps in random order, so without sorting the
			// key order changes from run to run (same reason as toolDefs: stable output
			// is what makes a run reproducible and comparable)
			keys := make([]string, 0, len(mem.data))
			for k := range mem.data {
				keys = append(keys, k)
			}
			sort.Strings(keys)

			var sb strings.Builder
			for _, k := range keys {
				fmt.Fprintf(&sb, "%s = %s\n", k, mem.data[k])
			}
			return strings.TrimSpace(sb.String()), nil
		},
	},
	"search_notes": {
		description: "检索本地运维笔记库，返回最相关的片段，用于回答需要查资料的问题",
		parameters: map[string]any{"type": "object", "properties": map[string]any{
			"query": map[string]any{"type": "string", "description": "要检索的问题或关键词"},
		}, "required": []string{"query"}},
		run: func(args json.RawMessage) (string, error) {
			var p struct {
				Query string `json:"query"`
			}
			if err := json.Unmarshal(args, &p); err != nil {
				return "", fmt.Errorf("参数解析失败：%v", err)
			}
			hits := retrieve.search(p.Query, 3)
			if len(hits) == 0 {
				return "笔记库中没有找到相关内容。", nil
			}
			var sb strings.Builder
			sb.WriteString("（以下是检索到的笔记片段，请严格只依据这些内容回答，不要编造笔记里没有的细节）\n")
			for i, h := range hits {
				fmt.Fprintf(&sb, "[片段%d %s%.2f] %s\n", i+1, retrieve.metric(), h.score, h.text)
			}
			return strings.TrimSpace(sb.String()), nil
		},
	},
}

// toolDefs is identical to Part 3's (sorted by name for a stable request body; see Part 3 §2).
func toolDefs() []Tool {
	names := make([]string, 0, len(tools))
	for name := range tools {
		names = append(names, name)
	}
	sort.Strings(names)

	out := make([]Tool, 0, len(names))
	for _, name := range names {
		t := tools[name]
		var td Tool
		td.Type = "function"
		td.Function.Name = name
		td.Function.Description = t.description
		td.Function.Parameters = t.parameters
		out = append(out, td)
	}
	return out
}

// ===========================================
// Sessions and the agent loop
// ===========================================

func main() {
	// Cross-process persistence demo: run `go run main.go` once to write the memory
	// (then the process exits), then `go run main.go recall` to run session B alone —
	// nothing is shared but memory.json.
	if len(os.Args) > 1 && os.Args[1] == "recall" {
		fmt.Println("===== 会话 B（新进程）：换一个全新会话（历史不共享），问它记得我吗 =====")
		runSessionB()
		return
	}

	fmt.Println("===== 会话 A：让 Agent 记住一些事实 =====")
	runSessionA()
	fmt.Println()
	fmt.Println("===== 会话 B：换一个全新会话（历史不共享），问它记得我吗 =====")
	runSessionB()
	fmt.Println()
	fmt.Println("===== 会话 C：查笔记回答问题（轻量 RAG）=====")
	runSessionC()
}

func runSessionA() {
	runAgent([]Message{{Role: "user", Content: "请调用两次 remember 工具：第一次 key=用户:名字、value=小明；第二次 key=用户:语言、value=Go。只做这两次调用，不要做其他事。"}})
}

func runSessionB() {
	runAgent([]Message{{Role: "user", Content: "这是一个全新的会话。请调用 list_memory 一次把长期记忆全部列出来，然后直接根据列出的内容回答我：你记得关于我的什么？列完就不用再调用其他工具了。"}})
}

func runSessionC() {
	runAgent([]Message{{Role: "user", Content: "问题：Intel Arc 显卡怎么在 Ollama 里启用 GPU 加速？请先 search_notes 检索运维笔记，再严格基于检索到的内容回答。"}})
}

func runAgent(messages []Message) {
	const maxRounds = 8 // fuse, same as Parts 3/4: stop the model from looping forever
	for round := 1; round <= maxRounds; round++ {
		resp, err := chat(messages)
		if err != nil {
			fmt.Println("❌ 请求失败:", err)
			return
		}
		if len(resp.Choices) == 0 {
			fmt.Println("❌ 无响应")
			return
		}
		msg := resp.Choices[0].Message
		finish := resp.Choices[0].FinishReason
		messages = append(messages, msg)
		// Same as Parts 2 and 3: go by tool_calls being non-empty (finish_reason is
		// only a cross-check — some upstreams still report "stop" while calling tools)
		if len(msg.ToolCalls) == 0 {
			fmt.Println("🤖", msg.Content)
			return
		}
		if finish != "tool_calls" {
			fmt.Printf("   ⚠️ finish_reason=%q 但本轮有 tool_calls，按 tool_calls 继续执行\n", finish)
		}
		fmt.Printf("🔧 第 %d 轮模型请求 %d 个工具\n", round, len(msg.ToolCalls))
		for _, tc := range msg.ToolCalls {
			t, ok := tools[tc.Function.Name]
			result := ""
			if !ok {
				result = fmt.Sprintf("未知工具 %q（模型幻觉了工具名）", tc.Function.Name)
			} else {
				result, err = t.run(json.RawMessage(tc.Function.Arguments))
				if err != nil {
					result = fmt.Sprintf("工具执行出错：%v。请修正参数后重试，或放弃这一步。", err)
				}
			}
			if result == "" {
				result = "(工具无输出)" // as in Part 3: a role="tool" message must carry content (Part 2 §3)
			}
			fmt.Printf("   - %s(%s)\n     => %s\n", tc.Function.Name, tc.Function.Arguments, result)
			messages = append(messages, Message{Role: "tool", ToolCallID: tc.ID, Content: result})
		}
	}
}

// ===========================================
// HTTP calls (chat and embed)
// ===========================================

var httpClient = &http.Client{Timeout: 5 * time.Minute} // request timeout: a stuck upstream cannot hang us forever

func chat(messages []Message) (ChatResponse, error) {
	reqBody := ChatRequest{
		Model:       "llama3.1:8b",
		Messages:    messages,
		Tools:       toolDefs(),
		Stream:      false,
		Temperature: 0, // greedy decoding keeps demo output stable; a chat scenario can go back to 0.7
	}
	jsonData, _ := json.Marshal(reqBody)
	resp, err := httpClient.Post("http://localhost:11434/v1/chat/completions", "application/json", bytes.NewBuffer(jsonData))
	if err != nil {
		return ChatResponse{}, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(resp.Body)
		return ChatResponse{}, fmt.Errorf("upstream %d: %s", resp.StatusCode, string(b))
	}
	body, _ := io.ReadAll(resp.Body)
	var result ChatResponse
	if err := json.Unmarshal(body, &result); err != nil {
		return ChatResponse{}, fmt.Errorf("解析失败: %s", string(body))
	}
	return result, nil
}

func embedTexts(texts []string) ([][]float64, error) {
	payload, _ := json.Marshal(map[string]any{"model": "nomic-embed-text", "input": texts})
	resp, err := httpClient.Post("http://localhost:11434/api/embed", "application/json", bytes.NewBuffer(payload))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("embed %d: %s", resp.StatusCode, string(b))
	}
	var out struct {
		Embeddings [][]float64 `json:"embeddings"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return nil, err
	}
	return out.Embeddings, nil
}