gott/Services/VectorSearchService.cs

using System;
using System.Text;
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);

            CREATE TABLE IF NOT EXISTS memory_file_embeddings (
                id           INTEGER PRIMARY KEY AUTOINCREMENT,
                agent_id     TEXT NOT NULL,
                file_path    TEXT NOT NULL,
                chunk_index  INTEGER NOT NULL,
                text_preview TEXT,
                vector       BLOB NOT NULL,
                updated_at   TEXT NOT NULL
            );
            CREATE INDEX IF NOT EXISTS idx_memfile_agent ON memory_file_embeddings(agent_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);
    }

    /// <summary>
    /// Generiert und speichert ein Embedding für eine Memory-Datei (oder aktualisiert es).
    /// </summary>
    public async Task EmbedMemoryFileAsync(string agentId, string filePath, string text)
    {
        if (!_embedding.IsAvailable || string.IsNullOrWhiteSpace(text)) return;
        if (text.Length < 10) return;

        // Chunks erstellen: Memory-Dateien können lang sein → in 500-Zeichen-Chunks aufteilen
        var chunks = ChunkText(text, 500);

        using var conn = CreateConnection();

        // Alte Embeddings für diese Datei löschen (bei Update)
        await conn.ExecuteAsync(
            "DELETE FROM memory_file_embeddings WHERE agent_id = @agentId AND file_path = @filePath",
            new { agentId, filePath });

        for (int i = 0; i < chunks.Count; i++)
        {
            float[] vector;
            await _embeddingLock.WaitAsync();
            try { vector = _embedding.GenerateEmbedding(chunks[i]); }
            finally { _embeddingLock.Release(); }
            if (vector.Length == 0) continue;

            var preview = chunks[i].Length > 200 ? chunks[i].Substring(0, 200) : chunks[i];
            await conn.ExecuteAsync(@"
                INSERT INTO memory_file_embeddings
                    (agent_id, file_path, chunk_index, text_preview, vector, updated_at)
                VALUES
                    (@agentId, @filePath, @chunkIndex, @preview, @vector, @updatedAt)",
                new
                {
                    agentId,
                    filePath,
                    chunkIndex = i,
                    preview,
                    vector = EmbeddingService.ToBytes(vector),
                    updatedAt = DateTime.UtcNow.ToString("o")
                });
        }
    }

    /// <summary>
    /// Sucht über Memory-Dateien eines Agents.
    /// </summary>
    public async Task<List<MemorySearchResult>> SearchMemoryFilesAsync(string query, string agentId, int topN = 5, float minScore = 0.4f)
    {
        if (!_embedding.IsAvailable || string.IsNullOrEmpty(agentId)) return new List<MemorySearchResult>();

        float[] queryVec;
        await _embeddingLock.WaitAsync();
        try { queryVec = _embedding.GenerateEmbedding(query); }
        finally { _embeddingLock.Release(); }
        if (queryVec.Length == 0) return new List<MemorySearchResult>();

        using var conn = CreateConnection();
        var rows = await conn.QueryAsync<MemoryFileRow>(
            "SELECT agent_id, file_path, chunk_index, text_preview, vector, updated_at FROM memory_file_embeddings WHERE agent_id = @agentId",
            new { agentId });

        var results = new List<MemorySearchResult>();
        foreach (var row in rows)
        {
            var vec = EmbeddingService.FromBytes(row.Vector);
            var score = EmbeddingService.CosineSimilarity(queryVec, vec);
            if (score >= minScore)
            {
                results.Add(new MemorySearchResult
                {
                    AgentId = row.AgentId,
                    FilePath = row.FilePath,
                    ChunkIndex = row.ChunkIndex,
                    TextPreview = row.TextPreview,
                    Score = score,
                    UpdatedAt = row.UpdatedAt
                });
            }
        }

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

    static List<string> ChunkText(string text, int chunkSize)
    {
        var chunks = new List<string>();
        if (string.IsNullOrEmpty(text)) return chunks;

        // Versuche an Absätzen zu trennen
        var paragraphs = text.Split(new[] { "\n\n", "\r\n\r\n" }, StringSplitOptions.RemoveEmptyEntries);
        var current = new StringBuilder();

        foreach (var para in paragraphs)
        {
            if (current.Length + para.Length > chunkSize && current.Length > 0)
            {
                chunks.Add(current.ToString().Trim());
                current.Clear();
            }
            if (current.Length > 0) current.AppendLine();
            current.Append(para);
        }
        if (current.Length > 0)
            chunks.Add(current.ToString().Trim());

        return chunks;
    }

    public class MemoryFileRow
    {
        public string AgentId { get; set; }
        public string FilePath { get; set; }
        public int ChunkIndex { get; set; }
        public string TextPreview { get; set; }
        public byte[] Vector { get; set; }
        public string UpdatedAt { get; set; }
    }

    public class MemorySearchResult
    {
        public string AgentId { get; set; }
        public string FilePath { get; set; }
        public int ChunkIndex { get; set; }
        public string TextPreview { get; set; }
        public float Score { get; set; }
        public string UpdatedAt { get; set; }
    }

    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; }
    }
}