Skip to content
This repository was archived by the owner on Jun 3, 2026. It is now read-only.

Commit fa570b6

Browse files
ishaanxguptaishaanxgupta
committed
add modularity
1 parent bbfb665 commit fa570b6

16 files changed

Lines changed: 1038 additions & 822 deletions

‎xmem-go/cmd/xmem/bootstrap.go‎

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,189 @@
1+
package main
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"log/slog"
7+
8+
"github.com/xortexai/xmem-go/internal/agents"
9+
"github.com/xortexai/xmem-go/internal/config"
10+
"github.com/xortexai/xmem-go/internal/database"
11+
"github.com/xortexai/xmem-go/internal/graph"
12+
"github.com/xortexai/xmem-go/internal/jobs"
13+
"github.com/xortexai/xmem-go/internal/models"
14+
"github.com/xortexai/xmem-go/internal/pipelines"
15+
"github.com/xortexai/xmem-go/internal/storage"
16+
"github.com/xortexai/xmem-go/internal/weaver"
17+
)
18+
19+
type runtimeDeps struct {
20+
ingest *pipelines.IngestPipeline
21+
retrieval *pipelines.RetrievalPipeline
22+
keyStore database.APIKeyStore
23+
jobStore jobs.Store
24+
}
25+
26+
func buildRuntime(ctx context.Context, settings config.Settings, logger *slog.Logger) (runtimeDeps, error) {
27+
embedder, err := buildEmbedder(settings, logger)
28+
if err != nil {
29+
return runtimeDeps{}, err
30+
}
31+
32+
vectorStore, snippetStore, err := buildVectorStores(ctx, settings, embedder, logger)
33+
if err != nil {
34+
return runtimeDeps{}, err
35+
}
36+
37+
temporalStore, err := buildTemporalStore(ctx, settings, logger)
38+
if err != nil {
39+
return runtimeDeps{}, err
40+
}
41+
42+
model := models.NewRegistry(settings)
43+
ingest, retrieval := buildPipelines(model, vectorStore, snippetStore, temporalStore, embedder)
44+
45+
keyStore, jobStore, err := buildAppStores(ctx, settings, logger)
46+
if err != nil {
47+
return runtimeDeps{}, err
48+
}
49+
50+
return runtimeDeps{
51+
ingest: ingest,
52+
retrieval: retrieval,
53+
keyStore: keyStore,
54+
jobStore: jobStore,
55+
}, nil
56+
}
57+
58+
func buildEmbedder(settings config.Settings, logger *slog.Logger) (storage.Embedder, error) {
59+
fallback := storage.HashEmbedder{Dimension: settings.PineconeDimension}
60+
if settings.EmbeddingProvider != "openai" {
61+
return fallback, nil
62+
}
63+
64+
openAIEmbedder, err := storage.NewOpenAIEmbedder(settings)
65+
if err == nil {
66+
logger.Info("using OpenAI embedder", "model", settings.OpenAIEmbeddingModel, "dimension", settings.PineconeDimension)
67+
return openAIEmbedder, nil
68+
}
69+
if production(settings) {
70+
return nil, fmt.Errorf("openai embedder initialization failed: %w", err)
71+
}
72+
logger.Warn("openai embedder unavailable, using hash embedder", "error", err)
73+
return fallback, nil
74+
}
75+
76+
func buildVectorStores(ctx context.Context, settings config.Settings, embedder storage.Embedder, logger *slog.Logger) (storage.VectorStore, storage.VectorStore, error) {
77+
vectorStore := storage.VectorStore(storage.NewMemoryVectorStore())
78+
snippetStore := storage.VectorStore(storage.NewMemoryVectorStore())
79+
if settings.VectorStoreProvider != "pinecone" {
80+
return vectorStore, snippetStore, nil
81+
}
82+
83+
pineconeStore, err := storage.NewPineconeVectorStore(ctx, settings, embedder, settings.PineconeNamespace)
84+
if err != nil {
85+
if production(settings) {
86+
return nil, nil, fmt.Errorf("pinecone vector store initialization failed: %w", err)
87+
}
88+
logger.Warn("pinecone unavailable, using memory vector store", "error", err)
89+
} else {
90+
vectorStore = pineconeStore
91+
logger.Info("using Pinecone vector store", "namespace", settings.PineconeNamespace)
92+
}
93+
94+
snippetNamespace := settings.PineconeNamespace + "-snippets"
95+
pineconeSnippets, err := storage.NewPineconeVectorStore(ctx, settings, embedder, snippetNamespace)
96+
if err != nil {
97+
if production(settings) {
98+
return nil, nil, fmt.Errorf("pinecone snippet store initialization failed: %w", err)
99+
}
100+
logger.Warn("pinecone snippet store unavailable, using memory vector store", "error", err)
101+
} else {
102+
snippetStore = pineconeSnippets
103+
logger.Info("using Pinecone snippet vector store", "namespace", snippetNamespace)
104+
}
105+
106+
return vectorStore, snippetStore, nil
107+
}
108+
109+
func buildTemporalStore(ctx context.Context, settings config.Settings, logger *slog.Logger) (graph.TemporalStore, error) {
110+
fallback := graph.NewMemoryTemporalStore()
111+
if settings.Neo4jPassword == "" {
112+
return fallback, nil
113+
}
114+
115+
neoStore, err := graph.NewNeo4jTemporalStore(ctx, settings)
116+
if err == nil {
117+
logger.Info("using Neo4j temporal store")
118+
return neoStore, nil
119+
}
120+
if production(settings) {
121+
return nil, fmt.Errorf("neo4j initialization failed: %w", err)
122+
}
123+
logger.Warn("neo4j unavailable, using memory temporal store", "error", err)
124+
return fallback, nil
125+
}
126+
127+
func buildPipelines(model models.ChatModel, vectorStore storage.VectorStore, snippetStore storage.VectorStore, temporalStore graph.TemporalStore, embedder storage.Embedder) (*pipelines.IngestPipeline, *pipelines.RetrievalPipeline) {
128+
w := &weaver.Weaver{
129+
VectorStore: vectorStore,
130+
SnippetVectorStore: snippetStore,
131+
Embedder: embedder,
132+
TemporalStore: temporalStore,
133+
}
134+
135+
ingest := &pipelines.IngestPipeline{
136+
ModelName: model.Name(),
137+
Weaver: w,
138+
Classifier: agents.ClassifierAgent{Model: model},
139+
Profiler: agents.ProfilerAgent{Model: model},
140+
Temporal: agents.TemporalAgent{Model: model},
141+
Summarizer: agents.SummarizerAgent{Model: model},
142+
Image: agents.ImageAgent{Model: model},
143+
Snippet: agents.SnippetAgent{Model: model},
144+
Judge: agents.JudgeAgent{Model: model, VectorStore: vectorStore, TopK: 3},
145+
}
146+
retrieval := &pipelines.RetrievalPipeline{
147+
Model: model,
148+
VectorStore: vectorStore,
149+
SnippetStore: snippetStore,
150+
TemporalStore: temporalStore,
151+
}
152+
return ingest, retrieval
153+
}
154+
155+
func buildAppStores(ctx context.Context, settings config.Settings, logger *slog.Logger) (database.APIKeyStore, jobs.Store, error) {
156+
keyStore := database.APIKeyStore(database.NewMemoryAPIKeyStore())
157+
jobStore := jobs.Store(jobs.NewMemoryStore())
158+
if settings.AppStoreProvider != "mongo" {
159+
return keyStore, jobStore, nil
160+
}
161+
162+
mongoStore, err := database.NewMongoAPIKeyStore(ctx, settings)
163+
if err != nil {
164+
if production(settings) {
165+
return nil, nil, fmt.Errorf("mongodb api key store initialization failed: %w", err)
166+
}
167+
logger.Warn("mongodb unavailable, using memory API key store", "error", err)
168+
} else {
169+
keyStore = mongoStore
170+
logger.Info("using MongoDB API key store")
171+
}
172+
173+
durableStore, err := database.NewMongoDurableJobStore(ctx, settings)
174+
if err != nil {
175+
if production(settings) {
176+
return nil, nil, fmt.Errorf("mongodb durable job store initialization failed: %w", err)
177+
}
178+
logger.Warn("mongodb durable job store unavailable, using memory job store", "error", err)
179+
} else {
180+
jobStore = durableStore
181+
logger.Info("using MongoDB durable job store")
182+
}
183+
184+
return keyStore, jobStore, nil
185+
}
186+
187+
func production(settings config.Settings) bool {
188+
return settings.Environment == "production"
189+
}

