Size: 10.6 KB Modified: 29/06/2026 11:55 PM
using System.Collections.Concurrent;
using System.Text.RegularExpressions;

namespace KaiChat.Services;

public class SearchResult
{
    public string ConversationTitle { get; set; } = "";
    public string Snippet { get; set; } = "";
    public float Relevance { get; set; }
    public string Source { get; set; } = "";
    public string SourcePath { get; set; } = "";
    public string? ConversationId { get; set; }
    public int? CompactionNumber { get; set; }
    public string? PreviousArchivePath { get; set; }
    public string Trace { get; set; } = "";
}

public interface ISearchService
{
    Task<List<SearchResult>> SearchAsync(string query, int maxResults = 10);
    Task BuildIndexAsync();
    Task RebuildIndexAsync();
    int IndexedConversations { get; }
}

public class SearchService : ISearchService
{
    private readonly string _baseFolder;
    private readonly string _extractedDir;
    private readonly string _archiveDir;

    private readonly IConversationService _convService;
    private readonly ILogger<SearchService> _logger;
    private readonly ConcurrentDictionary<string, SearchDocument> _index = new();
    private int _indexedCount;

    private static readonly HashSet<string> Stopwords = new(StringComparer.OrdinalIgnoreCase)
    {
        "the", "a", "an", "is", "it", "and", "or", "but", "in", "on", "at",
        "to", "for", "of", "with", "by", "from", "as", "was", "are", "were",
        "be", "been", "being", "have", "has", "had", "do", "does", "did",
        "will", "would", "could", "should", "may", "might", "shall", "can",
        "this", "that", "these", "those", "i", "you", "he", "she", "it",
        "we", "they", "me", "him", "her", "us", "them", "my", "your",
        "his", "its", "our", "their", "not", "no", "nor", "so", "if",
        "then", "than", "too", "very", "just", "about", "up", "out"
    };

    public SearchService(
        IConfiguration config,
        IConversationService convService,
        ILogger<SearchService> logger)
    {
        _baseFolder = config.GetValue<string>("KaiChat:BaseFolder") ?? ".";
        _extractedDir = Path.Combine(_baseFolder, "KaiChat", "ExtractedChats");
        _archiveDir = Path.Combine(_baseFolder, "KaiChat", "Archive");
        _convService = convService;
        _logger = logger;
    }

    public int IndexedConversations => _indexedCount;

    public Task RebuildIndexAsync() => BuildIndexAsync();

    public async Task BuildIndexAsync()
    {
        _index.Clear();
        _indexedCount = 0;

        if (Directory.Exists(_extractedDir))
        {
            foreach (var file in Directory.GetFiles(_extractedDir, "*.md")
                         .Where(f => !Path.GetFileName(f).StartsWith("_")))
            {
                var content = await File.ReadAllTextAsync(file);
                var title = Path.GetFileNameWithoutExtension(file);
                title = Regex.Replace(title, @"^\d+_", "").Replace('_', ' ').Trim();
                AddDocument(new SearchDocument
                {
                    Key = $"past:{file}",
                    Title = title,
                    Content = $"{title}\n{content}",
                    Source = "past",
                    SourcePath = ToRelativePath(file)
                });
            }
        }

        if (Directory.Exists(_archiveDir))
        {
            foreach (var file in Directory.GetFiles(_archiveDir, "*.md"))
            {
                var content = await File.ReadAllTextAsync(file);
                var metadata = ReadArchiveMetadata(content);
                var title = metadata.TryGetValue("conversation_title", out var archiveTitle) && !string.IsNullOrWhiteSpace(archiveTitle)
                    ? archiveTitle
                    : Path.GetFileNameWithoutExtension(file);
                var compactionNumber = metadata.TryGetValue("compaction_number", out var compactionText) &&
                                       int.TryParse(compactionText, out var parsedCompaction)
                    ? parsedCompaction
                    : (int?)null;

                AddDocument(new SearchDocument
                {
                    Key = $"archive:{file}",
                    Title = compactionNumber.HasValue ? $"{title} (compaction #{compactionNumber})" : title,
                    Content = $"{title}\n{content}",
                    Source = "archive",
                    SourcePath = ToRelativePath(file),
                    ConversationId = metadata.GetValueOrDefault("conversation_id"),
                    CompactionNumber = compactionNumber,
                    PreviousArchivePath = metadata.GetValueOrDefault("previous_archive_path")
                });
            }
        }

        foreach (var conv in _convService.ListConversations())
        {
            var builder = new System.Text.StringBuilder();
            builder.AppendLine(conv.Title);
            if (!string.IsNullOrEmpty(conv.CompactedSummary))
                builder.AppendLine(conv.CompactedSummary);

            if (conv.CompactionArchives.Count > 0)
            {
                builder.AppendLine("Compaction archive chain:");
                foreach (var entry in conv.CompactionArchives.OrderBy(a => a.CompactionNumber))
                    builder.AppendLine($"Compaction #{entry.CompactionNumber}: {entry.ArchivePath}");
            }

            foreach (var msg in conv.Messages)
                builder.AppendLine($"[{msg.Role}] {msg.Content}");

            AddDocument(new SearchDocument
            {
                Key = $"current:{conv.Id}",
                Title = conv.Title,
                Content = builder.ToString(),
                Source = "current",
                SourcePath = _convService.GetConversationRelativePath(conv.Id),
                ConversationId = conv.Id
            });
        }

        _logger.LogInformation("Search index built: {Count} documents", _indexedCount);
    }

