PNArchives/Memoh
1
0
Fork 0
mirror of https://github.com/memohai/Memoh.git synced 2026-04-27 07:16:19 +09:00
Code Issues Packages Projects Releases Wiki Activity

feat: embedding router

Browse Source
This commit is contained in:
Ran
2026-01-26 05:10:53 +07:00
parent c332ce7749
commit 3ff0e2c4dd
22 changed files with 2572 additions and 392 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/memory/service.go
+308 -22
View File
@@ -5,23 +5,33 @@ import (
"crypto/md5"
"encoding/hex"
"fmt"
"math"
"sort"
"strings"
"time"
"github.com/google/uuid"
"github.com/memohai/memoh/internal/embeddings"
)
type Service struct {
llm *LLMClient
embedder Embedder
store *QdrantStore
llm *LLMClient
embedder embeddings.Embedder
store *QdrantStore
resolver *embeddings.Resolver
defaultTextModelID string
defaultMultimodalModelID string
}
func NewService(llm *LLMClient, embedder Embedder, store *QdrantStore) *Service {
func NewService(llm *LLMClient, embedder embeddings.Embedder, store *QdrantStore, resolver *embeddings.Resolver, defaultTextModelID, defaultMultimodalModelID string) *Service {
return &Service{
llm: llm,
embedder: embedder,
store: store,
llm: llm,
embedder: embedder,
store: store,
resolver: resolver,
defaultTextModelID: defaultTextModelID,
defaultMultimodalModelID: defaultMultimodalModelID,
}
}
@@ -122,27 +132,128 @@ func (s *Service) Search(ctx context.Context, req SearchRequest) (SearchResponse
return SearchResponse{}, fmt.Errorf("query is required")
}
filters := buildSearchFilters(req)
vector, err := s.embedder.Embed(ctx, req.Query)
if err != nil {
return SearchResponse{}, err
modality := ""
if raw, ok := filters["modality"].(string); ok {
modality = strings.ToLower(strings.TrimSpace(raw))
}
points, scores, err := s.store.Search(ctx, vector, req.Limit, filters)
if err != nil {
return SearchResponse{}, err
}
results := make([]MemoryItem, 0, len(points))
for idx, point := range points {
item := payloadToMemoryItem(point.ID, point.Payload)
if idx < len(scores) {
item.Score = scores[idx]
var (
vector []float32
store *QdrantStore
vectorName string
err error
)
if modality == embeddings.TypeMultimodal {
if s.resolver == nil {
return SearchResponse{}, fmt.Errorf("embeddings resolver not configured")
}
results = append(results, item)
result, err := s.resolver.Embed(ctx, embeddings.Request{
Type: embeddings.TypeMultimodal,
Input: embeddings.Input{
Text: req.Query,
},
})
if err != nil {
return SearchResponse{}, err
}
vector = result.Embedding
store = s.store
vectorName = s.vectorNameForMultimodal()
} else {
vector, err = s.embedder.Embed(ctx, req.Query)
if err != nil {
return SearchResponse{}, err
}
store = s.store
vectorName = s.vectorNameForText()
}
if len(req.Sources) == 0 {
points, scores, err := store.Search(ctx, vector, req.Limit, filters, vectorName)
if err != nil {
return SearchResponse{}, err
}
results := make([]MemoryItem, 0, len(points))
for idx, point := range points {
item := payloadToMemoryItem(point.ID, point.Payload)
if idx < len(scores) {
item.Score = scores[idx]
}
results = append(results, item)
}
return SearchResponse{Results: results}, nil
}
pointsBySource, scoresBySource, err := store.SearchBySources(ctx, vector, req.Limit, filters, req.Sources, vectorName)
if err != nil {
return SearchResponse{}, err
}
results := fuseByRankFusion(pointsBySource, scoresBySource)
return SearchResponse{Results: results}, nil
}
func (s *Service) EmbedUpsert(ctx context.Context, req EmbedUpsertRequest) (EmbedUpsertResponse, error) {
if s.resolver == nil {
return EmbedUpsertResponse{}, fmt.Errorf("embeddings resolver not configured")
}
if req.UserID == "" && req.AgentID == "" && req.RunID == "" {
return EmbedUpsertResponse{}, fmt.Errorf("user_id, agent_id or run_id is required")
}
req.Type = strings.TrimSpace(req.Type)
req.Provider = strings.TrimSpace(req.Provider)
req.Model = strings.TrimSpace(req.Model)
req.Input.Text = strings.TrimSpace(req.Input.Text)
req.Input.ImageURL = strings.TrimSpace(req.Input.ImageURL)
req.Input.VideoURL = strings.TrimSpace(req.Input.VideoURL)
result, err := s.resolver.Embed(ctx, embeddings.Request{
Type: req.Type,
Provider: req.Provider,
Model: req.Model,
Input: embeddings.Input{
Text: req.Input.Text,
ImageURL: req.Input.ImageURL,
VideoURL: req.Input.VideoURL,
},
})
if err != nil {
return EmbedUpsertResponse{}, err
}
if s.store == nil {
return EmbedUpsertResponse{}, fmt.Errorf("qdrant store not configured")
}
vectorName := ""
if s.store != nil && s.store.usesNamedVectors {
vectorName = result.Model
}
id := uuid.NewString()
filters := buildEmbedFilters(req)
payload := buildEmbeddingPayload(req, filters)
if metadata, ok := payload["metadata"].(map[string]interface{}); ok && result.Model != "" {
metadata["model_id"] = result.Model
}
if err := s.store.Upsert(ctx, []qdrantPoint{{
ID: id,
Vector: result.Embedding,
VectorName: vectorName,
Payload: payload,
}}); err != nil {
return EmbedUpsertResponse{}, err
}
item := payloadToMemoryItem(id, payload)
return EmbedUpsertResponse{
Item: item,
Provider: result.Provider,
Model: result.Model,
Dimensions: result.Dimensions,
}, nil
}
func (s *Service) Update(ctx context.Context, req UpdateRequest) (MemoryItem, error) {
if strings.TrimSpace(req.MemoryID) == "" {
return MemoryItem{}, fmt.Errorf("memory_id is required")
@@ -171,6 +282,7 @@ func (s *Service) Update(ctx context.Context, req UpdateRequest) (MemoryItem, er
if err := s.store.Upsert(ctx, []qdrantPoint{{
ID: req.MemoryID,
Vector: vector,
VectorName: s.vectorNameForText(),
Payload: payload,
}}); err != nil {
return MemoryItem{}, err
@@ -270,7 +382,7 @@ func (s *Service) collectCandidates(ctx context.Context, facts []string, filters
if err != nil {
return nil, err
}
points, _, err := s.store.Search(ctx, vector, 5, filters)
points, _, err := s.store.Search(ctx, vector, 5, filters, s.vectorNameForText())
if err != nil {
return nil, err
}
@@ -301,6 +413,7 @@ func (s *Service) applyAdd(ctx context.Context, text string, filters map[string]
if err := s.store.Upsert(ctx, []qdrantPoint{{
ID: id,
Vector: vector,
VectorName: s.vectorNameForText(),
Payload: payload,
}}); err != nil {
return MemoryItem{}, err
@@ -337,6 +450,7 @@ func (s *Service) applyUpdate(ctx context.Context, id, text string, filters map[
if err := s.store.Upsert(ctx, []qdrantPoint{{
ID: id,
Vector: vector,
VectorName: s.vectorNameForText(),
Payload: payload,
}}); err != nil {
return MemoryItem{}, err
@@ -403,6 +517,68 @@ func buildSearchFilters(req SearchRequest) map[string]interface{} {
return filters
}
func buildEmbedFilters(req EmbedUpsertRequest) map[string]interface{} {
filters := map[string]interface{}{}
for key, value := range req.Filters {
filters[key] = value
}
if req.UserID != "" {
filters["userId"] = req.UserID
}
if req.AgentID != "" {
filters["agentId"] = req.AgentID
}
if req.RunID != "" {
filters["runId"] = req.RunID
}
return filters
}
func buildEmbeddingPayload(req EmbedUpsertRequest, filters map[string]interface{}) map[string]interface{} {
text := req.Input.Text
payload := buildPayload(text, filters, req.Metadata, "")
payload["hash"] = hashEmbeddingInput(req.Input.Text, req.Input.ImageURL, req.Input.VideoURL)
if req.Source != "" {
payload["source"] = req.Source
}
modality := "text"
if req.Type != "" {
modality = strings.ToLower(req.Type)
}
payload["modality"] = modality
if payload["metadata"] == nil {
payload["metadata"] = map[string]interface{}{}
}
if metadata, ok := payload["metadata"].(map[string]interface{}); ok {
if req.Source != "" {
metadata["source"] = req.Source
}
metadata["modality"] = modality
if req.Input.ImageURL != "" {
metadata["image_url"] = req.Input.ImageURL
}
if req.Input.VideoURL != "" {
metadata["video_url"] = req.Input.VideoURL
}
}
return payload
}
func (s *Service) vectorNameForText() string {
if s.store == nil || !s.store.usesNamedVectors {
return ""
}
return strings.TrimSpace(s.defaultTextModelID)
}
func (s *Service) vectorNameForMultimodal() string {
if s.store == nil || !s.store.usesNamedVectors {
return ""
}
return strings.TrimSpace(s.defaultMultimodalModelID)
}
func buildPayload(text string, filters map[string]interface{}, metadata map[string]interface{}, createdAt string) map[string]interface{} {
if createdAt == "" {
createdAt = time.Now().UTC().Format(time.RFC3339)
@@ -450,6 +626,16 @@ func payloadToMemoryItem(id string, payload map[string]interface{}) MemoryItem {
}
if meta, ok := payload["metadata"].(map[string]interface{}); ok {
item.Metadata = meta
} else if payload["metadata"] == nil {
item.Metadata = map[string]interface{}{}
}
if item.Metadata != nil {
if source, ok := payload["source"].(string); ok && source != "" {
item.Metadata["source"] = source
}
if modality, ok := payload["modality"].(string); ok && modality != "" {
item.Metadata["modality"] = modality
}
}
return item
}
@@ -459,6 +645,16 @@ func hashMemory(text string) string {
return hex.EncodeToString(sum[:])
}
func hashEmbeddingInput(text, imageURL, videoURL string) string {
combined := strings.Join([]string{
strings.TrimSpace(text),
strings.TrimSpace(imageURL),
strings.TrimSpace(videoURL),
}, "|")
sum := md5.Sum([]byte(combined))
return hex.EncodeToString(sum[:])
}
func mergeMetadata(base interface{}, extra map[string]interface{}) map[string]interface{} {
merged := map[string]interface{}{}
if baseMap, ok := base.(map[string]interface{}); ok {
@@ -471,3 +667,93 @@ func mergeMetadata(base interface{}, extra map[string]interface{}) map[string]in
}
return merged
}
type rerankCandidate struct {
ID string
Payload map[string]interface{}
Score float64
Source string
Rank int
}
const (
fusionModeRRF = "rrf"
fusionModeCombMNZ = "combmnz"
fusionMode = fusionModeRRF
rrfK = 60.0
)
func fuseByRankFusion(pointsBySource map[string][]qdrantPoint, scoresBySource map[string][]float64) []MemoryItem {
candidates := map[string]*rerankCandidate{}
rrfScores := map[string]float64{}
combScores := map[string]float64{}
combCounts := map[string]int{}
for source, points := range pointsBySource {
scores := scoresBySource[source]
minScore := math.MaxFloat64
maxScore := -math.MaxFloat64
for idx, point := range points {
if idx >= len(scores) {
continue
}
score := scores[idx]
if score < minScore {
minScore = score
}
if score > maxScore {
maxScore = score
}
if _, ok := candidates[point.ID]; !ok {
candidates[point.ID] = &rerankCandidate{
ID: point.ID,
Payload: point.Payload,
}
}
}
if minScore == math.MaxFloat64 {
minScore = 0
}
if maxScore == -math.MaxFloat64 {
maxScore = minScore
}
for idx, point := range points {
if idx >= len(scores) {
continue
}
score := scores[idx]
rank := float64(idx + 1)
rrfScores[point.ID] += 1.0 / (rrfK + rank)
scoreNorm := normalizeScore(score, minScore, maxScore)
combScores[point.ID] += scoreNorm
combCounts[point.ID]++
}
}
items := make([]MemoryItem, 0, len(candidates))
for id, candidate := range candidates {
item := payloadToMemoryItem(candidate.ID, candidate.Payload)
switch fusionMode {
case fusionModeCombMNZ:
item.Score = combScores[id] * float64(combCounts[id])
default:
item.Score = rrfScores[id]
}
items = append(items, item)
}
sort.Slice(items, func(i, j int) bool {
return items[i].Score > items[j].Score
})
return items
}
func normalizeScore(score, minScore, maxScore float64) float64 {
if maxScore <= minScore {
return 1
}
return (score - minScore) / (maxScore - minScore)
}