openspec-memory-server/embedding.go - master - Ao/1c_AI_Spec

openspec-memory-server/embedding.go
163 строки4 KB
Максимчук Александр Олегович
gjcktlybq
08 июл 2026, 15:20
08 июл 2026, 15:200b5eb7f
package main

import (
	"context"
	"fmt"
	"math"
	"strings"
	"time"

	surrealdb "github.com/surrealdb/surrealdb.go"
)

const РазмерВектора = 128

type Эмбеддинг struct {
	ID         string    `json:"id"`
	DialogueID string    `json:"dialogue_id"`
	Vector     []float32 `json:"vector"`
	Text       string    `json:"text"`
	ChunkType  string    `json:"chunk_type"`
	CreatedAt  time.Time `json:"created_at"`
}

type РезультатПоиска struct {
	ID         string  `json:"id"`
	DialogueID string  `json:"dialogue_id"`
	Text       string  `json:"text"`
	ChunkType  string  `json:"chunk_type"`
	Score      float32 `json:"score"`
}

type ЧанкДляЭмбеддинга struct {
	Text      string
	ChunkType string
}

func СгенерироватьВектор(текст string) []float32 {
	токены := strings.Fields(strings.ToLower(текст))
_tf := make(map[string]float64)
	for _, токен := range токены {
		_tf[токен]++
	}
	total := float64(len(токены))
	for k := range _tf {
		_tf[k] /= total
	}

	вектор := make([]float32, РазмерВектора)
	for i, токен := range токены {
		if i >= РазмерВектора {
			break
		}
		вектор[i] = float32(_tf[токен])
	}

	норма := float32(0)
	for _, v := range вектор {
		норма += v * v
	}
	норма = float32(math.Sqrt(float64(норма)))

	if норма > 0 {
		for i := range вектор {
			вектор[i] /= норма
		}
	}

	return вектор
}

func СохранитьЭмбеддинги(ctx context.Context, dialogueID string, чанки []ЧанкДляЭмбеддинга) error {
	for _, чанк := range чанки {
		вектор := СгенерироватьВектор(чанк.Text)

		query := `CREATE embeddings CONTENT {
			dialogue_id: $dialogue_id,
			vector: $vector,
			text: $text,
			chunk_type: $chunk_type,
			created_at: $created_at
		}`

		vars := map[string]any{
			"dialogue_id": dialogueID,
			"vector":      вектор,
			"text":        чанк.Text,
			"chunk_type":  чанк.ChunkType,
			"created_at":  time.Now(),
		}

		if _, err := surrealdb.Query[any](ctx, SurrealDB, query, vars); err != nil {
			return fmt.Errorf("failed to save embedding: %w", err)
		}
	}
	return nil
}

func ПоискПохожих(ctx context.Context, запрос string, limit int) ([]РезультатПоиска, error) {
	query := `SELECT id, change_name, language, task_type, keywords
		FROM dialogues
		WHERE string::contains(change_name, $query)
		OR language = $query
		OR task_type = $query
		LIMIT $limit`

	vars := map[string]any{
		"query": запрос,
		"limit": limit,
	}

	result, err := surrealdb.Query[any](ctx, SurrealDB, query, vars)
	if err != nil {
		return nil, fmt.Errorf("search query failed: %w", err)
	}

	if result == nil || len(*result) == 0 {
		return []РезультатПоиска{}, nil
	}

	var результаты []РезультатПоиска
	for _, r := range *result {
		if r.Error != nil {
			continue
		}
		if resultSlice, ok := r.Result.([]any); ok {
			for _, item := range resultSlice {
				if itemMap, ok := item.(map[string]any); ok {
					результаты = append(результаты, РезультатПоиска{
						ID:         fmt.Sprintf("%v", itemMap["id"]),
						DialogueID: fmt.Sprintf("%v", itemMap["id"]),
						Text:       fmt.Sprintf("%v", itemMap["change_name"]),
						ChunkType:  fmt.Sprintf("%v", itemMap["task_type"]),
						Score:      float32(0),
					})
				}
			}
		}
	}

	return результаты, nil
}

func ИзвлечьЧанкиИзДиалога(диалог Диалог) []ЧанкДляЭмбеддинга {
	var чанки []ЧанкДляЭмбеддинга

	for _, в := range диалог.Questions {
		чанки = append(чанки, ЧанкДляЭмбеддинга{Text: в, ChunkType: "question"})
	}

	for _, о := range диалог.Answers {
		чанки = append(чанки, ЧанкДляЭмбеддинга{Text: о, ChunkType: "answer"})
	}

	if диалог.SpecContent != nil && *диалог.SpecContent != "" {
		чанки = append(чанки, ЧанкДляЭмбеддинга{Text: *диалог.SpecContent, ChunkType: "spec"})
	}

	for _, к := range диалог.Keywords {
		чанки = append(чанки, ЧанкДляЭмбеддинга{Text: к, ChunkType: "keyword"})
	}

	return чанки
}