    public Task<List<SearchResult>> SearchAsync(string query, int maxResults = 10)
    {
        var results = new List<(float Score, SearchResult Result)>();
        var queryTerms = Tokenize(query);

        if (queryTerms.Count == 0)
            return Task.FromResult(new List<SearchResult>());

        foreach (var (_, doc) in _index)
        {
            var contentTerms = Tokenize(doc.Content);
            if (contentTerms.Count == 0) continue;

            float score = 0;
            foreach (var qt in queryTerms)
            {
                var count = contentTerms.Count(ct => ct == qt);
                if (count > 0)
                    score += 1f + (float)Math.Log(count + 1);
            }

            if (score <= 0) continue;

            if (queryTerms.Any(qt => doc.Title.Contains(qt, StringComparison.OrdinalIgnoreCase)))
                score *= 2f;

            results.Add((score, new SearchResult
            {
                ConversationTitle = doc.Title.Length > 60 ? doc.Title[..57] + "..." : doc.Title,
                Snippet = ExtractSnippet(doc.Content, queryTerms[0]),
                Relevance = score,
                Source = doc.Source,
                SourcePath = doc.SourcePath,
                ConversationId = doc.ConversationId,
                CompactionNumber = doc.CompactionNumber,
                PreviousArchivePath = doc.PreviousArchivePath,
                Trace = BuildTrace(doc)
            }));
        }

        var sorted = results
            .OrderByDescending(r => r.Score)
            .Take(maxResults)
            .Select(r => r.Result)
            .ToList();

        return Task.FromResult(sorted);
    }

    private void AddDocument(SearchDocument doc)
    {
        if (_index.TryAdd(doc.Key, doc))
            _indexedCount++;
    }

    private string ToRelativePath(string file)
    {
        return Path.GetRelativePath(_baseFolder, file).Replace('\\', '/');
    }

    private static string BuildTrace(SearchDocument doc)
    {
        return doc.Source switch
        {
            "archive" => $"Archive `{doc.SourcePath}` for conversation `{doc.ConversationId}`"
                + (doc.CompactionNumber.HasValue ? $", compaction #{doc.CompactionNumber}" : "")
                + (!string.IsNullOrWhiteSpace(doc.PreviousArchivePath) ? $", previous archive `{doc.PreviousArchivePath}`" : ", first archive in chain"),
            "current" => $"Current conversation `{doc.ConversationId}` at `{doc.SourcePath}`",
            "past" => $"Imported past chat at `{doc.SourcePath}`",
            _ => doc.SourcePath
        };
    }

    private static Dictionary<string, string> ReadArchiveMetadata(string content)
    {
        var metadata = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
        using var reader = new StringReader(content);
        if (reader.ReadLine()?.Trim() != "---")
            return metadata;

        string? line;
        while ((line = reader.ReadLine()) != null)
        {
            if (line.Trim() == "---")
                break;

            var idx = line.IndexOf(':');
            if (idx <= 0) continue;

            var key = line[..idx].Trim();
            var value = UnquoteFrontMatter(line[(idx + 1)..].Trim());
            metadata[key] = value;
        }

        return metadata;
    }

    private static string UnquoteFrontMatter(string value)
    {
        if (value.Length >= 2 && value[0] == '"' && value[^1] == '"')
            return value[1..^1].Replace("\\\"", "\"").Replace("\\\\", "\\");
        return value;
    }

    private static List<string> Tokenize(string text)
    {
        if (string.IsNullOrWhiteSpace(text)) return new();
        return Regex.Matches(text.ToLowerInvariant(), @"\b[a-z]{2,}\b")
            .Select(m => m.Value)
            .Where(t => !Stopwords.Contains(t))
            .ToList();
    }

    private static string ExtractSnippet(string content, string term)
    {
        var idx = content.IndexOf(term, StringComparison.OrdinalIgnoreCase);
        if (idx == -1) return content.Length > 200 ? content[..200] + "..." : content;

        var start = Math.Max(0, idx - 100);
        var end = Math.Min(content.Length, idx + 200);
        var snippet = content[start..end].ReplaceLineEndings(" ");

        if (start > 0) snippet = "..." + snippet;
        if (end < content.Length) snippet += "...";

        return snippet.Trim();
    }

    private sealed class SearchDocument
    {
        public string Key { get; init; } = "";
        public string Title { get; init; } = "";
        public string Content { get; init; } = "";
        public string Source { get; init; } = "";
        public string SourcePath { get; init; } = "";
        public string? ConversationId { get; init; }
        public int? CompactionNumber { get; init; }
        public string? PreviousArchivePath { get; init; }
    }
}
Offline