Runtime/AgentCore/Rag/LightRagKnowledgeService.cs

using System;
using System.Threading.Tasks;
using LightRAG.Core;
using LightRAG.Core.Abstractions;
using LightRAG.Storage.FileBased;
using OnDeviceAgent.RagLlm;
namespace OnDeviceAgent.AgentCore
{

    public enum RagQueryMode
    {
        Naive,
        Local,
        Global,
        Hybrid,
        Mix,
    }

    public sealed class LightRagKnowledgeService : IDisposable
    {
        readonly string m_WorkingDir;
        readonly ILlmModel m_Llm;
        readonly ITextEmbedder m_Embedder;
        readonly IUnityMainThreadDispatcher m_Dispatcher;
        readonly LightRagConfig m_Config;
        readonly Action<string> m_Log;
        readonly RagQueryMode m_Mode;
        readonly int m_MixTopK;
        readonly float m_RelevanceThreshold;
        readonly IChatLlm m_Translator; // non-null when the LLM is a chat LLM (Ollama or on-device); used for KO->EN query translation

        LightRag m_Rag;
        NanoVectorDbStorage m_ChunksVdb;
        volatile bool m_Ready;

        public bool Ready => m_Ready;

        static QueryMode ToLightRag(RagQueryMode m)
        {
            switch (m)
            {
                case RagQueryMode.Local: return QueryMode.Local;
                case RagQueryMode.Global: return QueryMode.Global;
                case RagQueryMode.Hybrid: return QueryMode.Hybrid;
                case RagQueryMode.Mix: return QueryMode.Mix;
                default: return QueryMode.Naive;
            }
        }

        public LightRagKnowledgeService(
            string workingDir,
            ILlmModel llm,
            ITextEmbedder embedder,
            IUnityMainThreadDispatcher dispatcher,
            LightRagConfig config = null,
            Action<string> log = null,
            RagQueryMode queryMode = RagQueryMode.Naive,
            float relevanceThreshold = 0.78f,
            int mixTopK = 12)
        {
            m_WorkingDir = workingDir;
            m_Llm = llm;
            m_Embedder = embedder;
            m_Dispatcher = dispatcher;
            m_Config = config ?? new LightRagConfig();
            m_Log = log;
            m_Mode = queryMode;
            m_RelevanceThreshold = relevanceThreshold;
            m_MixTopK = mixTopK;
            m_Translator = llm as IChatLlm;
        }

        public LightRagKnowledgeService(
            string workingDir,
            string ollamaEndpoint,
            string ollamaModel,
            int numCtx,
            ITextEmbedder embedder,
            IUnityMainThreadDispatcher dispatcher,
            Action<string> log = null,
            RagQueryMode queryMode = RagQueryMode.Naive,
            float relevanceThreshold = 0.78f,
            int mixTopK = 12)
            : this(workingDir, new OllamaChatLlm(ollamaModel, ollamaEndpoint, numCtx), embedder, dispatcher,
                   null, log, queryMode, relevanceThreshold, mixTopK)
        {
        }

        public async Task InitializeAsync()
        {
            try
            {
                var embedding = m_Embedder.AsEmbeddingFunc(m_Dispatcher);
                // Inject a data-free tokenizer: the default TiktokenTokenizer (o200k_base) loads a reflection-only
                // data assembly that fails on Android/IL2CPP, which would abort RAG init. See CharApproxTokenizer.
                m_Rag = FileBasedLightRag.Create(m_WorkingDir, m_Llm, embedding, m_Config, tokenizer: new CharApproxTokenizer());
                await m_Rag.InitializeAsync().ConfigureAwait(false);

                m_ChunksVdb = new NanoVectorDbStorage(
                    m_WorkingDir, NameSpace.VectorStoreChunks, embedding, 0f);
                await m_ChunksVdb.InitializeAsync().ConfigureAwait(false);

                m_Ready = true;
                m_Log?.Invoke($"[RAG] LightRAG ready ({m_Mode} mode) at {m_WorkingDir}.");
            }
            catch (Exception e)
            {
                m_Log?.Invoke("[RAG] LightRAG init failed: " + e.Message);
            }
        }