‎xmem-go/cmd/xmem/main.go‎

Lines changed: 5 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -9,16 +9,8 @@ import (
99
"syscall"
1010
"time"
1111

12-
"github.com/xortexai/xmem-go/internal/agents"
1312
"github.com/xortexai/xmem-go/internal/api"
1413
"github.com/xortexai/xmem-go/internal/config"
15-
"github.com/xortexai/xmem-go/internal/database"
16-
"github.com/xortexai/xmem-go/internal/graph"
17-
"github.com/xortexai/xmem-go/internal/jobs"
18-
"github.com/xortexai/xmem-go/internal/models"
19-
"github.com/xortexai/xmem-go/internal/pipelines"
20-
"github.com/xortexai/xmem-go/internal/storage"
21-
"github.com/xortexai/xmem-go/internal/weaver"
2214
)
2315

2416
func main() {
@@ -29,98 +21,12 @@ func main() {
2921
os.Exit(1)
3022
}
3123

32-
var vectorStore storage.VectorStore = storage.NewMemoryVectorStore()
33-
var snippetStore storage.VectorStore = storage.NewMemoryVectorStore()
34-
var temporalStore graph.TemporalStore = graph.NewMemoryTemporalStore()
35-
var embedder storage.Embedder = storage.HashEmbedder{Dimension: settings.PineconeDimension}
36-
if settings.EmbeddingProvider == "openai" {
37-
if openAIEmbedder, err := storage.NewOpenAIEmbedder(settings); err == nil {
38-
embedder = openAIEmbedder
39-
logger.Info("using OpenAI embedder", "model", settings.OpenAIEmbeddingModel, "dimension", settings.PineconeDimension)
40-
} else if settings.Environment == "production" {
41-
logger.Error("openai embedder initialization failed", "error", err)
42-
os.Exit(1)
43-
} else {
44-
logger.Warn("openai embedder unavailable, using hash embedder", "error", err)
45-
}
46-
}
47-
model := models.NewRegistry(settings)
48-
49-
if settings.VectorStoreProvider == "pinecone" {
50-
if pineconeStore, err := storage.NewPineconeVectorStore(context.Background(), settings, embedder, settings.PineconeNamespace); err == nil {
51-
vectorStore = pineconeStore
52-
logger.Info("using Pinecone vector store", "namespace", settings.PineconeNamespace)
53-
} else if settings.Environment == "production" {
54-
logger.Error("pinecone initialization failed", "error", err)
55-
os.Exit(1)
56-
} else {
57-
logger.Warn("pinecone unavailable, using memory vector store", "error", err)
58-
}
59-
if pineconeSnippets, err := storage.NewPineconeVectorStore(context.Background(), settings, embedder, settings.PineconeNamespace+"-snippets"); err == nil {
60-
snippetStore = pineconeSnippets
61-
logger.Info("using Pinecone snippet vector store", "namespace", settings.PineconeNamespace+"-snippets")
62-
}
63-
}
64-
65-
if settings.Neo4jPassword != "" {
66-
if neoStore, err := graph.NewNeo4jTemporalStore(context.Background(), settings); err == nil {
67-
temporalStore = neoStore
68-
logger.Info("using Neo4j temporal store")
69-
} else if settings.Environment == "production" {
70-
logger.Error("neo4j initialization failed", "error", err)
71-
os.Exit(1)
72-
} else {
73-
logger.Warn("neo4j unavailable, using memory temporal store", "error", err)
74-
}
75-
}
76-
77-
w := &weaver.Weaver{
78-
VectorStore: vectorStore,
79-
SnippetVectorStore: snippetStore,
80-
Embedder: embedder,
81-
TemporalStore: temporalStore,
82-
}
83-
ingest := &pipelines.IngestPipeline{
84-
ModelName: model.Name(),
85-
Weaver: w,
86-
Classifier: agents.ClassifierAgent{Model: model},
87-
Profiler: agents.ProfilerAgent{Model: model},
88-
Temporal: agents.TemporalAgent{Model: model},
89-
Summarizer: agents.SummarizerAgent{Model: model},
90-
Image: agents.ImageAgent{Model: model},
91-
Snippet: agents.SnippetAgent{Model: model},
92-
Judge: agents.JudgeAgent{Model: model, VectorStore: vectorStore, TopK: 3},
93-
}
94-
retrieval := &pipelines.RetrievalPipeline{
95-
Model: model,
96-
VectorStore: vectorStore,
97-
SnippetStore: snippetStore,
98-
TemporalStore: temporalStore,
99-
}
100-
101-
var keyStore database.APIKeyStore = database.NewMemoryAPIKeyStore()
102-
var jobStore jobs.Store = jobs.NewMemoryStore()
103-
if settings.AppStoreProvider == "mongo" {
104-
if mongoStore, err := database.NewMongoAPIKeyStore(context.Background(), settings); err == nil {
105-
keyStore = mongoStore
106-
logger.Info("using MongoDB API key store")
107-
} else if settings.Environment == "production" {
108-
logger.Error("mongodb initialization failed", "error", err)
109-
os.Exit(1)
110-
} else {
111-
logger.Warn("mongodb unavailable, using memory API key store", "error", err)
112-
}
113-
if durableStore, err := database.NewMongoDurableJobStore(context.Background(), settings); err == nil {
114-
jobStore = durableStore
115-
logger.Info("using MongoDB durable job store")
116-
} else if settings.Environment == "production" {
117-
logger.Error("mongodb durable job initialization failed", "error", err)
118-
os.Exit(1)
119-
} else {
120-
logger.Warn("mongodb durable job store unavailable, using memory job store", "error", err)
121-
}
24+
runtime, err := buildRuntime(context.Background(), settings, logger)
25+
if err != nil {
26+
logger.Error("runtime initialization failed", "error", err)
27+
os.Exit(1)
12228
}
123-
server := api.NewServer(settings, logger, ingest, retrieval, keyStore, jobStore)
29+
server := api.NewServer(settings, logger, runtime.ingest, runtime.retrieval, runtime.keyStore, runtime.jobStore)
12430
httpServer := &http.Server{
12531
Addr: settings.Addr(),
12632
Handler: server.Handler(),

‎xmem-go/internal/api/auth.go‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
package api
2+
3+
import (
4+
"crypto/hmac"
5+
"crypto/sha256"
6+
"encoding/base64"
7+
"encoding/json"
8+
"net/http"
9+
"strings"
10+
11+
"github.com/xortexai/xmem-go/internal/database"
12+
)
13+
14+
func (s *Server) authenticate(r *http.Request) (User, int, string) {
15+
auth := r.Header.Get("Authorization")
16+
if !strings.HasPrefix(auth, "Bearer ") {
17+
return User{}, http.StatusUnauthorized, "Missing API key. Provide a Bearer token in the Authorization header."
18+
}
19+
token := strings.TrimSpace(strings.TrimPrefix(auth, "Bearer "))
20+
if token == "" {
21+
return User{}, http.StatusUnauthorized, "Missing API key. Provide a Bearer token in the Authorization header."
22+
}
23+
if !strings.HasPrefix(token, "xmem_") {
24+
if user, ok := s.validateJWT(token); ok {
25+
return user, http.StatusOK, ""
26+
}
27+
}
28+
if doc, ok := s.keys.ValidateAPIKey(token); ok {
29+
if userDoc, exists := s.keys.GetUserByID(doc.UserID); exists {
30+
return User{ID: userDoc.ID, Name: userDoc.Name, Email: userDoc.Email, Username: userDoc.Username, APIKey: map[string]any{"id": doc.ID, "scopes": doc.Scopes, "org_id": doc.OrgID, "project_id": doc.ProjectID}}, http.StatusOK, ""
31+
}
32+
}
33+
for _, key := range s.settings.APIKeys {
34+
if database.ConstantTimeEqual(token, key) {
35+
return User{ID: database.StaticUserID(token), Name: "Static Key User", Email: "static@xmem.ai"}, http.StatusOK, ""
36+
}
37+
}
38+
return User{}, http.StatusForbidden, "Invalid API key or token."
39+
}
40+
41+
func (s *Server) validateJWT(token string) (User, bool) {
42+
parts := strings.Split(token, ".")
43+
if len(parts) != 3 || s.settings.JWTAlgorithm != "HS256" {
44+
return User{}, false
45+
}
46+
signingInput := parts[0] + "." + parts[1]
47+
mac := hmac.New(sha256.New, []byte(s.settings.JWTSecretKey))
48+
_, _ = mac.Write([]byte(signingInput))
49+
expected := base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
50+
if !hmac.Equal([]byte(expected), []byte(parts[2])) {
51+
return User{}, false
52+
}
53+
payloadBytes, err := base64.RawURLEncoding.DecodeString(parts[1])
54+
if err != nil {
55+
return User{}, false
56+
}
57+
var payload map[string]any
58+
if err := json.Unmarshal(payloadBytes, &payload); err != nil {
59+
return User{}, false
60+
}
61+
if payload["type"] != "access" {
62+
return User{}, false
63+
}
64+
sub, _ := payload["sub"].(string)
65+
if sub == "" {
66+
return User{}, false
67+
}
68+
if userDoc, exists := s.keys.GetUserByID(sub); exists {
69+
return User{ID: userDoc.ID, Name: userDoc.Name, Email: userDoc.Email, Username: userDoc.Username}, true
70+
}
71+
return User{ID: sub, Name: sub, Email: ""}, true
72+
}
73+
74+
func effectiveUserID(user User) string {
75+
if user.Username != "" {
76+
return user.Username
77+
}
78+
if user.Name != "" {
79+
return user.Name
80+
}
81+
return user.ID
82+
}

0 commit comments

Comments
 (0)