Runtime/AgentCore/Rag/KnowledgeRagComponent.cs

using System;
using System.Threading;
using System.Threading.Tasks;
using UnityEngine;
using OnDeviceAgent.Inference;
using OnDeviceAgent.RagLlm;
namespace OnDeviceAgent.AgentCore
{

    [DisallowMultipleComponent]
    public sealed class KnowledgeRagComponent : MonoBehaviour
    {
        [Header("Knowledge RAG: LightRAG knowledge graph (multilingual-e5-small embeddings)")]
        [SerializeField] string m_TextModelPath = "Model/E5/e5_small_fp16.sentis";
        [SerializeField] string m_TokenizerFile = "Model/E5/tokenizer.json";
        [SerializeField, Tooltip("StreamingAssets-relative LightRAG DB folder built by OnDeviceAgent ▸ RAG ▸ Ingest Knowledge to DB.")]
        string m_DbDir = "VoiceAgent/DB";
        [SerializeField, Min(1)] int m_RagTopK = 6;
        [SerializeField, Range(0f, 1f), Tooltip("Min e5 query↔chunk cosine for a query to count as on-topic for the indexed knowledge base. " +
            "Below this, LightRAG is skipped so off-topic questions don't bloat the small-model prompt. 0 = always run RAG. " +
            "Calibrated 2026-06: on-topic top-cosine ≥0.82, off-topic ≤0.79 → 0.78 cleanly separates them.")]
        float m_RagRelevanceThreshold = 0.78f;

        [Header("Retrieval mode (LightRAG)")]
        [SerializeField, Tooltip("LightRAG retrieval mode. Naive = e5 vector RAG (best for specific factual lookups; " +
            "graph context is noise there). Mix = graph knowledge-graph RAG (best for high-level / cross-section " +
            "synthesis). Local/Global/Hybrid = LightRAG graph variants.")]
        RagQueryMode m_QueryMode = RagQueryMode.Naive;
        [SerializeField, Min(1), Tooltip("Chunk budget for graph (Mix-family) modes; synthesis gathers across sections. Naive uses Top-K.")]
        int m_MixTopK = 12;

        [Header("Ingest (Editor only, used by OnDeviceAgent ▸ RAG ▸ Ingest Knowledge to DB)")]
        [SerializeField, Tooltip("Ollama endpoint used during knowledge ingest. Cloud-routed models (*-cloud) reach through the same endpoint.")]
        string m_IngestEndpoint = "http://localhost:11434";
        [SerializeField, OllamaModelDropdown(endpointField: "m_IngestEndpoint"),
            Tooltip("LLM used for entity / relation extraction during ingest. A larger model improves KG quality and can differ from the runtime model.")]
        string m_IngestModel = "gemma4:e2b";

        public string IngestEndpoint => m_IngestEndpoint;
        public string IngestModel => m_IngestModel;

        LightRagKnowledgeService m_Rag;
        // Disposed via m_Rag, but exposed so the agent reuses this instance to gate skills instead of loading a second model copy.
        ITextEmbedder m_Embedder;

        public ITextEmbedder TextEmbedder => m_Embedder;

        public bool IsEnabled => enabled;

        public bool Initialize(string endpoint, string model, IUnityMainThreadDispatcher dispatcher)
        {
            if (!enabled) return false;
            if (dispatcher == null) { Debug.LogError("[RAG] dispatcher missing."); return false; }
            if (m_Rag != null) return true;

            var dbWorkingDir = StreamingDbInstaller.Install(m_DbDir, Debug.Log);
            var embedder = new E5TextEmbedder(log: Debug.Log);
            if (!embedder.Load(m_TextModelPath, m_TokenizerFile))
            {
                embedder.Dispose();
                Debug.LogWarning("[RAG] e5 text embedder load failed, knowledge RAG disabled.");
                return false;
            }
            m_Embedder = embedder;

    #if UNITY_ANDROID && !UNITY_EDITOR
            // Route RAG's LLM through the shared LiteRT-LM singleton - Ollama HTTP doesn't exist on-device and reusing avoids a second model load.
            // Match the main transport's SoC-gated backend so this doesn't request a conflicting one for the shared engine.
            var androidTransport = new AndroidLlmTransport(dispatcher, LiteRtModelProvisioner.UseNpu ? "NPU" : "GPU", Debug.Log);
            IChatLlm ragLlm = new AndroidChatLlm(
                (sys, prompt, ct) => androidTransport.RunAsync(sys, prompt, null, null, ct), model);
            m_Rag = new LightRagKnowledgeService(dbWorkingDir, ragLlm, embedder, dispatcher,
                config: null, log: Debug.Log,
                queryMode: m_QueryMode, relevanceThreshold: m_RagRelevanceThreshold, mixTopK: m_MixTopK);
    #else
            m_Rag = new LightRagKnowledgeService(dbWorkingDir, endpoint, model, 8192, embedder, dispatcher, Debug.Log,
                queryMode: m_QueryMode, relevanceThreshold: m_RagRelevanceThreshold, mixTopK: m_MixTopK);
    #endif
            _ = m_Rag.InitializeAsync();
            return true;
        }

        public void Wire(BuiltinToolDependencies deps)
        {
            if (m_Rag == null || deps == null) return;
            deps.SearchKnowledge = SearchKnowledgeTextAsync;
        }

        async Task<string> SearchKnowledgeTextAsync(string query, CancellationToken ct)
        {
            if (m_Rag == null) return null;
            return await m_Rag.SearchAsync(query, m_RagTopK);
        }

        void OnDestroy()
        {
            m_Rag?.Dispose();
            m_Rag = null;
            m_Embedder = null;
        }
    }
}