        public async Task<string> SearchAsync(string query, int k)
        {
            if (!m_Ready || string.IsNullOrWhiteSpace(query)) return null;
            try
            {
                // KO->EN translation: e5 cross-lingual hit@5 62% direct vs 85% translated (measured 2026-06).
                var enQuery = await ToEnglishAsync(query).ConfigureAwait(false);

                if (!await IsRelevantAsync(enQuery).ConfigureAwait(false))
                    return null;

                var topK = m_Mode == RagQueryMode.Naive ? k : Math.Max(k, m_MixTopK);
                var param = new QueryParam
                {
                    Mode = ToLightRag(m_Mode),
                    OnlyNeedContext = true,
                    ChunkTopK = topK,
                    TopK = topK,
                };
                var result = await m_Rag.QueryAsync(enQuery, param).ConfigureAwait(false);
                var content = result?.Content;
                if (string.IsNullOrWhiteSpace(content))
                {
                    m_Log?.Invoke($"[RAG] top-{topK} retrieved nothing for \"{enQuery}\".");
                    return null;
                }
                m_Log?.Invoke($"[RAG] top-{topK} context for \"{enQuery}\" ({content.Length} chars):\n{TrimForLog(content)}");
                return content;
            }
            catch (Exception e)
            {
                m_Log?.Invoke("[RAG] text query failed: " + e.Message);
                return null;
            }
        }

        // Expects an already-English query; SearchAsync translates once and passes the result here.
        async Task<bool> IsRelevantAsync(string query)
        {
            if (m_RelevanceThreshold <= 0f || !m_Ready || m_ChunksVdb == null || string.IsNullOrWhiteSpace(query))
                return true;
            try
            {
                var qEmbed = await m_Dispatcher.RunOnMainAsync(() => m_Embedder.EmbedTextAsync(query, isQuery: true)).ConfigureAwait(false);
                if (qEmbed == null) return true;
                var probe = await m_ChunksVdb.QueryAsync(string.Empty, 1, queryEmbedding: qEmbed).ConfigureAwait(false);
                var top = (probe != null && probe.Count > 0 && probe[0].TryGetValue("distance", out var d) && d != null)
                    ? Convert.ToSingle(d) : 0f;
                var relevant = top >= m_RelevanceThreshold;
                if (!relevant)
                    m_Log?.Invoke($"[RAG] gated out (top cosine {top:0.00} < {m_RelevanceThreshold:0.00}): \"{query}\"");
                return relevant;
            }
            catch (Exception e)
            {
                m_Log?.Invoke("[RAG] relevance probe failed (allowing query): " + e.Message);
                return true;
            }
        }

        async Task<string> ToEnglishAsync(string query)
        {
            if (m_Translator == null || !NeedsTranslation(query)) return query;
            try
            {
                var en = await m_Translator.TranslateToEnglishAsync(query).ConfigureAwait(false);
                if (!string.IsNullOrWhiteSpace(en) && !string.Equals(en, query, StringComparison.Ordinal))
                    m_Log?.Invoke($"[RAG] translated query: \"{query}\" -> \"{en}\"");
                return string.IsNullOrWhiteSpace(en) ? query : en;
            }
            catch (Exception e)
            {
                m_Log?.Invoke("[RAG] query translation failed (using original): " + e.Message);
                return query;
            }
        }

        static bool NeedsTranslation(string s)
        {
            if (string.IsNullOrEmpty(s)) return false;
            foreach (var c in s)
                if (c > 0x7F) return true;
            return false;
        }

        // Cap the retrieved context in the log so a large top-k dump doesn't flood the console.
        const int MaxLoggedContextChars = 1500;
        static string TrimForLog(string s)
        {
            return s.Length <= MaxLoggedContextChars ? s : s.Substring(0, MaxLoggedContextChars) + " …(truncated)";
        }

        public async Task<int> InsertAsync(string content, string source = "runtime")
        {
            if (m_Rag == null || string.IsNullOrWhiteSpace(content)) return 0;
            var processed = await m_Rag.InsertAsync(content, source).ConfigureAwait(false);
            if (m_ChunksVdb != null)
                await m_ChunksVdb.InitializeAsync().ConfigureAwait(false);
            return processed;
        }

        public Task FinalizeAsync() => m_Rag != null ? m_Rag.FinalizeAsync() : Task.CompletedTask;

        public void Dispose()
        {
            m_Embedder?.Dispose();
            (m_Llm as IDisposable)?.Dispose();
            m_Ready = false;
        }
    }
}