Skip to content

Commit fb6a7d3

Browse files
hackertronhackertron
authored
feat: add persistent memory with SQLite, embeddings, and hybrid retrieval (#10)
* feat: add persistent memory with SQLite, embeddings, and hybrid retrieval Implement Step 5 of the architecture — persistent memory backed by SQLite (pure Go, no CGO) with OpenAI embeddings and hybrid vector+FTS retrieval using reciprocal rank fusion. Replace the context budget stub with real rolling summarization. New packages/files: - internal/memory/ — SQLite DB layer, OpenAI embedding backend, memory store (MemoryRetrieval), session store, vector/FTS retrieval with RRF - internal/tool/memory_search.go — memory_search tool (ReadOnly) - internal/tool/memory_save.go — memory_save tool (SideEffecting) Modified: - RegisterBuiltins now accepts optional MemoryRetrieval for memory tools - AgentRuntime gains SetMemory() for persistent conversation history and real summarization when context budget is exceeded - Session gains CompactWithSummary() for context window compaction - CLI wires memory DB, embedder, session store, and tools
1 parent a010aa4 commit fb6a7d3

17 files changed

Lines changed: 1994 additions & 24 deletions

‎cmd/yantra/main.go‎

Lines changed: 50 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,13 @@ package main
33
import (
44
"context"
55
"fmt"
6+
"log/slog"
67
"os"
78
"os/signal"
89
"path/filepath"
910
"syscall"
1011

12+
"github.com/hackertron/Yantra/internal/memory"
1113
"github.com/hackertron/Yantra/internal/provider"
1214
"github.com/hackertron/Yantra/internal/runtime"
1315
"github.com/hackertron/Yantra/internal/tool"
@@ -208,16 +210,61 @@ func runAgent(ctx context.Context, prompt, systemPrompt, workspace string) error
208210

209211
policy := tool.NewWorkspacePolicy(cfg.Tools.Shell)
210212
reg := tool.NewRegistry(policy)
211-
if err := tool.RegisterBuiltins(reg, cfg.Tools); err != nil {
212-
return fmt.Errorf("registering tools: %w", err)
213-
}
214213

215214
absWorkspace, err := filepath.Abs(workspace)
216215
if err != nil {
217216
return fmt.Errorf("resolving workspace: %w", err)
218217
}
219218

219+
// Set up memory if enabled.
220+
var mem types.MemoryRetrieval
221+
var memDB *memory.DB
222+
var sessionID string
223+
224+
if cfg.Memory.Enabled {
225+
dbPath := cfg.Memory.DBPath
226+
if dbPath == "" {
227+
dbPath = ".yantra/memory.db"
228+
}
229+
if !filepath.IsAbs(dbPath) {
230+
dbPath = filepath.Join(absWorkspace, dbPath)
231+
}
232+
233+
memDB, err = memory.OpenDB(dbPath)
234+
if err != nil {
235+
slog.Warn("failed to open memory DB, continuing without memory", "error", err)
236+
} else {
237+
embedder, err := memory.NewEmbeddingBackend(cfg.Memory)
238+
if err != nil {
239+
slog.Warn("failed to create embedding backend, continuing without embeddings", "error", err)
240+
}
241+
242+
store := memory.NewStore(memDB, embedder, cfg.Memory.Retrieval)
243+
mem = store
244+
245+
// Create a session for this run.
246+
sessionStore := memory.NewSessionStore(memDB)
247+
sess, err := sessionStore.Create(ctx, "cli-run")
248+
if err != nil {
249+
slog.Warn("failed to create session", "error", err)
250+
} else {
251+
sessionID = sess.ID
252+
}
253+
}
254+
}
255+
256+
if memDB != nil {
257+
defer memDB.Close()
258+
}
259+
260+
if err := tool.RegisterBuiltins(reg, cfg.Tools, mem); err != nil {
261+
return fmt.Errorf("registering tools: %w", err)
262+
}
263+
220264
rt := runtime.New(p, reg, cfg.Runtime, absWorkspace)
265+
if mem != nil {
266+
rt.SetMemory(mem, sessionID)
267+
}
221268

222269
progress := make(chan types.ProgressEvent, 32)
223270
go func() {

‎go.mod‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,37 +12,47 @@ require (
1212
github.com/openai/openai-go/v3 v3.24.0
1313
github.com/spf13/cobra v1.10.2
1414
google.golang.org/genai v1.48.0
15+
modernc.org/sqlite v1.46.1
1516
)
1617

1718
require (
1819
cloud.google.com/go v0.116.0 // indirect
1920
cloud.google.com/go/auth v0.9.3 // indirect
2021
cloud.google.com/go/compute/metadata v0.5.0 // indirect
22+
github.com/dustin/go-humanize v1.0.1 // indirect
2123
github.com/fatih/structs v1.1.0 // indirect
2224
github.com/fsnotify/fsnotify v1.9.0 // indirect
2325
github.com/go-viper/mapstructure/v2 v2.4.0 // indirect
2426
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
2527
github.com/google/go-cmp v0.6.0 // indirect
2628
github.com/google/s2a-go v0.1.8 // indirect
29+
github.com/google/uuid v1.6.0 // indirect
2730
github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect
2831
github.com/gorilla/websocket v1.5.3 // indirect
2932
github.com/inconshreveable/mousetrap v1.1.0 // indirect
3033
github.com/knadh/koanf/maps v0.1.2 // indirect
34+
github.com/mattn/go-isatty v0.0.20 // indirect
3135
github.com/mitchellh/copystructure v1.2.0 // indirect
3236
github.com/mitchellh/reflectwalk v1.0.2 // indirect
37+
github.com/ncruces/go-strftime v1.0.0 // indirect
3338
github.com/pelletier/go-toml v1.9.5 // indirect
39+
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
3440
github.com/spf13/pflag v1.0.9 // indirect
3541
github.com/tidwall/gjson v1.18.0 // indirect
3642
github.com/tidwall/match v1.1.1 // indirect
3743
github.com/tidwall/pretty v1.2.1 // indirect
3844
github.com/tidwall/sjson v1.2.5 // indirect
3945
go.opencensus.io v0.24.0 // indirect
4046
golang.org/x/crypto v0.40.0 // indirect
47+
golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 // indirect
4148
golang.org/x/net v0.41.0 // indirect
42-
golang.org/x/sync v0.16.0 // indirect
43-
golang.org/x/sys v0.34.0 // indirect
49+
golang.org/x/sync v0.17.0 // indirect
50+
golang.org/x/sys v0.37.0 // indirect
4451
golang.org/x/text v0.27.0 // indirect
4552
google.golang.org/genproto/googleapis/rpc v0.0.0-20240903143218-8af14fe29dc1 // indirect
4653
google.golang.org/grpc v1.66.2 // indirect
4754
google.golang.org/protobuf v1.34.2 // indirect
55+
modernc.org/libc v1.67.6 // indirect
56+
modernc.org/mathutil v1.7.1 // indirect
57+
modernc.org/memory v1.11.0 // indirect
4858
)

‎go.sum‎

Lines changed: 53 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
1717
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
1818
github.com/dnaeon/go-vcr v1.2.0 h1:zHCHvJYTMh1N7xnV7zf1m1GPBF9Ad0Jk/whtQ1663qI=
1919
github.com/dnaeon/go-vcr v1.2.0/go.mod h1:R4UdLID7HZT3taECzJs4YgbbH6PIGXB6W/sc5OLb6RQ=
20+
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
21+
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
2022
github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
2123
github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
2224
github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1mIlRU8Am5FuJP05cCM98=
@@ -49,13 +51,19 @@ github.com/google/go-cmp v0.5.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/
4951
github.com/google/go-cmp v0.5.3/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
5052
github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
5153
github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
54+
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs=
55+
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA=
5256
github.com/google/s2a-go v0.1.8 h1:zZDs9gcbt9ZPLV0ndSyQk6Kacx2g/X+SKYovpnz3SMM=
5357
github.com/google/s2a-go v0.1.8/go.mod h1:6iNWHTpQ+nfNRN5E00MSdfDwVesa8hhS32PhPO8deJA=
5458
github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
59+
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
60+
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
5561
github.com/googleapis/enterprise-certificate-proxy v0.3.4 h1:XYIDZApgAnrN1c855gTgghdIA6Stxb52D5RnLI1SLyw=
5662
github.com/googleapis/enterprise-certificate-proxy v0.3.4/go.mod h1:YKe7cfqYXjKGpGvmSg28/fFvhNzinZQm8DGnaburhGA=
5763
github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
5864
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
65+
github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
66+
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
5967
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
6068
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
6169
github.com/knadh/koanf/maps v0.1.2 h1:RBfmAW5CnZT+PJ1CVc1QSJKf4Xu9kxfQgYVQSu8hpbo=
@@ -70,17 +78,23 @@ github.com/knadh/koanf/providers/structs v1.0.0 h1:DznjB7NQykhqCar2LvNug3MuxEQsZ
7078
github.com/knadh/koanf/providers/structs v1.0.0/go.mod h1:kjo5TFtgpaZORlpoJqcbeLowM2cINodv8kX+oFAeQ1w=
7179
github.com/knadh/koanf/v2 v2.3.2 h1:Ee6tuzQYFwcZXQpc2MiVeC6qHMandf5SMUJJNoFp/c4=
7280
github.com/knadh/koanf/v2 v2.3.2/go.mod h1:gRb40VRAbd4iJMYYD5IxZ6hfuopFcXBpc9bbQpZwo28=
81+
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
82+
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
7383
github.com/mitchellh/copystructure v1.2.0 h1:vpKXTN4ewci03Vljg/q9QvCGUDttBOGBIa15WveJJGw=
7484
github.com/mitchellh/copystructure v1.2.0/go.mod h1:qLl+cE2AmVv+CoeAwDPye/v+N2HKCj9FbZEVFJRxO9s=
7585
github.com/mitchellh/reflectwalk v1.0.2 h1:G2LzWKi524PWgd3mLHV8Y5k7s6XUvT0Gef6zxSIeXaQ=
7686
github.com/mitchellh/reflectwalk v1.0.2/go.mod h1:mSTlrgnPZtwu0c4WaC2kGObEpuNDbx0jmZXqmk4esnw=
87+
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
88+
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
7789
github.com/openai/openai-go/v3 v3.24.0 h1:08x6GnYiB+AAejTo6yzPY8RkZMJQ8NpreiOyM5QfyYU=
7890
github.com/openai/openai-go/v3 v3.24.0/go.mod h1:cdufnVK14cWcT9qA1rRtrXx4FTRsgbDPW7Ia7SS5cZo=
7991
github.com/pelletier/go-toml v1.9.5 h1:4yBQzkHv+7BHq2PQUZF3Mx0IYxG7LsP222s7Agd3ve8=
8092
github.com/pelletier/go-toml v1.9.5/go.mod h1:u1nR/EPcESfeI/szUZKdtJ0xRNbUoANCkoOuaOx1Y+c=
8193
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
8294
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
8395
github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
96+
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
97+
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
8498
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
8599
github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU=
86100
github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4=
@@ -112,9 +126,13 @@ golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPh
112126
golang.org/x/crypto v0.40.0 h1:r4x+VvoG5Fm+eJcxMaY8CQM7Lb0l1lsmjGBQ6s8BfKM=
113127
golang.org/x/crypto v0.40.0/go.mod h1:Qr1vMER5WyS2dfPHAlsOj01wgLbsyWtFn/aY+5+ZdxY=
114128
golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
129+
golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 h1:mgKeJMpvi0yx/sU5GsxQ7p6s2wtOnGAHZWCHUM4KGzY=
130+
golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546/go.mod h1:j/pmGrbnkbPtQfxEe5D0VQhZC6qKbfKifgD0oM7sR70=
115131
golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
116132
golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU=
117133
golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
134+
golang.org/x/mod v0.29.0 h1:HV8lRxZC4l2cr3Zq1LvtOsi/ThTgWnUk/y64QSs8GwA=
135+
golang.org/x/mod v0.29.0/go.mod h1:NyhrlYXJ2H4eJiRy/WDBO6HMqZQ6q9nk4JzS3NuCK+w=
118136
golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
119137
golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
120138
golang.org/x/net v0.0.0-20190213061140-3a22650c66bd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
@@ -127,14 +145,15 @@ golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAG
127145
golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
128146
golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
129147
golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
130-
golang.org/x/sync v0.16.0 h1:ycBJEhp9p4vXvUZNszeOq0kGTPghopOL8q0fq3vstxw=
131-
golang.org/x/sync v0.16.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
148+
golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug=
149+
golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
132150
golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
133151
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
134152
golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
135153
golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
136-
golang.org/x/sys v0.34.0 h1:H5Y5sJ2L2JRdyv7ROF1he/lPdvFsd0mJHFw2ThKHxLA=
137-
golang.org/x/sys v0.34.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
154+
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
155+
golang.org/x/sys v0.37.0 h1:fdNQudmxPjkdUTPnLn5mdQv7Zwvbvpaxqs831goi9kQ=
156+
golang.org/x/sys v0.37.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
138157
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
139158
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
140159
golang.org/x/text v0.27.0 h1:4fGWRpyh641NLlecmyl4LOe6yDdfaYNrGb2zdfo4JV4=
@@ -144,6 +163,8 @@ golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGm
144163
golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY=
145164
golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
146165
golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q=
166+
golang.org/x/tools v0.38.0 h1:Hx2Xv8hISq8Lm16jvBZ2VQf+RLmbd7wVUsALibYI/IQ=
167+
golang.org/x/tools v0.38.0/go.mod h1:yEsQ/d/YK8cjh0L6rZlY8tgtlKiBNTL14pGDJPJpYQs=
147168
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
148169
google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM=
149170
google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
@@ -180,3 +201,31 @@ gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
180201
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
181202
honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
182203
honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
204+
modernc.org/cc/v4 v4.27.1 h1:9W30zRlYrefrDV2JE2O8VDtJ1yPGownxciz5rrbQZis=
205+
modernc.org/cc/v4 v4.27.1/go.mod h1:uVtb5OGqUKpoLWhqwNQo/8LwvoiEBLvZXIQ/SmO6mL0=
206+
modernc.org/ccgo/v4 v4.30.1 h1:4r4U1J6Fhj98NKfSjnPUN7Ze2c6MnAdL0hWw6+LrJpc=
207+
modernc.org/ccgo/v4 v4.30.1/go.mod h1:bIOeI1JL54Utlxn+LwrFyjCx2n2RDiYEaJVSrgdrRfM=
208+
modernc.org/fileutil v1.3.40 h1:ZGMswMNc9JOCrcrakF1HrvmergNLAmxOPjizirpfqBA=
209+
modernc.org/fileutil v1.3.40/go.mod h1:HxmghZSZVAz/LXcMNwZPA/DRrQZEVP9VX0V4LQGQFOc=
210+
modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI=
211+
modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito=
212+
modernc.org/gc/v3 v3.1.1 h1:k8T3gkXWY9sEiytKhcgyiZ2L0DTyCQ/nvX+LoCljoRE=
213+
modernc.org/gc/v3 v3.1.1/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY=
214+
modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks=
215+
modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI=
216+
modernc.org/libc v1.67.6 h1:eVOQvpModVLKOdT+LvBPjdQqfrZq+pC39BygcT+E7OI=
217+
modernc.org/libc v1.67.6/go.mod h1:JAhxUVlolfYDErnwiqaLvUqc8nfb2r6S6slAgZOnaiE=
218+
modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
219+
modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
220+
modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI=
221+
modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw=
222+
modernc.org/opt v0.1.4 h1:2kNGMRiUjrp4LcaPuLY2PzUfqM/w9N23quVwhKt5Qm8=
223+
modernc.org/opt v0.1.4/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns=
224+
modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w=
225+
modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE=
226+
modernc.org/sqlite v1.46.1 h1:eFJ2ShBLIEnUWlLy12raN0Z1plqmFX9Qe3rjQTKt6sU=
227+
modernc.org/sqlite v1.46.1/go.mod h1:CzbrU2lSB1DKUusvwGz7rqEKIq+NUd8GWuBBZDs9/nA=
228+
modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0=
229+
modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A=
230+
modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y=
231+
modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM=

‎internal/memory/embedding.go‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
package memory
2+
3+
import (
4+
"fmt"
5+
"os"
6+
7+
"github.com/hackertron/Yantra/internal/types"
8+
)
9+
10+
// NewEmbeddingBackend creates an EmbeddingBackend based on config.
11+
// Returns nil (not an error) if no embedding backend can be configured.
12+
func NewEmbeddingBackend(cfg types.MemoryConfig) (types.EmbeddingBackend, error) {
13+
switch cfg.EmbeddingBackend {
14+
case "openai", "":
15+
apiKey := os.Getenv("OPENAI_API_KEY")
16+
if apiKey == "" {
17+
return nil, nil // graceful: no embeddings available
18+
}
19+
model := cfg.Embedding.Model
20+
return NewOpenAIEmbedder(apiKey, model)
21+
default:
22+
return nil, fmt.Errorf("memory: unsupported embedding backend: %q", cfg.EmbeddingBackend)
23+
}
24+
}

‎internal/memory/embedding_openai.go‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
package memory
2+
3+
import (
4+
"context"
5+
"fmt"
6+
7+
"github.com/hackertron/Yantra/internal/types"
8+
"github.com/openai/openai-go/v3"
9+
"github.com/openai/openai-go/v3/option"
10+
)
11+
12+
// OpenAIEmbedder implements EmbeddingBackend using OpenAI's embedding API.
13+
type OpenAIEmbedder struct {
14+
client *openai.Client
15+
model string
16+
dims int
17+
}
18+
19+
// NewOpenAIEmbedder creates an OpenAI embedding backend.
20+
func NewOpenAIEmbedder(apiKey, model string) (*OpenAIEmbedder, error) {
21+
if apiKey == "" {
22+
return nil, fmt.Errorf("memory: OPENAI_API_KEY is required for OpenAI embeddings")
23+
}
24+
if model == "" {
25+
model = "text-embedding-3-small"
26+
}
27+
28+
client := openai.NewClient(option.WithAPIKey(apiKey))
29+
30+
// Dimension mapping for known models.
31+
dims := 1536
32+
switch model {
33+
case "text-embedding-3-small":
34+
dims = 1536
35+
case "text-embedding-3-large":
36+
dims = 3072
37+
case "text-embedding-ada-002":
38+
dims = 1536
39+
}
40+
41+
return &OpenAIEmbedder{
42+
client: &client,
43+
model: model,
44+
dims: dims,
45+
}, nil
46+
}
47+
48+
func (e *OpenAIEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
49+
resp, err := e.client.Embeddings.New(ctx, openai.EmbeddingNewParams{
50+
Model: openai.EmbeddingModel(e.model),
51+
Input: openai.EmbeddingNewParamsInputUnion{
52+
OfString: openai.String(text),
53+
},
54+
})
55+
if err != nil {
56+
return nil, &types.MemoryError{Op: "embed", Message: "openai embedding failed", Err: err}
57+
}
58+
59+
if len(resp.Data) == 0 {
60+
return nil, &types.MemoryError{Op: "embed", Message: "no embedding data returned"}
61+
}
62+
63+
// Convert float64 → float32.
64+
f64 := resp.Data[0].Embedding
65+
out := make([]float32, len(f64))
66+
for i, v := range f64 {
67+
out[i] = float32(v)
68+
}
69+
return out, nil
70+
}
71+
72+
func (e *OpenAIEmbedder) Dimensions() int {
73+
return e.dims
74+
}
75+
76+
var _ types.EmbeddingBackend = (*OpenAIEmbedder)(nil)

0 commit comments

Comments
 (0)