first commit
This commit is contained in:
@@ -0,0 +1,126 @@
|
||||
package ice
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"ai-agent/internal/memory"
|
||||
)
|
||||
|
||||
type Assembler struct {
|
||||
embedder *Embedder
|
||||
convStore *Store
|
||||
memStore *memory.Store
|
||||
budgetCfg BudgetConfig
|
||||
sessionID string
|
||||
}
|
||||
|
||||
func (a *Assembler) Assemble(ctx context.Context, query string) (string, error) {
|
||||
budget := a.budgetCfg.Calculate(0)
|
||||
type convResult struct {
|
||||
chunks []ContextChunk
|
||||
err error
|
||||
}
|
||||
type memResult struct {
|
||||
chunks []ContextChunk
|
||||
}
|
||||
convCh := make(chan convResult, 1)
|
||||
memCh := make(chan memResult, 1)
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(2)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
chunks, err := a.retrieveConversations(ctx, query, budget.Conversation)
|
||||
convCh <- convResult{chunks: chunks, err: err}
|
||||
}()
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
chunks := a.retrieveMemories(query, budget.Memory)
|
||||
memCh <- memResult{chunks: chunks}
|
||||
}()
|
||||
wg.Wait()
|
||||
close(convCh)
|
||||
close(memCh)
|
||||
cr := <-convCh
|
||||
mr := <-memCh
|
||||
if cr.err != nil {
|
||||
return "", fmt.Errorf("conversation retrieval: %w", cr.err)
|
||||
}
|
||||
return formatContext(cr.chunks, mr.chunks), nil
|
||||
}
|
||||
|
||||
func (a *Assembler) retrieveConversations(ctx context.Context, query string, tokenBudget int) ([]ContextChunk, error) {
|
||||
if tokenBudget <= 0 {
|
||||
return nil, nil
|
||||
}
|
||||
queryEmb, err := a.embedder.Embed(ctx, query)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
results := a.convStore.Search(queryEmb, a.sessionID, 20)
|
||||
var chunks []ContextChunk
|
||||
usedTokens := 0
|
||||
for _, r := range results {
|
||||
tokens := estimateTokens(r.Entry.Content)
|
||||
if usedTokens+tokens > tokenBudget {
|
||||
continue
|
||||
}
|
||||
chunks = append(chunks, ContextChunk{
|
||||
Source: SourceConversation,
|
||||
Content: r.Entry.Content,
|
||||
Score: r.Score,
|
||||
Tokens: tokens,
|
||||
})
|
||||
usedTokens += tokens
|
||||
}
|
||||
return chunks, nil
|
||||
}
|
||||
|
||||
func (a *Assembler) retrieveMemories(query string, tokenBudget int) []ContextChunk {
|
||||
if a.memStore == nil || tokenBudget <= 0 {
|
||||
return nil
|
||||
}
|
||||
memories := a.memStore.Recall(query, 10)
|
||||
var chunks []ContextChunk
|
||||
usedTokens := 0
|
||||
for _, m := range memories {
|
||||
tokens := estimateTokens(m.Content)
|
||||
if usedTokens+tokens > tokenBudget {
|
||||
continue
|
||||
}
|
||||
content := m.Content
|
||||
if len(m.Tags) > 0 {
|
||||
content += " [" + strings.Join(m.Tags, ", ") + "]"
|
||||
}
|
||||
chunks = append(chunks, ContextChunk{
|
||||
Source: SourceMemory,
|
||||
Content: content,
|
||||
Tokens: tokens,
|
||||
})
|
||||
usedTokens += tokens
|
||||
}
|
||||
return chunks
|
||||
}
|
||||
|
||||
func formatContext(convChunks, memChunks []ContextChunk) string {
|
||||
var sb strings.Builder
|
||||
if len(convChunks) > 0 {
|
||||
sb.WriteString("\n## Relevant Past Conversations\n\n")
|
||||
for _, c := range convChunks {
|
||||
sb.WriteString("- ")
|
||||
sb.WriteString(c.Content)
|
||||
sb.WriteString("\n")
|
||||
}
|
||||
}
|
||||
if len(memChunks) > 0 {
|
||||
sb.WriteString("\n## Remembered Facts\n\n")
|
||||
for _, c := range memChunks {
|
||||
sb.WriteString("- ")
|
||||
sb.WriteString(c.Content)
|
||||
sb.WriteString("\n")
|
||||
}
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
package ice
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestFormatContext(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
convChunks []ContextChunk
|
||||
memChunks []ContextChunk
|
||||
wantConv bool // should contain "Relevant Past Conversations"
|
||||
wantMem bool // should contain "Remembered Facts"
|
||||
wantEmpty bool
|
||||
}{
|
||||
{
|
||||
name: "both conversation and memory chunks",
|
||||
convChunks: []ContextChunk{
|
||||
{Source: SourceConversation, Content: "past chat about Go"},
|
||||
},
|
||||
memChunks: []ContextChunk{
|
||||
{Source: SourceMemory, Content: "user prefers dark mode"},
|
||||
},
|
||||
wantConv: true,
|
||||
wantMem: true,
|
||||
},
|
||||
{
|
||||
name: "conversations only",
|
||||
convChunks: []ContextChunk{
|
||||
{Source: SourceConversation, Content: "previous discussion"},
|
||||
},
|
||||
memChunks: nil,
|
||||
wantConv: true,
|
||||
wantMem: false,
|
||||
},
|
||||
{
|
||||
name: "memories only",
|
||||
convChunks: nil,
|
||||
memChunks: []ContextChunk{
|
||||
{Source: SourceMemory, Content: "user name is Alice"},
|
||||
},
|
||||
wantConv: false,
|
||||
wantMem: true,
|
||||
},
|
||||
{
|
||||
name: "both empty",
|
||||
convChunks: nil,
|
||||
memChunks: nil,
|
||||
wantEmpty: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := formatContext(tt.convChunks, tt.memChunks)
|
||||
|
||||
if tt.wantEmpty {
|
||||
if got != "" {
|
||||
t.Errorf("expected empty string, got %q", got)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
hasConv := strings.Contains(got, "Relevant Past Conversations")
|
||||
hasMem := strings.Contains(got, "Remembered Facts")
|
||||
|
||||
if hasConv != tt.wantConv {
|
||||
t.Errorf("has conversations section = %v, want %v", hasConv, tt.wantConv)
|
||||
}
|
||||
if hasMem != tt.wantMem {
|
||||
t.Errorf("has memories section = %v, want %v", hasMem, tt.wantMem)
|
||||
}
|
||||
|
||||
// Verify content is present in output.
|
||||
for _, c := range tt.convChunks {
|
||||
if !strings.Contains(got, c.Content) {
|
||||
t.Errorf("output missing conversation content %q", c.Content)
|
||||
}
|
||||
}
|
||||
for _, c := range tt.memChunks {
|
||||
if !strings.Contains(got, c.Content) {
|
||||
t.Errorf("output missing memory content %q", c.Content)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
package ice
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"ai-agent/internal/llm"
|
||||
"ai-agent/internal/memory"
|
||||
)
|
||||
|
||||
var autoMemorySystemPrompt = "Extract any important facts, user preferences, decisions, or action items from this exchange.\n" +
|
||||
"Output one item per line in the format: TYPE: content\n" +
|
||||
"Where TYPE is one of: FACT, DECISION, PREFERENCE, TODO\n" +
|
||||
"If there is nothing worth remembering, output exactly: NONE"
|
||||
|
||||
var autoMemoryUserTemplate = "User: %s\nAssistant: %s"
|
||||
|
||||
type AutoMemory struct {
|
||||
client llm.Client
|
||||
memStore *memory.Store
|
||||
}
|
||||
|
||||
func (am *AutoMemory) Detect(ctx context.Context, userMsg, assistantMsg string) error {
|
||||
if am.memStore == nil {
|
||||
return nil
|
||||
}
|
||||
if len(userMsg) < 20 && len(assistantMsg) < 50 {
|
||||
return nil
|
||||
}
|
||||
prompt := fmt.Sprintf(autoMemoryUserTemplate, userMsg, assistantMsg)
|
||||
var response strings.Builder
|
||||
err := am.client.ChatStream(ctx, llm.ChatOptions{
|
||||
System: autoMemorySystemPrompt,
|
||||
Messages: []llm.Message{
|
||||
{Role: "user", Content: prompt},
|
||||
},
|
||||
}, func(chunk llm.StreamChunk) error {
|
||||
response.WriteString(chunk.Text)
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("auto-memory LLM call: %w", err)
|
||||
}
|
||||
return am.parseAndSave(response.String())
|
||||
}
|
||||
|
||||
func (am *AutoMemory) parseAndSave(response string) error {
|
||||
lines := strings.Split(strings.TrimSpace(response), "\n")
|
||||
for _, line := range lines {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" || strings.EqualFold(line, "NONE") {
|
||||
continue
|
||||
}
|
||||
parts := strings.SplitN(line, ": ", 2)
|
||||
if len(parts) != 2 {
|
||||
continue
|
||||
}
|
||||
typeName := strings.TrimSpace(parts[0])
|
||||
content := strings.TrimSpace(parts[1])
|
||||
if content == "" {
|
||||
continue
|
||||
}
|
||||
tag := strings.ToLower(typeName)
|
||||
switch tag {
|
||||
case "fact", "decision", "preference", "todo":
|
||||
// Valid type.
|
||||
default:
|
||||
continue
|
||||
}
|
||||
if _, err := am.memStore.Save(content, []string{tag, "auto"}); err != nil {
|
||||
return fmt.Errorf("save auto-memory: %w", err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
package ice
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"ai-agent/internal/memory"
|
||||
)
|
||||
|
||||
func TestParseAndSave(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
wantCount int
|
||||
wantTags [][]string
|
||||
}{
|
||||
{
|
||||
name: "valid FACT and DECISION lines",
|
||||
input: "FACT: user likes Go\nDECISION: use postgres",
|
||||
wantCount: 2,
|
||||
wantTags: [][]string{{"fact", "auto"}, {"decision", "auto"}},
|
||||
},
|
||||
{
|
||||
name: "NONE saves nothing",
|
||||
input: "NONE",
|
||||
wantCount: 0,
|
||||
},
|
||||
{
|
||||
name: "empty lines are skipped",
|
||||
input: "\n\n\n",
|
||||
wantCount: 0,
|
||||
},
|
||||
{
|
||||
name: "invalid type is skipped",
|
||||
input: "UNKNOWN: something",
|
||||
wantCount: 0,
|
||||
},
|
||||
{
|
||||
name: "missing colon format is skipped",
|
||||
input: "this has no colon",
|
||||
wantCount: 0,
|
||||
},
|
||||
{
|
||||
name: "PREFERENCE type",
|
||||
input: "PREFERENCE: dark mode",
|
||||
wantCount: 1,
|
||||
wantTags: [][]string{{"preference", "auto"}},
|
||||
},
|
||||
{
|
||||
name: "TODO type",
|
||||
input: "TODO: fix the bug",
|
||||
wantCount: 1,
|
||||
wantTags: [][]string{{"todo", "auto"}},
|
||||
},
|
||||
{
|
||||
name: "mixed valid and invalid",
|
||||
input: "FACT: real fact\nBAD: not valid\nTODO: real todo",
|
||||
wantCount: 2,
|
||||
wantTags: [][]string{{"fact", "auto"}, {"todo", "auto"}},
|
||||
},
|
||||
{
|
||||
name: "empty content after type is skipped",
|
||||
input: "FACT: ",
|
||||
wantCount: 0,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
memPath := filepath.Join(dir, "memories.json")
|
||||
ms := memory.NewStore(memPath)
|
||||
am := &AutoMemory{memStore: ms}
|
||||
err := am.parseAndSave(tt.input)
|
||||
if err != nil {
|
||||
t.Fatalf("parseAndSave returned error: %v", err)
|
||||
}
|
||||
if ms.Count() != tt.wantCount {
|
||||
t.Errorf("memory count = %d, want %d", ms.Count(), tt.wantCount)
|
||||
}
|
||||
if tt.wantTags != nil {
|
||||
recent := ms.Recent(tt.wantCount)
|
||||
for i, j := 0, len(recent)-1; i < j; i, j = i+1, j-1 {
|
||||
recent[i], recent[j] = recent[j], recent[i]
|
||||
}
|
||||
for i, wantTags := range tt.wantTags {
|
||||
if i >= len(recent) {
|
||||
t.Errorf("missing memory at index %d", i)
|
||||
continue
|
||||
}
|
||||
got := recent[i].Tags
|
||||
if len(got) != len(wantTags) {
|
||||
t.Errorf("memory[%d] tags = %v, want %v", i, got, wantTags)
|
||||
continue
|
||||
}
|
||||
for j := range wantTags {
|
||||
if got[j] != wantTags[j] {
|
||||
t.Errorf("memory[%d] tag[%d] = %q, want %q", i, j, got[j], wantTags[j])
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
package ice
|
||||
|
||||
// BudgetConfig controls how the context window is divided among sources.
|
||||
type BudgetConfig struct {
|
||||
NumCtx int
|
||||
SystemReserve int // tokens reserved for system prompt
|
||||
RecentReserve int // tokens reserved for recent conversation
|
||||
ConversationPct float64 // fraction of remaining budget for past conversations
|
||||
MemoryPct float64 // fraction of remaining budget for memories
|
||||
CodePct float64 // fraction of remaining budget for code context
|
||||
}
|
||||
|
||||
// DefaultBudgetConfig returns sensible defaults for a given context window.
|
||||
func DefaultBudgetConfig(numCtx int) BudgetConfig {
|
||||
return BudgetConfig{
|
||||
NumCtx: numCtx,
|
||||
SystemReserve: 1500,
|
||||
RecentReserve: 2000,
|
||||
ConversationPct: 0.40,
|
||||
MemoryPct: 0.20,
|
||||
CodePct: 0.40,
|
||||
}
|
||||
}
|
||||
|
||||
// Calculate allocates token budgets given how many tokens the current prompt uses.
|
||||
func (bc BudgetConfig) Calculate(promptTokens int) Budget {
|
||||
// Use 75% of numCtx as total available.
|
||||
available := int(float64(bc.NumCtx) * 0.75)
|
||||
available -= bc.SystemReserve
|
||||
available -= bc.RecentReserve
|
||||
available -= promptTokens
|
||||
|
||||
if available < 0 {
|
||||
available = 0
|
||||
}
|
||||
|
||||
return Budget{
|
||||
Total: available,
|
||||
System: bc.SystemReserve,
|
||||
Recent: bc.RecentReserve,
|
||||
Conversation: int(float64(available) * bc.ConversationPct),
|
||||
Memory: int(float64(available) * bc.MemoryPct),
|
||||
Code: int(float64(available) * bc.CodePct),
|
||||
}
|
||||
}
|
||||
|
||||
// estimateTokens returns a rough token count for a string (chars / 4).
|
||||
func estimateTokens(s string) int {
|
||||
n := len(s) / 4
|
||||
if n == 0 && len(s) > 0 {
|
||||
n = 1
|
||||
}
|
||||
return n
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
package ice
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestBudgetConfig_Calculate(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
cfg BudgetConfig
|
||||
promptTokens int
|
||||
wantTotal int
|
||||
wantConv int
|
||||
wantMemory int
|
||||
wantCode int
|
||||
}{
|
||||
{
|
||||
name: "normal allocation",
|
||||
cfg: BudgetConfig{
|
||||
NumCtx: 8192,
|
||||
SystemReserve: 1500,
|
||||
RecentReserve: 2000,
|
||||
ConversationPct: 0.40,
|
||||
MemoryPct: 0.20,
|
||||
CodePct: 0.40,
|
||||
},
|
||||
promptTokens: 500,
|
||||
// available = int(8192*0.75) - 1500 - 2000 - 500 = 6144 - 4000 = 2144
|
||||
wantTotal: 2144,
|
||||
wantConv: 857, // int(2144 * 0.40) = 857
|
||||
wantMemory: 428, // int(2144 * 0.20) = 428
|
||||
wantCode: 857, // int(2144 * 0.40) = 857
|
||||
},
|
||||
{
|
||||
name: "large prompt clamps to zero",
|
||||
cfg: BudgetConfig{
|
||||
NumCtx: 8192,
|
||||
SystemReserve: 1500,
|
||||
RecentReserve: 2000,
|
||||
ConversationPct: 0.40,
|
||||
MemoryPct: 0.20,
|
||||
CodePct: 0.40,
|
||||
},
|
||||
promptTokens: 99999,
|
||||
wantTotal: 0,
|
||||
wantConv: 0,
|
||||
wantMemory: 0,
|
||||
wantCode: 0,
|
||||
},
|
||||
{
|
||||
name: "exact boundary available is zero",
|
||||
cfg: BudgetConfig{
|
||||
NumCtx: 8192,
|
||||
SystemReserve: 1500,
|
||||
RecentReserve: 2000,
|
||||
ConversationPct: 0.40,
|
||||
MemoryPct: 0.20,
|
||||
CodePct: 0.40,
|
||||
},
|
||||
// int(8192*0.75) - 1500 - 2000 = 2644
|
||||
promptTokens: 2644,
|
||||
wantTotal: 0,
|
||||
wantConv: 0,
|
||||
wantMemory: 0,
|
||||
wantCode: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
b := tt.cfg.Calculate(tt.promptTokens)
|
||||
if b.Total != tt.wantTotal {
|
||||
t.Errorf("Total = %d, want %d", b.Total, tt.wantTotal)
|
||||
}
|
||||
if b.Conversation != tt.wantConv {
|
||||
t.Errorf("Conversation = %d, want %d", b.Conversation, tt.wantConv)
|
||||
}
|
||||
if b.Memory != tt.wantMemory {
|
||||
t.Errorf("Memory = %d, want %d", b.Memory, tt.wantMemory)
|
||||
}
|
||||
if b.Code != tt.wantCode {
|
||||
t.Errorf("Code = %d, want %d", b.Code, tt.wantCode)
|
||||
}
|
||||
if b.System != tt.cfg.SystemReserve {
|
||||
t.Errorf("System = %d, want %d", b.System, tt.cfg.SystemReserve)
|
||||
}
|
||||
if b.Recent != tt.cfg.RecentReserve {
|
||||
t.Errorf("Recent = %d, want %d", b.Recent, tt.cfg.RecentReserve)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEstimateTokens(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
want int
|
||||
}{
|
||||
{
|
||||
name: "len/4 heuristic",
|
||||
input: "hello world",
|
||||
want: 2, // 11/4 = 2
|
||||
},
|
||||
{
|
||||
name: "single char clamps to 1",
|
||||
input: "a",
|
||||
want: 1, // 1/4 = 0, clamp to 1
|
||||
},
|
||||
{
|
||||
name: "empty string",
|
||||
input: "",
|
||||
want: 0,
|
||||
},
|
||||
{
|
||||
name: "exactly 4 chars",
|
||||
input: "abcd",
|
||||
want: 1, // 4/4 = 1
|
||||
},
|
||||
{
|
||||
name: "three chars clamps to 1",
|
||||
input: "abc",
|
||||
want: 1, // 3/4 = 0, clamp to 1
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := estimateTokens(tt.input)
|
||||
if got != tt.want {
|
||||
t.Errorf("estimateTokens(%q) = %d, want %d", tt.input, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
package ice
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"ai-agent/internal/llm"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultEmbedModel = "nomic-embed-text"
|
||||
maxBatchSize = 32
|
||||
)
|
||||
|
||||
type Embedder struct {
|
||||
client llm.Client
|
||||
model string
|
||||
}
|
||||
|
||||
func NewEmbedder(client llm.Client, model string) *Embedder {
|
||||
if model == "" {
|
||||
model = defaultEmbedModel
|
||||
}
|
||||
return &Embedder{client: client, model: model}
|
||||
}
|
||||
|
||||
func (e *Embedder) Embed(ctx context.Context, text string) ([]float32, error) {
|
||||
vecs, err := e.EmbedBatch(ctx, []string{text})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(vecs) == 0 {
|
||||
return nil, fmt.Errorf("empty embedding response")
|
||||
}
|
||||
return vecs[0], nil
|
||||
}
|
||||
|
||||
func (e *Embedder) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error) {
|
||||
if len(texts) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
var all [][]float32
|
||||
for i := 0; i < len(texts); i += maxBatchSize {
|
||||
end := i + maxBatchSize
|
||||
if end > len(texts) {
|
||||
end = len(texts)
|
||||
}
|
||||
batch := texts[i:end]
|
||||
vecs, err := e.client.Embed(ctx, e.model, batch)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("embed batch [%d:%d]: %w", i, end, err)
|
||||
}
|
||||
all = append(all, vecs...)
|
||||
}
|
||||
return all, nil
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
package ice
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
"ai-agent/internal/llm"
|
||||
"ai-agent/internal/memory"
|
||||
)
|
||||
|
||||
type EngineConfig struct {
|
||||
EmbedModel string
|
||||
StorePath string
|
||||
NumCtx int
|
||||
}
|
||||
|
||||
type Engine struct {
|
||||
embedder *Embedder
|
||||
store *Store
|
||||
memStore *memory.Store
|
||||
budgetCfg BudgetConfig
|
||||
sessionID string
|
||||
turnIndex int
|
||||
autoMemory *AutoMemory
|
||||
}
|
||||
|
||||
func NewEngine(client llm.Client, memStore *memory.Store, cfg EngineConfig) (*Engine, error) {
|
||||
storePath := cfg.StorePath
|
||||
if storePath == "" {
|
||||
home, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("determine home dir: %w", err)
|
||||
}
|
||||
storePath = filepath.Join(home, ".config", "ai-agent", "conversations.json")
|
||||
}
|
||||
embedModel := cfg.EmbedModel
|
||||
if embedModel == "" {
|
||||
embedModel = defaultEmbedModel
|
||||
}
|
||||
sessionID := fmt.Sprintf("s_%d", time.Now().UnixNano())
|
||||
return &Engine{
|
||||
embedder: NewEmbedder(client, embedModel),
|
||||
store: NewStore(storePath),
|
||||
memStore: memStore,
|
||||
budgetCfg: DefaultBudgetConfig(cfg.NumCtx),
|
||||
sessionID: sessionID,
|
||||
autoMemory: &AutoMemory{client: client, memStore: memStore},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (e *Engine) AssembleContext(ctx context.Context, query string) (string, error) {
|
||||
a := &Assembler{
|
||||
embedder: e.embedder,
|
||||
convStore: e.store,
|
||||
memStore: e.memStore,
|
||||
budgetCfg: e.budgetCfg,
|
||||
sessionID: e.sessionID,
|
||||
}
|
||||
return a.Assemble(ctx, query)
|
||||
}
|
||||
|
||||
func (e *Engine) IndexMessage(ctx context.Context, role, content string) error {
|
||||
if content == "" {
|
||||
return nil
|
||||
}
|
||||
text := content
|
||||
if len(text) > 2000 {
|
||||
text = text[:2000]
|
||||
}
|
||||
emb, err := e.embedder.Embed(ctx, text)
|
||||
if err != nil {
|
||||
return fmt.Errorf("embed message: %w", err)
|
||||
}
|
||||
e.turnIndex++
|
||||
_, err = e.store.Add(e.sessionID, role, text, emb, e.turnIndex)
|
||||
return err
|
||||
}
|
||||
|
||||
func (e *Engine) IndexSummary(ctx context.Context, summary string) error {
|
||||
if summary == "" {
|
||||
return nil
|
||||
}
|
||||
emb, err := e.embedder.Embed(ctx, summary)
|
||||
if err != nil {
|
||||
return fmt.Errorf("embed summary: %w", err)
|
||||
}
|
||||
_, err = e.store.Add(e.sessionID, "summary", summary, emb, e.turnIndex)
|
||||
return err
|
||||
}
|
||||
|
||||
func (e *Engine) DetectAutoMemory(ctx context.Context, userMsg, assistantMsg string) {
|
||||
if e.autoMemory == nil {
|
||||
return
|
||||
}
|
||||
go func() {
|
||||
_ = e.autoMemory.Detect(ctx, userMsg, assistantMsg)
|
||||
}()
|
||||
}
|
||||
|
||||
func (e *Engine) Flush() error {
|
||||
return e.store.Flush()
|
||||
}
|
||||
|
||||
func (e *Engine) Store() *Store {
|
||||
return e.store
|
||||
}
|
||||
|
||||
func (e *Engine) SessionID() string {
|
||||
return e.sessionID
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package ice
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestEngineConfigDefaults(t *testing.T) {
|
||||
// Test embed model default
|
||||
embedModel := ""
|
||||
if embedModel == "" {
|
||||
embedModel = defaultEmbedModel
|
||||
}
|
||||
if embedModel != defaultEmbedModel {
|
||||
t.Errorf("embedModel = %q, want %q", embedModel, defaultEmbedModel)
|
||||
}
|
||||
|
||||
// Test custom embed model
|
||||
cfg := EngineConfig{
|
||||
EmbedModel: "custom-model",
|
||||
}
|
||||
if cfg.EmbedModel != "custom-model" {
|
||||
t.Errorf("EmbedModel = %q, want %q", cfg.EmbedModel, "custom-model")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBudgetConfigCalculate(t *testing.T) {
|
||||
cfg := DefaultBudgetConfig(16384)
|
||||
|
||||
budget := cfg.Calculate(100)
|
||||
// 16384 * 0.75 = 12288
|
||||
// 12288 - 1500 - 2000 - 100 = 8688
|
||||
if budget.Total != 8688 {
|
||||
t.Errorf("Total = %d, want %d", budget.Total, 8688)
|
||||
}
|
||||
if budget.System != 1500 {
|
||||
t.Errorf("System = %d, want %d", budget.System, 1500)
|
||||
}
|
||||
if budget.Recent != 2000 {
|
||||
t.Errorf("Recent = %d, want %d", budget.Recent, 2000)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBudgetConfigCalculateNegative(t *testing.T) {
|
||||
// With small context, should not panic and return zeros
|
||||
cfg := DefaultBudgetConfig(1000)
|
||||
budget := cfg.Calculate(500)
|
||||
|
||||
// 1000 * 0.75 = 750
|
||||
// 750 - 1500 - 2000 - 500 = -3250 -> clamped to 0
|
||||
if budget.Total != 0 {
|
||||
t.Errorf("Total should be 0 when budget is negative, got %d", budget.Total)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBudgetConfigPercentages(t *testing.T) {
|
||||
cfg := DefaultBudgetConfig(16384)
|
||||
budget := cfg.Calculate(100)
|
||||
|
||||
// Check percentages: ConversationPct=0.40, MemoryPct=0.20, CodePct=0.40
|
||||
// available = 12288 - 1500 - 2000 - 100 = 8688
|
||||
// Conversation = 8688 * 0.40 = 3475
|
||||
// Memory = 8688 * 0.20 = 1737
|
||||
// Code = 8688 * 0.40 = 3475
|
||||
if budget.Conversation != 3475 {
|
||||
t.Errorf("Conversation = %d, want %d", budget.Conversation, 3475)
|
||||
}
|
||||
if budget.Memory != 1737 {
|
||||
t.Errorf("Memory = %d, want %d", budget.Memory, 1737)
|
||||
}
|
||||
if budget.Code != 3475 {
|
||||
t.Errorf("Code = %d, want %d", budget.Code, 3475)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,167 @@
|
||||
package ice
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// timeNow is a variable for testing.
|
||||
var timeNow = time.Now
|
||||
|
||||
const minSimilarityThreshold = 0.3
|
||||
|
||||
// Store is a flat-file vector store for conversation history.
|
||||
// It holds all entries in memory and persists to a JSON file.
|
||||
type Store struct {
|
||||
mu sync.Mutex
|
||||
path string
|
||||
entries []ConversationEntry
|
||||
nextID int
|
||||
dirty bool
|
||||
}
|
||||
|
||||
// NewStore loads an existing store from path or creates an empty one.
|
||||
func NewStore(path string) *Store {
|
||||
s := &Store{path: path}
|
||||
s.load()
|
||||
return s
|
||||
}
|
||||
|
||||
// Add appends a new conversation entry and returns its ID.
|
||||
func (s *Store) Add(sessionID, role, content string, embedding []float32, turnIndex int) (int, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
s.nextID++
|
||||
entry := ConversationEntry{
|
||||
ID: s.nextID,
|
||||
SessionID: sessionID,
|
||||
Role: role,
|
||||
Content: content,
|
||||
Embedding: embedding,
|
||||
TurnIndex: turnIndex,
|
||||
}
|
||||
// Use a zero-value check to set CreatedAt (avoids importing time in every call site).
|
||||
entry.CreatedAt = timeNow()
|
||||
s.entries = append(s.entries, entry)
|
||||
s.dirty = true
|
||||
return s.nextID, nil
|
||||
}
|
||||
|
||||
// Search returns the top-K entries most similar to queryEmbedding.
|
||||
// Entries from excludeSession are skipped. Results are sorted by score descending.
|
||||
func (s *Store) Search(queryEmbedding []float32, excludeSession string, topK int) []ScoredEntry {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if len(queryEmbedding) == 0 || len(s.entries) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
var scored []ScoredEntry
|
||||
for _, e := range s.entries {
|
||||
if e.SessionID == excludeSession {
|
||||
continue
|
||||
}
|
||||
if len(e.Embedding) == 0 {
|
||||
continue
|
||||
}
|
||||
sim := cosineSimilarity(queryEmbedding, e.Embedding)
|
||||
if sim >= minSimilarityThreshold {
|
||||
scored = append(scored, ScoredEntry{Entry: e, Score: sim})
|
||||
}
|
||||
}
|
||||
|
||||
sort.Slice(scored, func(i, j int) bool {
|
||||
return scored[i].Score > scored[j].Score
|
||||
})
|
||||
|
||||
if len(scored) > topK {
|
||||
scored = scored[:topK]
|
||||
}
|
||||
return scored
|
||||
}
|
||||
|
||||
// Flush persists any pending changes to disk.
|
||||
func (s *Store) Flush() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if !s.dirty {
|
||||
return nil
|
||||
}
|
||||
return s.persist()
|
||||
}
|
||||
|
||||
// Count returns the total number of stored entries.
|
||||
func (s *Store) Count() int {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return len(s.entries)
|
||||
}
|
||||
|
||||
// load reads entries from the JSON file.
|
||||
func (s *Store) load() {
|
||||
data, err := os.ReadFile(s.path)
|
||||
if err != nil {
|
||||
return // File doesn't exist yet.
|
||||
}
|
||||
|
||||
var entries []ConversationEntry
|
||||
if err := json.Unmarshal(data, &entries); err != nil {
|
||||
return // Corrupt file, start empty.
|
||||
}
|
||||
|
||||
s.entries = entries
|
||||
for _, e := range s.entries {
|
||||
if e.ID > s.nextID {
|
||||
s.nextID = e.ID
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// persist writes all entries to the JSON file.
|
||||
func (s *Store) persist() error {
|
||||
dir := filepath.Dir(s.path)
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
return fmt.Errorf("create ice store dir: %w", err)
|
||||
}
|
||||
|
||||
data, err := json.Marshal(s.entries)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal ice store: %w", err)
|
||||
}
|
||||
|
||||
if err := os.WriteFile(s.path, data, 0o644); err != nil {
|
||||
return fmt.Errorf("write ice store: %w", err)
|
||||
}
|
||||
|
||||
s.dirty = false
|
||||
return nil
|
||||
}
|
||||
|
||||
// cosineSimilarity computes the cosine similarity between two vectors.
|
||||
func cosineSimilarity(a, b []float32) float32 {
|
||||
if len(a) != len(b) || len(a) == 0 {
|
||||
return 0
|
||||
}
|
||||
|
||||
var dot, normA, normB float64
|
||||
for i := range a {
|
||||
dot += float64(a[i]) * float64(b[i])
|
||||
normA += float64(a[i]) * float64(a[i])
|
||||
normB += float64(b[i]) * float64(b[i])
|
||||
}
|
||||
|
||||
denom := math.Sqrt(normA) * math.Sqrt(normB)
|
||||
if denom == 0 {
|
||||
return 0
|
||||
}
|
||||
return float32(dot / denom)
|
||||
}
|
||||
@@ -0,0 +1,232 @@
|
||||
package ice
|
||||
|
||||
import (
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestCosineSimilarity(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
a, b []float32
|
||||
want float32
|
||||
tol float32
|
||||
}{
|
||||
{
|
||||
name: "identical vectors",
|
||||
a: []float32{1, 2, 3},
|
||||
b: []float32{1, 2, 3},
|
||||
want: 1.0,
|
||||
tol: 1e-6,
|
||||
},
|
||||
{
|
||||
name: "orthogonal vectors",
|
||||
a: []float32{1, 0},
|
||||
b: []float32{0, 1},
|
||||
want: 0.0,
|
||||
tol: 1e-6,
|
||||
},
|
||||
{
|
||||
name: "opposite vectors",
|
||||
a: []float32{1, 0},
|
||||
b: []float32{-1, 0},
|
||||
want: -1.0,
|
||||
tol: 1e-6,
|
||||
},
|
||||
{
|
||||
name: "different lengths returns 0",
|
||||
a: []float32{1, 0},
|
||||
b: []float32{1, 0, 0},
|
||||
want: 0,
|
||||
tol: 0,
|
||||
},
|
||||
{
|
||||
name: "zero vector returns 0",
|
||||
a: []float32{0, 0},
|
||||
b: []float32{1, 1},
|
||||
want: 0,
|
||||
tol: 0,
|
||||
},
|
||||
{
|
||||
name: "known value with tolerance",
|
||||
a: []float32{1, 1},
|
||||
b: []float32{1, 0},
|
||||
// 1/(sqrt(2)*1) ≈ 0.7071
|
||||
want: float32(1.0 / math.Sqrt(2)),
|
||||
tol: 1e-4,
|
||||
},
|
||||
{
|
||||
name: "empty vectors",
|
||||
a: []float32{},
|
||||
b: []float32{},
|
||||
want: 0,
|
||||
tol: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := cosineSimilarity(tt.a, tt.b)
|
||||
diff := got - tt.want
|
||||
if diff < 0 {
|
||||
diff = -diff
|
||||
}
|
||||
if diff > tt.tol {
|
||||
t.Errorf("cosineSimilarity(%v, %v) = %f, want %f (±%f)",
|
||||
tt.a, tt.b, got, tt.want, tt.tol)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Add_And_Count(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "store.json")
|
||||
|
||||
s := NewStore(path)
|
||||
|
||||
if s.Count() != 0 {
|
||||
t.Fatalf("new store Count = %d, want 0", s.Count())
|
||||
}
|
||||
|
||||
id1, err := s.Add("sess1", "user", "hello", []float32{1, 0, 0}, 0)
|
||||
if err != nil {
|
||||
t.Fatalf("Add returned error: %v", err)
|
||||
}
|
||||
if id1 != 1 {
|
||||
t.Errorf("first Add returned id=%d, want 1", id1)
|
||||
}
|
||||
if s.Count() != 1 {
|
||||
t.Errorf("Count after first Add = %d, want 1", s.Count())
|
||||
}
|
||||
|
||||
id2, err := s.Add("sess1", "assistant", "world", []float32{0, 1, 0}, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("Add returned error: %v", err)
|
||||
}
|
||||
if id2 != 2 {
|
||||
t.Errorf("second Add returned id=%d, want 2", id2)
|
||||
}
|
||||
if s.Count() != 2 {
|
||||
t.Errorf("Count after second Add = %d, want 2", s.Count())
|
||||
}
|
||||
}
|
||||
|
||||
func TestStore_Search(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "store.json")
|
||||
s := NewStore(path)
|
||||
|
||||
// Add entries with known embeddings.
|
||||
s.Add("sess1", "user", "entry A", []float32{1, 0, 0}, 0)
|
||||
s.Add("sess1", "user", "entry B", []float32{0, 1, 0}, 1)
|
||||
s.Add("sess2", "user", "entry C", []float32{0.9, 0.1, 0}, 0)
|
||||
s.Add("sess2", "user", "entry D", []float32{0, 0, 1}, 1) // orthogonal to query
|
||||
|
||||
t.Run("similarity filtering and sorting", func(t *testing.T) {
|
||||
// Query similar to entries A and C, exclude no session.
|
||||
results := s.Search([]float32{1, 0, 0}, "", 10)
|
||||
if len(results) == 0 {
|
||||
t.Fatal("expected results, got 0")
|
||||
}
|
||||
// Entry A should be highest (identical to query).
|
||||
if results[0].Entry.Content != "entry A" {
|
||||
t.Errorf("top result = %q, want 'entry A'", results[0].Entry.Content)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("session exclusion", func(t *testing.T) {
|
||||
results := s.Search([]float32{1, 0, 0}, "sess1", 10)
|
||||
for _, r := range results {
|
||||
if r.Entry.SessionID == "sess1" {
|
||||
t.Errorf("excluded session sess1 should not appear in results")
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("min threshold 0.3", func(t *testing.T) {
|
||||
// Entry D: [0,0,1] is orthogonal to [1,0,0] → similarity 0.
|
||||
results := s.Search([]float32{1, 0, 0}, "", 10)
|
||||
for _, r := range results {
|
||||
if r.Score < minSimilarityThreshold {
|
||||
t.Errorf("result %q has score %f below threshold %f",
|
||||
r.Entry.Content, r.Score, minSimilarityThreshold)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("topK limit", func(t *testing.T) {
|
||||
results := s.Search([]float32{1, 0, 0}, "", 1)
|
||||
if len(results) > 1 {
|
||||
t.Errorf("topK=1 but got %d results", len(results))
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty store returns nil", func(t *testing.T) {
|
||||
emptyPath := filepath.Join(dir, "empty.json")
|
||||
empty := NewStore(emptyPath)
|
||||
results := empty.Search([]float32{1, 0}, "", 5)
|
||||
if results != nil {
|
||||
t.Errorf("empty store search should return nil, got %v", results)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty query embedding returns nil", func(t *testing.T) {
|
||||
results := s.Search([]float32{}, "", 5)
|
||||
if results != nil {
|
||||
t.Errorf("empty query should return nil, got %v", results)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestStore_Flush_Persistence(t *testing.T) {
|
||||
t.Run("round trip", func(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "store.json")
|
||||
|
||||
s1 := NewStore(path)
|
||||
s1.Add("sess1", "user", "hello", []float32{1, 0}, 0)
|
||||
s1.Add("sess1", "assistant", "world", []float32{0, 1}, 1)
|
||||
|
||||
if err := s1.Flush(); err != nil {
|
||||
t.Fatalf("Flush: %v", err)
|
||||
}
|
||||
|
||||
// Reload from same path.
|
||||
s2 := NewStore(path)
|
||||
if s2.Count() != 2 {
|
||||
t.Errorf("reloaded store Count = %d, want 2", s2.Count())
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("corrupt JSON recovery", func(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "store.json")
|
||||
|
||||
// Write corrupt JSON.
|
||||
os.WriteFile(path, []byte("not valid json{{{"), 0o644)
|
||||
|
||||
s := NewStore(path)
|
||||
if s.Count() != 0 {
|
||||
t.Errorf("corrupt store Count = %d, want 0", s.Count())
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("nextID restoration", func(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "store.json")
|
||||
|
||||
s1 := NewStore(path)
|
||||
s1.Add("sess1", "user", "first", []float32{1}, 0)
|
||||
s1.Add("sess1", "user", "second", []float32{1}, 1)
|
||||
s1.Flush()
|
||||
|
||||
s2 := NewStore(path)
|
||||
id, _ := s2.Add("sess1", "user", "third", []float32{1}, 2)
|
||||
if id != 3 {
|
||||
t.Errorf("continued id = %d, want 3", id)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package ice
|
||||
|
||||
import "time"
|
||||
|
||||
// SourceKind identifies where a context chunk came from.
|
||||
type SourceKind int
|
||||
|
||||
const (
|
||||
SourceConversation SourceKind = iota
|
||||
SourceMemory
|
||||
)
|
||||
|
||||
// ConversationEntry is a single stored message with its embedding.
|
||||
type ConversationEntry struct {
|
||||
ID int `json:"id"`
|
||||
SessionID string `json:"session_id"`
|
||||
Role string `json:"role"` // "user", "assistant", "summary"
|
||||
Content string `json:"content"`
|
||||
Embedding []float32 `json:"embedding"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
TurnIndex int `json:"turn_index"`
|
||||
}
|
||||
|
||||
// ScoredEntry pairs a conversation entry with its similarity score.
|
||||
type ScoredEntry struct {
|
||||
Entry ConversationEntry
|
||||
Score float32
|
||||
}
|
||||
|
||||
// ContextChunk is a piece of assembled context ready for the prompt.
|
||||
type ContextChunk struct {
|
||||
Source SourceKind
|
||||
Content string
|
||||
Score float32
|
||||
Tokens int
|
||||
}
|
||||
|
||||
// Budget holds the token allocation for each context source.
|
||||
type Budget struct {
|
||||
Total int
|
||||
System int
|
||||
Conversation int
|
||||
Memory int
|
||||
Code int
|
||||
Recent int
|
||||
}
|
||||
Reference in New Issue
Block a user