RAG.cs

Go to the documentation of this file.

 
using System;
using System.IO.Compression;
using System.Threading.Tasks;
using UnityEditor;
using UnityEngine;
 
namespace LLMUnity
{
    public enum SearchMethods
    {
        DBSearch,
        SimpleSearch,
    }
 
    public class NoChunking {}
 
    public enum ChunkingMethods
    {
        NoChunking,
        TokenSplitter,
        WordSplitter,
        SentenceSplitter
    }
 
    [Serializable]
    public class RAG : Searchable
    {
        [Tooltip("Search method type to use for RAG. SimpleSearch is a simple brute-force search, while DBSearch is a fast Approximate Nearest Neighbor (ANN) method (recommended!).")]
        public SearchMethods searchType = SearchMethods.SimpleSearch;
        [Tooltip("Search method GameObject.")]
        public SearchMethod search;
        [Tooltip("Chunking method type to use for RAG for splitting the inputs into chunks. This is useful to have a more consistent meaning within each data part.")]
        public ChunkingMethods chunkingType = ChunkingMethods.NoChunking;
        [Tooltip("Chunking method GameObject.")]
        public Chunking chunking;
 
        public void Init(SearchMethods searchMethod = SearchMethods.SimpleSearch, ChunkingMethods chunkingMethod = ChunkingMethods.NoChunking, LLM llm = null)
        {
            searchType = searchMethod;
            chunkingType = chunkingMethod;
            UpdateGameObjects();
            search.SetLLM(llm);
        }
 
        public void ReturnChunks(bool returnChunks)
        {
            if (chunking != null) chunking.ReturnChunks(returnChunks);
        }
 
        protected void ConstructSearch()
        {
            search = ConstructComponent<SearchMethod>(Type.GetType("LLMUnity." + searchType.ToString()), (previous, current) => current.llmEmbedder.llm = previous.llmEmbedder.llm);
            if (chunking != null) chunking.SetSearch(search);
        }
 
        protected void ConstructChunking()
        {
            Type type = null;
            if (chunkingType != ChunkingMethods.NoChunking) type = Type.GetType("LLMUnity." + chunkingType.ToString());
            chunking = ConstructComponent<Chunking>(type);
            if (chunking != null) chunking.SetSearch(search);
        }
 
        public override void UpdateGameObjects()
        {
            if (this == null) return;
            ConstructSearch();
            ConstructChunking();
        }
 
        protected Searchable GetSearcher()
        {
            if (chunking != null) return chunking;
            if (search != null) return search;
            LLMUnitySetup.LogError("The search GameObject is null", true);
            return null;
        }
 
#if UNITY_EDITOR
        private void OnValidateUpdate()
        {
            EditorApplication.delayCall -= OnValidateUpdate;
            UpdateGameObjects();
        }
 
        public virtual void OnValidate()
        {
            if (!Application.isPlaying) EditorApplication.delayCall += OnValidateUpdate;
        }
 
#endif
 
        public override string Get(int key) { return GetSearcher().Get(key); }
        public override async Task<int> Add(string inputString, string group = "") { return await GetSearcher().Add(inputString, group); }
        public override int Remove(string inputString, string group = "") { return GetSearcher().Remove(inputString, group); }
        public override void Remove(int key) { GetSearcher().Remove(key); }
        public override int Count() { return GetSearcher().Count(); }
        public override int Count(string group) { return GetSearcher().Count(group); }
        public override void Clear() { GetSearcher().Clear(); }
        public override async Task<int> IncrementalSearch(string queryString, string group = "") { return await GetSearcher().IncrementalSearch(queryString, group);}
        public override (string[], float[], bool) IncrementalFetch(int fetchKey, int k) { return GetSearcher().IncrementalFetch(fetchKey, k);}
        public override (int[], float[], bool) IncrementalFetchKeys(int fetchKey, int k) { return GetSearcher().IncrementalFetchKeys(fetchKey, k);}
        public override void IncrementalSearchComplete(int fetchKey) { GetSearcher().IncrementalSearchComplete(fetchKey);}
        public override void Save(ZipArchive archive) { GetSearcher().Save(archive); }
        public override void Load(ZipArchive archive) { GetSearcher().Load(archive); }
    }
}