gott/Services/VectorSearchService.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Dapper;
using Microsoft.Data.Sqlite;

namespace Gott.Services;

/// <summary>
/// Verwaltet Message-Embeddings in SQLite und bietet semantische Suche.
/// </summary>
public class VectorSearchService
{
    readonly string _connectionString;
    readonly EmbeddingService _embedding;
    readonly SemaphoreSlim _embeddingLock = new(1, 1);

    public VectorSearchService(string dbPath, EmbeddingService embedding)
    {
        _connectionString = $"Data Source={dbPath}";
        _embedding = embedding;
        InitializeSchema();
    }

    SqliteConnection CreateConnection()
    {
        var conn = new SqliteConnection(_connectionString);
        conn.Open();
        return conn;
    }

    void InitializeSchema()
    {
        using var conn = CreateConnection();
        conn.Execute(@"
            CREATE TABLE IF NOT EXISTS message_embeddings (
                id           INTEGER PRIMARY KEY AUTOINCREMENT,
                message_id   INTEGER NOT NULL UNIQUE,
                chain_id     TEXT NOT NULL,
                session_id   TEXT NOT NULL,
                agent_id     TEXT,
                role         TEXT NOT NULL,
                text_preview TEXT,
                vector       BLOB NOT NULL,
                created_at   TEXT NOT NULL
            );
            CREATE INDEX IF NOT EXISTS idx_emb_session ON message_embeddings(session_id);
            CREATE INDEX IF NOT EXISTS idx_emb_agent ON message_embeddings(agent_id);
            CREATE INDEX IF NOT EXISTS idx_emb_message ON message_embeddings(message_id);
        ");
        Console.WriteLine("[VectorSearch] Schema initialized.");
    }

    /// <summary>
    /// Generiert und speichert ein Embedding für eine Message.
    /// Thread-safe durch SemaphoreSlim (ONNX InferenceSession ist nicht thread-safe).
    /// </summary>
    public async Task EmbedMessageAsync(long messageId, string chainId, string sessionId,
        string agentId, string role, string text)
    {
        if (!_embedding.IsAvailable || string.IsNullOrWhiteSpace(text)) return;
        if (text.Length < 10) return; // Zu kurze Texte sind nicht sinnvoll

        // Prüfe ob schon embedded
        using var conn = CreateConnection();
        var exists = await conn.ExecuteScalarAsync<int>(
            "SELECT COUNT(*) FROM message_embeddings WHERE message_id = @messageId",
            new { messageId });
        if (exists > 0) return;

        // Embedding generieren (thread-safe)
        float[] vector;
        await _embeddingLock.WaitAsync();
        try
        {
            vector = _embedding.GenerateEmbedding(text);
        }
        finally
        {
            _embeddingLock.Release();
        }

        if (vector.Length == 0) return;

        var preview = text.Length > 200 ? text.Substring(0, 200) : text;

        await conn.ExecuteAsync(@"
            INSERT OR IGNORE INTO message_embeddings
                (message_id, chain_id, session_id, agent_id, role, text_preview, vector, created_at)
            VALUES
                (@messageId, @chainId, @sessionId, @agentId, @role, @preview, @vector, @createdAt)",
            new
            {
                messageId,
                chainId,
                sessionId,
                agentId = agentId ?? "",
                role,
                preview,
                vector = EmbeddingService.ToBytes(vector),
                createdAt = DateTime.UtcNow.ToString("o")
            });
    }

    /// <summary>
    /// Semantische Suche: Findet die N ähnlichsten Messages zu einer Query.
    /// Durchsucht alle Sessions eines Agents (oder alle wenn agentId null).
    /// </summary>
    public async Task<List<SearchResult>> SearchAsync(string query, string agentId = null,
        int topN = 10, float minScore = 0.3f)
    {
        if (!_embedding.IsAvailable) return new List<SearchResult>();

        // Query embedden
        float[] queryVec;
        await _embeddingLock.WaitAsync();
        try
        {
            queryVec = _embedding.GenerateEmbedding(query);
        }
        finally
        {
            _embeddingLock.Release();
        }

        if (queryVec.Length == 0) return new List<SearchResult>();

        // Alle Embeddings laden und Cosine Similarity berechnen
        using var conn = CreateConnection();
        string sql;
        object param;

        if (!string.IsNullOrEmpty(agentId))
        {
            sql = "SELECT message_id, chain_id, session_id, agent_id, role, text_preview, vector, created_at FROM message_embeddings WHERE agent_id = @agentId";
            param = new { agentId };
        }
        else
        {
            sql = "SELECT message_id, chain_id, session_id, agent_id, role, text_preview, vector, created_at FROM message_embeddings";
            param = new { };
        }

        var rows = await conn.QueryAsync<EmbeddingRow>(sql, param);
        var results = new List<SearchResult>();

        foreach (var row in rows)
        {
            var vec = EmbeddingService.FromBytes(row.Vector);
            var score = EmbeddingService.CosineSimilarity(queryVec, vec);
            if (score >= minScore)
            {
                results.Add(new SearchResult
                {
                    MessageId = row.MessageId,
                    ChainId = row.ChainId,
                    SessionId = row.SessionId,
                    AgentId = row.AgentId,
                    Role = row.Role,
                    TextPreview = row.TextPreview,
                    Score = score,
                    CreatedAt = row.CreatedAt
                });
            }
        }

        return results.OrderByDescending(r => r.Score).Take(topN).ToList();
    }

    /// <summary>
    /// Gibt die Anzahl der gespeicherten Embeddings zurück.
    /// </summary>
    public int GetEmbeddingCount()
    {
        using var conn = CreateConnection();
        return conn.ExecuteScalar<int>("SELECT COUNT(*) FROM message_embeddings");
    }

    /// <summary>
    /// Gibt die IDs aller bereits eingebetteten Messages zurück.
    /// </summary>
    public async Task<HashSet<long>> GetEmbeddedMessageIdsAsync()
    {
        using var conn = CreateConnection();
        var ids = await conn.QueryAsync<long>("SELECT message_id FROM message_embeddings");
        return new HashSet<long>(ids);
    }

    public class EmbeddingRow
    {
        public long MessageId { get; set; }
        public string ChainId { get; set; }
        public string SessionId { get; set; }
        public string AgentId { get; set; }
        public string Role { get; set; }
        public string TextPreview { get; set; }
        public byte[] Vector { get; set; }
        public string CreatedAt { get; set; }
    }

    public class SearchResult
    {
        public long MessageId { get; set; }
        public string ChainId { get; set; }
        public string SessionId { get; set; }
        public string AgentId { get; set; }
        public string Role { get; set; }
        public string TextPreview { get; set; }
        public float Score { get; set; }
        public string CreatedAt { get; set; }
    }
}