27 Commits

Author SHA1 Message Date
Grail Finder
2c495253c2 Chore: panic is not solved (WIP) 2026-03-07 09:40:58 +03:00
Grail Finder
118a0a0d55 Chore: cleanup logs 2026-03-07 09:08:01 +03:00
Grail Finder
44633d64c6 Chore: move to own file 2026-03-07 08:52:10 +03:00
Grail Finder
0598e3e86d Feat: kokoro onnx (WIP) 2026-03-07 08:35:44 +03:00
Grail Finder
014e297ae3 Chore: linter complaints 2026-03-06 19:57:44 +03:00
Grail Finder
5f273681df Chore: remove plan doc 2026-03-06 19:03:26 +03:00
Grail Finder
17b68bc21f Enha (rag): async writes 2026-03-06 18:58:23 +03:00
Grail Finder
edfd43c52a Doc: update 2026-03-06 13:45:12 +03:00
Grail Finder
62ec55505c Enha (rag): query each doc 2026-03-06 13:17:49 +03:00
Grail Finder
f9866bcf5a Feat (rag): hybrid search attempt 2026-03-06 11:20:50 +03:00
Grail Finder
822cc48834 Fix: avoid panic if statuslinewidget not loaded yet 2026-03-06 10:37:08 +03:00
Grail Finder
4ef0a21511 Enha (onnx): unload model if noop for 30s 2026-03-06 09:32:45 +03:00
Grail Finder
d2caebdb4f Enha (onnx): use gpu 2026-03-06 09:11:25 +03:00
Grail Finder
e1f2a8cd7b Chore: remove unused RagEnabled var 2026-03-06 07:46:15 +03:00
Grail Finder
efc92d884c Chore: onnx library lookup 2026-03-05 20:02:46 +03:00
Grail Finder
ac8c8bb055 Enha: onnx config vars 2026-03-05 19:20:21 +03:00
Grail Finder
c2c107c786 Dep: make-fetch onnx embed gemma 2026-03-05 16:05:03 +03:00
Grail Finder
c2757653a3 Fix: buildable 2026-03-05 14:49:59 +03:00
Grail Finder
4bd6883966 WIP 2026-03-05 14:38:26 +03:00
Grail Finder
7c56e27dbe Dep: trying sugarme tokenizer 2026-03-05 14:27:19 +03:00
Grail Finder
fbc955ca37 Enha: local onnx 2026-03-05 14:13:58 +03:00
Grail Finder
c65c11bcfb Fix: shellmode tab completion 2026-03-05 11:36:35 +03:00
Grail Finder
04f1fd464b Chore: remove cluedo sysprompt 2026-03-05 11:17:01 +03:00
Grail Finder
6e9c453ee0 Enha: explicit app.Draw per textView update for smooth streaming 2026-03-05 10:35:17 +03:00
Grail Finder
645b7351a8 Fix: add different kind of notifiction for fullscreen mode 2026-03-05 09:09:13 +03:00
Grail Finder
57088565bd Fix (notification): being closed by prev notification early 2026-03-05 08:51:04 +03:00
Grail Finder
4b6769e531 Fix (notification): non-blocking way to notify 2026-03-05 08:43:50 +03:00
23 changed files with 1848 additions and 235 deletions

3
.gitignore vendored
View File

@@ -3,6 +3,8 @@
testlog
history/
*.db
*.db-shm
*.db-wal
config.toml
sysprompts/*
!sysprompts/alice_bob_carl.json
@@ -15,3 +17,4 @@ gflt
chat_exports/*.json
ragimport
.env
onnx/

126
Makefile
View File

@@ -1,4 +1,4 @@
.PHONY: setconfig run lint lintall install-linters setup-whisper build-whisper download-whisper-model docker-up docker-down docker-logs noextra-run installdelve checkdelve
.PHONY: setconfig run lint lintall install-linters setup-whisper build-whisper download-whisper-model docker-up docker-down docker-logs noextra-run installdelve checkdelve fetch-onnx install-onnx-deps fetch-kokoro-voices install-espeak
run: setconfig
go build -tags extra -o gf-lt && ./gf-lt
@@ -30,6 +30,108 @@ lint: ## Run linters. Use make install-linters first.
lintall: lint
noblanks ./...
fetch-onnx:
mkdir -p onnx/embedgemma && curl -o onnx/embedgemma/config.json -L https://huggingface.co/onnx-community/embeddinggemma-300m-ONNX/resolve/main/config.json && curl -o onnx/embedgemma/tokenizer.json -L https://huggingface.co/onnx-community/embeddinggemma-300m-ONNX/resolve/main/tokenizer.json && curl -o onnx/embedgemma/model_q4.onnx -L https://huggingface.co/onnx-community/embeddinggemma-300m-ONNX/resolve/main/onnx/model_q4.onnx && curl -o onnx/embedgemma/model_q4.onnx_data -L https://huggingface.co/onnx-community/embeddinggemma-300m-ONNX/resolve/main/onnx/model_q4.onnx_data?download=true
fetch-kokoro-onnx:
mkdir -p onnx/kokoro && curl -o onnx/kokoro/config.json -L https://huggingface.co/onnx-community/Kokoro-82M-v1.0-ONNX/resolve/main/config.json && curl -o onnx/kokoro/tokenizer.json -L https://huggingface.co/onnx-community/Kokoro-82M-v1.0-ONNX/resolve/main/tokenizer.json && curl -o onnx/kokoro/model_quantized.onnx -L https://huggingface.co/onnx-community/Kokoro-82M-v1.0-ONNX/resolve/main/onnx/model_quantized.onnx && curl -o onnx/kokoro/voices.bin -L https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0/voices-v1.0.bin
install-onnx-deps: ## Install ONNX Runtime with CUDA support (or CPU fallback)
@echo "=== ONNX Runtime Installer ===" && \
echo "" && \
echo "Checking for existing ONNX Runtime..." && \
if ldconfig -p 2>/dev/null | grep -q libonnxruntime.so.1; then \
echo "ONNX Runtime is already installed:" && \
ldconfig -p 2>/dev/null | grep libonnxruntime && \
echo "" && \
echo "Skipping installation. To reinstall, remove existing libs first:" && \
echo " sudo rm -f /usr/local/lib/libonnxruntime*.so*" && \
exit 0; \
fi && \
echo "No ONNX Runtime found. Proceeding with installation..." && \
echo "" && \
echo "Detecting CUDA version..." && \
HAS_CUDA=0 && \
if command -v nvidia-smi >/dev/null 2>&1; then \
CUDA_INFO=$$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -1) && \
if [ -n "$$CUDA_INFO" ]; then \
echo "Found NVIDIA GPU with driver: $$CUDA_INFO" && \
HAS_CUDA=1; \
else \
echo "NVIDIA driver found but could not detect CUDA version"; \
fi; \
else \
echo "No NVIDIA GPU detected (nvidia-smi not found)"; \
fi && \
echo "" && \
echo "Determining ONNX Runtime version..." && \
ARCH=$$(uname -m) && \
if [ "$$ARCH" = "x86_64" ]; then \
ONNX_ARCH="x64"; \
elif [ "$$ARCH" = "aarch64" ] || [ "$$ARCH" = "arm64" ]; then \
ONNX_ARCH="aarch64"; \
else \
echo "Unsupported architecture: $$ARCH" && \
exit 1; \
fi && \
echo "Detected architecture: $$ARCH (ONNX runtime: $$ONNX_ARCH)" && \
if [ "$$HAS_CUDA" = "1" ]; then \
echo "Installing ONNX Runtime with CUDA support..."; \
ONNX_VERSION="1.24.2"; \
else \
echo "Installing ONNX Runtime (CPU version)..."; \
ONNX_VERSION="1.24.2"; \
fi && \
FILENAME="onnxruntime-linux-$${ONNX_ARCH}-${ONNX_VERSION}.tgz" && \
URL="https://github.com/microsoft/onnxruntime/releases/download/v$${ONNX_VERSION}/$${FILENAME}" && \
echo "Downloading $${URL}..." && \
mkdir -p /tmp/onnx-install && \
curl -L -o /tmp/onnx-install/$${FILENAME} "$${URL}" || { \
echo "Failed to download ONNX Runtime v$${ONNX_VERSION}. Trying v1.18.0..." && \
ONNX_VERSION="1.18.0" && \
FILENAME="onnxruntime-linux-$${ONNX_ARCH}-${ONNX_VERSION}.tgz" && \
URL="https://github.com/microsoft/onnxruntime/releases/download/v$${ONNX_VERSION}/$${FILENAME}" && \
curl -L -o /tmp/onnx-install/$${FILENAME} "$${URL}" || { \
echo "ERROR: Failed to download ONNX Runtime from GitHub" && \
echo "" && \
echo "Please install manually:" && \
echo " 1. Go to https://github.com/microsoft/onnxruntime/releases" && \
echo " 2. Download onnxruntime-linux-$${ONNX_ARCH}-VERSION.tgz" && \
echo " 3. Extract and copy to /usr/local/lib:" && \
echo " tar -xzf onnxruntime-linux-$${ONNX_ARCH}-VERSION.tgz" && \
echo " sudo cp -r onnxruntime-linux-$${ONNX_ARCH}-VERSION/lib/* /usr/local/lib/" && \
echo " sudo ldconfig" && \
exit 1; \
}; \
} && \
echo "Extracting..." && \
cd /tmp/onnx-install && tar -xzf $${FILENAME} && \
echo "Installing to /usr/local/lib..." && \
ONNX_DIR=$$(find /tmp/onnx-install -maxdepth 1 -type d -name "onnxruntime-linux-*") && \
if [ -d "$${ONNX_DIR}/lib" ]; then \
cp -r $${ONNX_DIR}/lib/* /usr/local/lib/ 2>/dev/null || sudo cp -r $${ONNX_DIR}/lib/* /usr/local/lib/; \
else \
echo "ERROR: Could not find lib directory in extracted archive" && \
exit 1; \
fi && \
echo "Updating library cache..." && \
sudo ldconfig 2>/dev/null || ldconfig && \
echo "" && \
echo "=== Installation complete! ===" && \
echo "" && \
echo "Installed libraries:" && \
ldconfig -p | grep libonnxruntime || echo "(libraries may require logout/relogin to appear)" && \
echo "" && \
if [ "$$HAS_CUDA" = "1" ]; then \
echo "NOTE: CUDA-enabled ONNX Runtime installed."; \
echo "Ensure you also have CUDA libraries installed:"; \
echo " - libcudnn, libcublas, libcurand"; \
else \
echo "NOTE: CPU-only ONNX Runtime installed."; \
echo "For GPU support, install CUDA and re-run this script."; \
fi && \
rm -rf /tmp/onnx-install
# Whisper STT Setup (in batteries directory)
setup-whisper: build-whisper download-whisper-model
@@ -95,3 +197,25 @@ docker-logs-whisper: ## View logs from Whisper STT service only
docker-logs-kokoro: ## View logs from Kokoro TTS service only
@echo "Displaying logs from Kokoro TTS service..."
docker-compose -f batteries/docker-compose.yml logs -f kokoro-tts
# Kokoro ONNX TTS Setup
install-espeak: ## Install espeak-ng for phoneme tokenization
@echo "=== Installing espeak-ng ===" && \
if command -v espeak-ng >/dev/null 2>&1; then \
echo "espeak-ng is already installed:" && \
espeak-ng --version && \
exit 0; \
fi && \
echo "Installing espeak-ng..." && \
sudo apt-get update && \
sudo apt-get install -y espeak-ng espeak && \
echo "espeak-ng installed successfully!" && \
espeak-ng --version
fetch-kokoro-voices: ## Download Kokoro voice files (PyTorch format)
@echo "=== Downloading Kokoro voices ===" && \
mkdir -p onnx/kokoro/voices && \
echo "Downloading af_bella voice..." && \
curl -L -o onnx/kokoro/voices/af_bella.pt https://raw.githubusercontent.com/hexgrad/kokoro/main/kokoro/voices/af_heart.pt && \
echo "Voice file downloaded to onnx/kokoro/voices/" && \
ls -lh onnx/kokoro/voices/

22
bot.go
View File

@@ -418,7 +418,9 @@ func fetchLCPModelsWithStatus() (*models.LCPModels, error) {
if err := json.NewDecoder(resp.Body).Decode(data); err != nil {
return nil, err
}
localModelsMu.Lock()
localModelsData = data
localModelsMu.Unlock()
return data, nil
}
@@ -1393,12 +1395,16 @@ func updateModelLists() {
}
}
// if llama.cpp started after gf-lt?
localModelsMu.Lock()
LocalModels, err = fetchLCPModelsWithLoadStatus()
localModelsMu.Unlock()
ml, err := fetchLCPModelsWithLoadStatus()
if err != nil {
logger.Warn("failed to fetch llama.cpp models", "error", err)
}
localModelsMu.Lock()
LocalModels = ml
localModelsMu.Unlock()
for statusLineWidget == nil {
time.Sleep(time.Millisecond * 100)
}
// set already loaded model in llama.cpp
if strings.Contains(cfg.CurrentAPI, "localhost") || strings.Contains(cfg.CurrentAPI, "127.0.0.1") {
localModelsMu.Lock()
@@ -1493,14 +1499,20 @@ func init() {
// load cards
basicCard.Role = cfg.AssistantRole
logLevel.Set(slog.LevelInfo)
logger = slog.New(slog.NewTextHandler(logfile, &slog.HandlerOptions{Level: logLevel}))
logger = slog.New(slog.NewTextHandler(logfile, &slog.HandlerOptions{Level: logLevel, AddSource: true}))
store = storage.NewProviderSQL(cfg.DBPATH, logger)
if store == nil {
cancel()
os.Exit(1)
return
}
ragger = rag.New(logger, store, cfg)
ragger, err = rag.New(logger, store, cfg)
if err != nil {
logger.Error("failed to create RAG", "error", err)
}
if ragger != nil && ragger.FallbackMessage() != "" && app != nil {
showToast("RAG", "ONNX unavailable, using API: "+ragger.FallbackMessage())
}
// https://github.com/coreydaley/ggerganov-llama.cpp/blob/master/examples/server/README.md
// load all chats in memory
if _, err := loadHistoryChats(); err != nil {

View File

@@ -13,6 +13,9 @@ OpenRouterChatAPI = "https://openrouter.ai/api/v1/chat/completions"
# embeddings
EmbedURL = "http://localhost:8082/v1/embeddings"
HFToken = ""
EmbedModelPath = "onnx/embedgemma/model_q4.onnx"
EmbedTokenizerPath = "onnx/embedgemma/tokenizer.json"
EmbedDims = 768
#
ShowSys = true
LogFile = "log.txt"
@@ -24,9 +27,9 @@ ChunkLimit = 100000
AutoScrollEnabled = true
AutoCleanToolCallsFromCtx = false
# rag settings
RAGEnabled = false
RAGBatchSize = 1
RAGWordLimit = 80
RAGOverlapWords = 16
RAGDir = "ragimport"
# extra tts
TTS_ENABLED = false

View File

@@ -36,11 +36,14 @@ type Config struct {
// embeddings
EmbedURL string `toml:"EmbedURL"`
HFToken string `toml:"HFToken"`
EmbedModelPath string `toml:"EmbedModelPath"`
EmbedTokenizerPath string `toml:"EmbedTokenizerPath"`
EmbedDims int `toml:"EmbedDims"`
// rag settings
RAGEnabled bool `toml:"RAGEnabled"`
RAGDir string `toml:"RAGDir"`
RAGBatchSize int `toml:"RAGBatchSize"`
RAGWordLimit uint32 `toml:"RAGWordLimit"`
RAGOverlapWords uint32 `toml:"RAGOverlapWords"`
// deepseek
DeepSeekChatAPI string `toml:"DeepSeekChatAPI"`
DeepSeekCompletionAPI string `toml:"DeepSeekCompletionAPI"`
@@ -58,6 +61,10 @@ type Config struct {
TTS_SPEED float32 `toml:"TTS_SPEED"`
TTS_PROVIDER string `toml:"TTS_PROVIDER"`
TTS_LANGUAGE string `toml:"TTS_LANGUAGE"`
// Kokoro ONNX TTS
KokoroModelPath string `toml:"KokoroModelPath"`
KokoroVoicesPath string `toml:"KokoroVoicesPath"`
KokoroVoice string `toml:"KokoroVoice"`
// STT
STT_TYPE string `toml:"STT_TYPE"` // WHISPER_SERVER, WHISPER_BINARY
STT_URL string `toml:"STT_URL"`

View File

@@ -71,9 +71,6 @@ This document explains how to set up and configure the application using the `co
#### EmbedURL (`"http://localhost:8082/v1/embeddings"`)
- The endpoint for embedding API, used for RAG (Retrieval Augmented Generation) functionality.
#### RAGEnabled (`false`)
- Enable or disable RAG functionality for enhanced context retrieval.
#### RAGBatchSize (`1`)
- Number of documents to process in each RAG batch.

421
extra/kokoro_onnx.go Normal file
View File

@@ -0,0 +1,421 @@
//go:build extra
// +build extra
package extra
import (
"bytes"
"fmt"
"gf-lt/models"
"gf-lt/onnx"
"log/slog"
"os/exec"
"strings"
"sync"
"time"
"github.com/gopxl/beep/v2"
"github.com/gopxl/beep/v2/speaker"
"github.com/gopxl/beep/v2/wav"
"github.com/neurosnap/sentences/english"
"github.com/yalue/onnxruntime_go"
)
// KokoroONNXOrator implements Kokoro TTS using ONNX runtime
type KokoroONNXOrator struct {
logger *slog.Logger
mu sync.Mutex
session *onnxruntime_go.DynamicAdvancedSession
phonemeMap map[string]int
espeakCmd string
voice string
speed float32
styleVector []float32
currentStream *beep.Ctrl
currentDone chan bool
textBuffer strings.Builder
interrupt bool
modelLoaded bool
modelPath string
voicesPath string
}
// Phoneme to token ID mapping from Kokoro tokenizer.json
var kokoroPhonemeMap = map[string]int{
"$": 0, ";": 1, ":": 2, ",": 3, ".": 4, "!": 5, "?": 6, "—": 9, "…": 10, "\"": 11, "(": 12, ")": 13, "“": 14, "”": 15, " ": 16, "̃": 17, "ˢ": 18, "ˤ": 19, "˦": 20, "˨": 21, "ᾝ": 22, "⭧": 23,
"A": 24, "I": 25, "O": 31, "Q": 33, "S": 35, "T": 36, "W": 39, "Y": 41, "ʲ": 42,
"a": 43, "b": 44, "c": 45, "d": 46, "e": 47, "f": 48, "h": 50, "i": 51, "j": 52, "k": 53, "l": 54, "m": 55, "n": 56, "o": 57, "p": 58, "q": 59, "r": 60, "s": 61, "t": 62, "u": 63, "v": 64, "w": 65, "x": 66, "y": 67, "z": 68,
"ɑ": 69, "ɐ": 70, "ɒ": 71, "æ": 72, "β": 75, "ɔ": 76, "ɕ": 77, "ç": 78, "ɖ": 80, "ð": 81, "˔": 82, "ə": 83, "ɚ": 85, "ɛ": 86, "ɜ": 87, "ɟ": 90, "ɡ": 92, "ɥ": 99, "ɨ": 101, "ɪ": 102, "ɝ": 103, "ɯ": 110, "ɰ": 111, "ŋ": 112, "ɳ": 113, "ɲ": 114, "ɴ": 115, "ø": 116, "ɸ": 118, "θ": 119, "œ": 120, "ɹ": 123, "ɾ": 125, "ɺ": 126, "ʁ": 128, "ɽ": 129, "ʂ": 130, "ʃ": 131, "ʈ": 132, "˧": 133, "ʊ": 135, "ʋ": 136, "ʌ": 138, "ɢ": 139, "ɣ": 140, "χ": 142, "ʎ": 143, "ʒ": 147, "ʔ": 148,
"ˈ": 156, "ˌ": 157, "ː": 158, "̰": 162, "̊": 164, "↕": 169, "→": 171, "↗": 172, "↘": 173, "ᶻ": 177,
}
func (o *KokoroONNXOrator) ensureInitialized(modelPath string) error {
if o.modelLoaded {
return nil
}
o.mu.Lock()
defer o.mu.Unlock()
if o.modelLoaded {
return nil
}
if modelPath == "" {
o.logger.Error("modelPath is empty, cannot load ONNX model")
return fmt.Errorf("modelPath is empty, set KokoroModelPath in config")
}
// Initialize ONNX runtime (shared with embedder)
if err := onnx.Init(); err != nil {
o.logger.Error("ONNX init failed", "error", err)
return fmt.Errorf("ONNX init failed: %w", err)
}
if onnx.HasCUDASupport() {
o.logger.Info("ONNX using CUDA")
} else {
o.logger.Info("ONNX using CPU fallback")
}
if o.phonemeMap == nil {
o.phonemeMap = kokoroPhonemeMap
}
if o.espeakCmd == "" {
o.espeakCmd = "espeak-ng"
if _, err := exec.LookPath(o.espeakCmd); err != nil {
o.espeakCmd = "espeak"
if _, err := exec.LookPath(o.espeakCmd); err != nil {
return fmt.Errorf("espeak-ng or espeak not found. Install with: sudo apt-get install espeak-ng")
}
}
}
o.logger.Info("using espeak command", "cmd", o.espeakCmd)
// Load voice embedding if not already loaded
if o.styleVector == nil {
voiceName := o.voice
if voiceName == "" {
voiceName = "af_bella"
}
if o.voicesPath != "" {
styleVec, err := onnx.LoadVoice(o.voicesPath, voiceName)
if err != nil {
o.logger.Warn("failed to load voice, using zeros", "error", err, "voice", voiceName)
o.styleVector = make([]float32, 256)
} else {
// Shape is (510, 1, 256), we want the last 256 values (or first? let's use mean or just pick one)
// Actually, let's average across all 510 to get a single 256-dim vector
if len(styleVec) != 510*256 {
o.logger.Error("voice embedding has unexpected size", "len", len(styleVec))
err = fmt.Errorf("voice embedding has unexpected size", "len", len(styleVec))
return err
}
o.styleVector = make([]float32, 256)
for i := 0; i < 256; i++ {
var sum float32
for j := 0; j < 510; j++ {
sum += styleVec[j*256+i]
}
o.styleVector[i] = sum / 510.0
}
o.logger.Info("loaded voice embedding", "voice", voiceName)
}
} else {
o.logger.Warn("no voices path configured, using zeros for style")
o.styleVector = make([]float32, 256)
}
}
opts, err := onnx.NewSessionOptions()
if err != nil {
return fmt.Errorf("failed to create session options: %w", err)
}
defer func() { _ = opts.Destroy() }()
if onnx.HasCUDASupport() {
o.logger.Info("session options created with CUDA")
} else {
o.logger.Info("session options created with CPU")
}
session, err := onnxruntime_go.NewDynamicAdvancedSession(
modelPath,
[]string{"input_ids", "style", "speed"},
[]string{"waveform"},
opts,
)
if err != nil {
o.logger.Error("failed to create ONNX session", "error", err)
return fmt.Errorf("failed to create ONNX session: %w", err)
}
o.session = session
o.modelLoaded = true
o.logger.Info("Kokoro ONNX model loaded successfully", "model", modelPath)
return nil
}
func (o *KokoroONNXOrator) textToPhonemes(text string) (string, error) {
cmd := exec.Command(o.espeakCmd, "-x", "-q", text)
output, err := cmd.Output()
if err != nil {
o.logger.Error("espeak failed", "error", err, "cmd", o.espeakCmd, "text", text)
return "", fmt.Errorf("espeak failed: %w", err)
}
phonemeStr := strings.TrimSpace(string(output))
return phonemeStr, nil
}
func (o *KokoroONNXOrator) phonemesToTokens(phonemeStr string) ([]int, error) {
if phonemeStr == "" {
o.logger.Error("empty phoneme string")
return nil, fmt.Errorf("empty phoneme string")
}
// Iterate over each character in the phoneme string
tokens := make([]int, 0)
for _, ch := range phonemeStr {
chStr := string(ch)
if tokenID, ok := o.phonemeMap[chStr]; ok {
tokens = append(tokens, tokenID)
}
}
if len(tokens) == 0 {
o.logger.Error("no phonemes mapped to tokens", "phonemeStr", phonemeStr)
return nil, fmt.Errorf("no valid phonemes mapped to tokens")
}
return tokens, nil
}
func (o *KokoroONNXOrator) generateAudio(text string) ([]float32, error) {
if err := o.ensureInitialized(o.modelPath); err != nil {
o.logger.Error("ensureInitialized failed", "error", err)
return nil, err
}
phonemeStr, err := o.textToPhonemes(text)
if err != nil {
o.logger.Error("phoneme conversion failed", "error", err)
return nil, fmt.Errorf("phoneme conversion failed: %w", err)
}
tokens, err := o.phonemesToTokens(phonemeStr)
if err != nil {
o.logger.Error("token conversion failed", "error", err)
return nil, fmt.Errorf("token conversion failed: %w", err)
}
if len(tokens) > 510 {
return nil, fmt.Errorf("text too long: %d tokens (max 510)", len(tokens))
}
tokens = append([]int{0}, tokens...)
tokens = append(tokens, 0)
inputIDs := make([]int64, len(tokens))
for i, t := range tokens {
inputIDs[i] = int64(t)
}
inputTensor, err := onnxruntime_go.NewTensor[int64](
onnxruntime_go.NewShape(1, int64(len(inputIDs))),
inputIDs,
)
if err != nil {
o.logger.Error("failed to create input tensor", "error", err)
return nil, fmt.Errorf("failed to create input tensor: %w", err)
}
defer func() { _ = inputTensor.Destroy() }()
styleTensor, err := onnxruntime_go.NewTensor[float32](
onnxruntime_go.NewShape(1, 256),
o.styleVector,
)
if err != nil {
o.logger.Error("failed to create style tensor", "error", err)
return nil, fmt.Errorf("failed to create style tensor: %w", err)
}
defer func() { _ = styleTensor.Destroy() }()
speedTensor, err := onnxruntime_go.NewTensor[float32](
onnxruntime_go.NewShape(1),
[]float32{o.speed},
)
if err != nil {
o.logger.Error("failed to create speed tensor", "error", err)
return nil, fmt.Errorf("failed to create speed tensor: %w", err)
}
defer func() { _ = speedTensor.Destroy() }()
outputTensor, err := onnxruntime_go.NewEmptyTensor[float32](
onnxruntime_go.NewShape(1, 512),
)
if err != nil {
o.logger.Error("failed to create output tensor", "error", err)
return nil, fmt.Errorf("failed to create output tensor: %w", err)
}
defer func() { _ = outputTensor.Destroy() }()
err = o.session.Run(
[]onnxruntime_go.Value{inputTensor, styleTensor, speedTensor},
[]onnxruntime_go.Value{outputTensor},
)
if err != nil {
o.logger.Error("ONNX inference failed", "error", err)
return nil, fmt.Errorf("ONNX inference failed: %w", err)
}
audioData := outputTensor.GetData()
if len(audioData) == 0 {
o.logger.Error("empty audio output from ONNX")
return nil, fmt.Errorf("empty audio output")
}
audio := make([]float32, len(audioData))
copy(audio, audioData)
return audio, nil
}
func (o *KokoroONNXOrator) Speak(text string) error {
audio, err := o.generateAudio(text)
if err != nil {
o.logger.Error("audio generation failed", "error", err)
return fmt.Errorf("audio generation failed: %w", err)
}
// Create streamer for encoding
encodeStreamer := beep.StreamerFunc(func(samples [][2]float64) (n int, ok bool) {
for i := range samples {
if i >= len(audio) {
return i, false
}
samples[i][0] = float64(audio[i])
samples[i][1] = float64(audio[i])
}
return len(audio), true
})
buf := &seekableBuffer{new(bytes.Buffer)}
err = wav.Encode(buf, encodeStreamer, beep.Format{
SampleRate: 24000,
NumChannels: 1,
Precision: 2,
})
if err != nil {
o.logger.Error("wav encoding failed", "error", err)
return fmt.Errorf("wav encoding failed: %w", err)
}
decodedStreamer, format, err := wav.Decode(bytes.NewReader(buf.Bytes()))
if err != nil {
o.logger.Error("wav decode failed", "error", err)
return fmt.Errorf("wav decode failed: %w", err)
}
defer decodedStreamer.Close()
if err := speaker.Init(format.SampleRate, format.SampleRate.N(time.Second/10)); err != nil {
o.logger.Error("speaker init failed", "error", err)
return fmt.Errorf("speaker init failed: %w", err)
}
o.logger.Info("playing audio", "sampleRate", format.SampleRate, "channels", format.NumChannels)
done := make(chan bool)
o.mu.Lock()
o.currentDone = done
o.currentStream = &beep.Ctrl{Streamer: beep.Seq(decodedStreamer, beep.Callback(func() {
o.mu.Lock()
close(done)
o.currentStream = nil
o.currentDone = nil
o.mu.Unlock()
})), Paused: false}
o.mu.Unlock()
speaker.Play(o.currentStream)
<-done
return nil
}
func (o *KokoroONNXOrator) Stop() {
speaker.Lock()
defer speaker.Unlock()
o.mu.Lock()
defer o.mu.Unlock()
if o.currentStream != nil {
o.currentStream.Streamer = nil
}
}
func (o *KokoroONNXOrator) GetLogger() *slog.Logger {
return o.logger
}
func (o *KokoroONNXOrator) stoproutine() {
for {
<-TTSDoneChan
o.Stop()
for len(TTSTextChan) > 0 {
<-TTSTextChan
}
o.mu.Lock()
o.textBuffer.Reset()
if o.currentDone != nil {
select {
case o.currentDone <- true:
default:
}
}
o.interrupt = true
o.mu.Unlock()
}
}
func (o *KokoroONNXOrator) readroutine() {
tokenizer, _ := english.NewSentenceTokenizer(nil)
for {
select {
case chunk := <-TTSTextChan:
o.mu.Lock()
o.interrupt = false
_, err := o.textBuffer.WriteString(chunk)
if err != nil {
o.logger.Warn("failed to write to buffer", "error", err)
o.mu.Unlock()
continue
}
text := o.textBuffer.String()
sentences := tokenizer.Tokenize(text)
if len(sentences) <= 1 {
o.mu.Unlock()
continue
}
completeSentences := sentences[:len(sentences)-1]
remaining := sentences[len(sentences)-1].Text
o.textBuffer.Reset()
o.textBuffer.WriteString(remaining)
o.mu.Unlock()
for _, sentence := range completeSentences {
o.mu.Lock()
interrupted := o.interrupt
o.mu.Unlock()
if interrupted {
return
}
cleanedText := models.CleanText(sentence.Text)
if cleanedText == "" {
continue
}
o.logger.Info("KokoroONNX speak", "text", cleanedText)
if err := o.Speak(cleanedText); err != nil {
o.logger.Error("KokoroONNX tts failed", "text", cleanedText, "error", err)
}
}
case <-TTSFlushChan:
if len(TTSTextChan) > 0 {
for chunk := range TTSTextChan {
o.mu.Lock()
_, err := o.textBuffer.WriteString(chunk)
o.mu.Unlock()
if err != nil {
continue
}
if len(TTSTextChan) == 0 {
break
}
}
}
o.mu.Lock()
remaining := o.textBuffer.String()
remaining = models.CleanText(remaining)
o.textBuffer.Reset()
o.mu.Unlock()
if remaining == "" {
continue
}
sentencesRem := tokenizer.Tokenize(remaining)
for _, rs := range sentencesRem {
o.mu.Lock()
interrupt := o.interrupt
o.mu.Unlock()
if interrupt {
break
}
if err := o.Speak(rs.Text); err != nil {
o.logger.Error("tts failed", "text", rs.Text, "error", err)
}
}
}
}
}

View File

@@ -32,6 +32,14 @@ var (
// endsWithPunctuation = regexp.MustCompile(`[;.!?]$`)
)
type seekableBuffer struct {
*bytes.Buffer
}
func (s *seekableBuffer) Seek(offset int64, whence int) (int64, error) {
return 0, nil
}
type Orator interface {
Speak(text string) error
Stop()
@@ -194,6 +202,18 @@ func NewOrator(log *slog.Logger, cfg *config.Config) Orator {
go orator.readroutine()
go orator.stoproutine()
return orator
case "kokoro_onnx":
log.Info("Initializing Kokoro ONNX TTS", "modelPath", cfg.KokoroModelPath, "voicesPath", cfg.KokoroVoicesPath, "voice", cfg.KokoroVoice, "speed", cfg.TTS_SPEED)
orator := &KokoroONNXOrator{
logger: log,
modelPath: cfg.KokoroModelPath,
voicesPath: cfg.KokoroVoicesPath,
speed: cfg.TTS_SPEED,
voice: cfg.KokoroVoice,
}
go orator.readroutine()
go orator.stoproutine()
return orator
default:
language := cfg.TTS_LANGUAGE
if language == "" {

9
go.mod
View File

@@ -7,7 +7,6 @@ require (
github.com/GrailFinder/google-translate-tts v0.1.3
github.com/GrailFinder/searchagent v0.2.0
github.com/PuerkitoBio/goquery v1.11.0
github.com/deckarep/golang-set/v2 v2.8.0
github.com/gdamore/tcell/v2 v2.13.2
github.com/glebarez/go-sqlite v1.22.0
github.com/gopxl/beep/v2 v2.1.1
@@ -17,14 +16,18 @@ require (
github.com/neurosnap/sentences v1.1.2
github.com/playwright-community/playwright-go v0.5700.1
github.com/rivo/tview v0.42.0
github.com/sugarme/tokenizer v0.3.0
github.com/yalue/onnxruntime_go v1.27.0
github.com/yuin/goldmark v1.4.13
)
require (
github.com/andybalholm/cascadia v1.3.3 // indirect
github.com/deckarep/golang-set/v2 v2.8.0 // indirect
github.com/dustin/go-humanize v1.0.1 // indirect
github.com/ebitengine/oto/v3 v3.4.0 // indirect
github.com/ebitengine/purego v0.9.1 // indirect
github.com/emirpasic/gods v1.18.1 // indirect
github.com/gdamore/encoding v1.0.1 // indirect
github.com/go-jose/go-jose/v3 v3.0.4 // indirect
github.com/go-stack/stack v1.8.1 // indirect
@@ -33,10 +36,14 @@ require (
github.com/hajimehoshi/oto/v2 v2.3.1 // indirect
github.com/lucasb-eyer/go-colorful v1.3.0 // indirect
github.com/mattn/go-isatty v0.0.20 // indirect
github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db // indirect
github.com/ncruces/go-strftime v1.0.0 // indirect
github.com/patrickmn/go-cache v2.1.0+incompatible // indirect
github.com/pkg/errors v0.9.1 // indirect
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
github.com/rivo/uniseg v0.4.7 // indirect
github.com/schollz/progressbar/v2 v2.15.0 // indirect
github.com/sugarme/regexpset v0.0.0-20200920021344-4d4ec8eaf93c // indirect
golang.org/x/exp v0.0.0-20251209150349-8475f28825e9 // indirect
golang.org/x/net v0.48.0 // indirect
golang.org/x/sys v0.39.0 // indirect

15
go.sum
View File

@@ -21,6 +21,8 @@ github.com/ebitengine/oto/v3 v3.4.0 h1:br0PgASsEWaoWn38b2Goe7m1GKFYfNgnsjSd5Gg+/
github.com/ebitengine/oto/v3 v3.4.0/go.mod h1:IOleLVD0m+CMak3mRVwsYY8vTctQgOM0iiL6S7Ar7eI=
github.com/ebitengine/purego v0.9.1 h1:a/k2f2HQU3Pi399RPW1MOaZyhKJL9w/xFpKAg4q1s0A=
github.com/ebitengine/purego v0.9.1/go.mod h1:iIjxzd6CiRiOG0UyXP+V1+jWqUXVjPKLAI0mRfJZTmQ=
github.com/emirpasic/gods v1.18.1 h1:FXtiHYKDGKCW2KzwZKx0iC0PQmdlorYgdFG9jPXJ1Bc=
github.com/emirpasic/gods v1.18.1/go.mod h1:8tpGGwCnJ5H4r6BWwaV6OrWmMoPhUl5jm/FMNAnJvWQ=
github.com/gdamore/encoding v1.0.1 h1:YzKZckdBL6jVt2Gc+5p82qhrGiqMdG/eNs6Wy0u3Uhw=
github.com/gdamore/encoding v1.0.1/go.mod h1:0Z0cMFinngz9kS1QfMjCP8TY7em3bZYeeklsSDPivEo=
github.com/gdamore/tcell/v2 v2.13.2 h1:5j4srfF8ow3HICOv/61/sOhQtA25qxEB2XR3Q/Bhx2g=
@@ -61,10 +63,14 @@ github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWE
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
github.com/mattn/go-sqlite3 v1.14.22 h1:2gZY6PC6kBnID23Tichd1K+Z0oS6nE/XwU+Vz/5o4kU=
github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db h1:62I3jR2EmQ4l5rM/4FEfDWcRD+abF5XlKShorW5LRoQ=
github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db/go.mod h1:l0dey0ia/Uv7NcFFVbCLtqEBQbrT4OCwCSKTEv6enCw=
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
github.com/neurosnap/sentences v1.1.2 h1:iphYOzx/XckXeBiLIUBkPu2EKMJ+6jDbz/sLJZ7ZoUw=
github.com/neurosnap/sentences v1.1.2/go.mod h1:/pwU4E9XNL21ygMIkOIllv/SMy2ujHwpf8GQPu1YPbQ=
github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
github.com/playwright-community/playwright-go v0.5700.1 h1:PNFb1byWqrTT720rEO0JL88C6Ju0EmUnR5deFLvtP/U=
@@ -77,10 +83,19 @@ github.com/rivo/tview v0.42.0 h1:b/ftp+RxtDsHSaynXTbJb+/n/BxDEi+W3UfF5jILK6c=
github.com/rivo/tview v0.42.0/go.mod h1:cSfIYfhpSGCjp3r/ECJb+GKS7cGJnqV8vfjQPwoXyfY=
github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
github.com/schollz/progressbar/v2 v2.15.0 h1:dVzHQ8fHRmtPjD3K10jT3Qgn/+H+92jhPrhmxIJfDz8=
github.com/schollz/progressbar/v2 v2.15.0/go.mod h1:UdPq3prGkfQ7MOzZKlDRpYKcFqEMczbD7YmbPgpzKMI=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
github.com/sugarme/regexpset v0.0.0-20200920021344-4d4ec8eaf93c h1:pwb4kNSHb4K89ymCaN+5lPH/MwnfSVg4rzGDh4d+iy4=
github.com/sugarme/regexpset v0.0.0-20200920021344-4d4ec8eaf93c/go.mod h1:2gwkXLWbDGUQWeL3RtpCmcY4mzCtU13kb9UsAg9xMaw=
github.com/sugarme/tokenizer v0.3.0 h1:FE8DYbNSz/kSbgEo9l/RjgYHkIJYEdskumitFQBE9FE=
github.com/sugarme/tokenizer v0.3.0/go.mod h1:VJ+DLK5ZEZwzvODOWwY0cw+B1dabTd3nCB5HuFCItCc=
github.com/yalue/onnxruntime_go v1.27.0 h1:c1YSgDNtpf0WGtxj3YeRIb8VC5LmM1J+Ve3uHdteC1U=
github.com/yalue/onnxruntime_go v1.27.0/go.mod h1:b4X26A8pekNb1ACJ58wAXgNKeUCGEAQ9dmACut9Sm/4=
github.com/yuin/goldmark v1.4.13 h1:fVcFKWvrslecOb/tg+Cc05dkeYx540o0FuFt3nUVDoE=
github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=

View File

@@ -521,7 +521,7 @@ func updateFlexLayout() {
if shellMode {
flex.AddItem(shellInput, 0, 10, false)
} else {
flex.AddItem(textArea, 0, 10, false)
flex.AddItem(bottomFlex, 0, 10, true)
}
if positionVisible {
flex.AddItem(statusLineWidget, 0, 2, false)

View File

@@ -115,9 +115,6 @@ func makePropsTable(props map[string]float32) *tview.Table {
row++
}
// Add checkboxes
addCheckboxRow("RAG use", cfg.RAGEnabled, func(checked bool) {
cfg.RAGEnabled = checked
})
addCheckboxRow("Inject role", injectRole, func(checked bool) {
injectRole = checked
})

View File

@@ -7,8 +7,16 @@ import (
"fmt"
"gf-lt/config"
"gf-lt/models"
"gf-lt/onnx"
"log/slog"
"net/http"
"os"
"sync"
"time"
"github.com/sugarme/tokenizer"
"github.com/sugarme/tokenizer/pretrained"
"github.com/yalue/onnxruntime_go"
)
// Embedder defines the interface for embedding text
@@ -27,7 +35,9 @@ type APIEmbedder struct {
func NewAPIEmbedder(l *slog.Logger, cfg *config.Config) *APIEmbedder {
return &APIEmbedder{
logger: l,
client: &http.Client{},
client: &http.Client{
Timeout: 30 * time.Second,
},
cfg: cfg,
}
}
@@ -134,11 +144,230 @@ func (a *APIEmbedder) EmbedSlice(lines []string) ([][]float32, error) {
return embeddings, nil
}
// TODO: ONNXEmbedder implementation would go here
// This would require:
// 1. Loading ONNX models locally
// 2. Using a Go ONNX runtime (like gorgonia/onnx or similar)
// 3. Converting text to embeddings without external API calls
//
// For now, we'll focus on the API implementation which is already working in the current system,
// and can be extended later when we have ONNX runtime integration
type ONNXEmbedder struct {
session *onnxruntime_go.DynamicAdvancedSession
tokenizer *tokenizer.Tokenizer
tokenizerPath string
dims int
logger *slog.Logger
mu sync.Mutex
modelPath string
}
func NewONNXEmbedder(modelPath, tokenizerPath string, dims int, logger *slog.Logger) (*ONNXEmbedder, error) {
// Check if model and tokenizer files exist
if _, err := os.Stat(modelPath); err != nil {
return nil, fmt.Errorf("ONNX model not found: %w", err)
}
if _, err := os.Stat(tokenizerPath); err != nil {
return nil, fmt.Errorf("tokenizer not found: %w", err)
}
// Initialize ONNX runtime
if err := onnx.Init(); err != nil {
return nil, fmt.Errorf("ONNX init failed: %w", err)
}
if onnx.HasCUDASupport() {
logger.Info("ONNX CUDA support enabled")
} else {
logger.Info("ONNX using CPU fallback")
}
emb := &ONNXEmbedder{
tokenizerPath: tokenizerPath,
dims: dims,
logger: logger,
modelPath: modelPath,
}
return emb, nil
}
func (e *ONNXEmbedder) ensureInitialized() error {
if e.session != nil {
return nil
}
e.mu.Lock()
defer e.mu.Unlock()
if e.session != nil {
return nil
}
// Load tokenizer lazily
if e.tokenizer == nil {
tok, err := pretrained.FromFile(e.tokenizerPath)
if err != nil {
return fmt.Errorf("failed to load tokenizer: %w", err)
}
e.tokenizer = tok
}
// ONNX runtime already initialized by onnx.Init() in NewONNXEmbedder
if !onnx.IsReady() {
return errors.New("ONNX runtime not ready")
}
// Create session options
opts, err := onnx.NewSessionOptions()
if err != nil {
return fmt.Errorf("failed to create session options: %w", err)
}
defer func() {
_ = opts.Destroy()
}()
if onnx.HasCUDASupport() {
e.logger.Info("Using CUDA for ONNX inference")
} else {
e.logger.Info("Using CPU for ONNX inference")
}
// Create session with options
session, err := onnxruntime_go.NewDynamicAdvancedSession(
e.getModelPath(),
[]string{"input_ids", "attention_mask"},
[]string{"sentence_embedding"},
opts,
)
if err != nil {
return fmt.Errorf("failed to create ONNX session: %w", err)
}
e.session = session
return nil
}
func (e *ONNXEmbedder) getModelPath() string {
return e.modelPath
}
func (e *ONNXEmbedder) Destroy() error {
e.mu.Lock()
defer e.mu.Unlock()
if e.session != nil {
if err := e.session.Destroy(); err != nil {
return fmt.Errorf("failed to destroy ONNX session: %w", err)
}
e.session = nil
e.logger.Info("ONNX session destroyed, VRAM freed")
}
return nil
}
func (e *ONNXEmbedder) Embed(text string) ([]float32, error) {
if err := e.ensureInitialized(); err != nil {
return nil, err
}
// 1. Tokenize
encoding, err := e.tokenizer.EncodeSingle(text)
if err != nil {
return nil, fmt.Errorf("tokenization failed: %w", err)
}
// 2. Convert to int64 and create attention mask
ids := encoding.Ids
inputIDs := make([]int64, len(ids))
attentionMask := make([]int64, len(ids))
for i, id := range ids {
inputIDs[i] = int64(id)
attentionMask[i] = 1
}
// 3. Create input tensors (shape: [1, seq_len])
seqLen := int64(len(inputIDs))
inputIDsTensor, err := onnxruntime_go.NewTensor[int64](
onnxruntime_go.NewShape(1, seqLen),
inputIDs,
)
if err != nil {
return nil, fmt.Errorf("failed to create input_ids tensor: %w", err)
}
defer func() { _ = inputIDsTensor.Destroy() }()
maskTensor, err := onnxruntime_go.NewTensor[int64](
onnxruntime_go.NewShape(1, seqLen),
attentionMask,
)
if err != nil {
return nil, fmt.Errorf("failed to create attention_mask tensor: %w", err)
}
defer func() { _ = maskTensor.Destroy() }()
// 4. Create output tensor
outputTensor, err := onnxruntime_go.NewEmptyTensor[float32](
onnxruntime_go.NewShape(1, int64(e.dims)),
)
if err != nil {
return nil, fmt.Errorf("failed to create output tensor: %w", err)
}
defer func() { _ = outputTensor.Destroy() }()
// 5. Run inference
err = e.session.Run(
[]onnxruntime_go.Value{inputIDsTensor, maskTensor},
[]onnxruntime_go.Value{outputTensor},
)
if err != nil {
return nil, fmt.Errorf("inference failed: %w", err)
}
// 6. Copy output data
outputData := outputTensor.GetData()
embedding := make([]float32, len(outputData))
copy(embedding, outputData)
return embedding, nil
}
func (e *ONNXEmbedder) EmbedSlice(texts []string) ([][]float32, error) {
if err := e.ensureInitialized(); err != nil {
return nil, err
}
encodings := make([]*tokenizer.Encoding, len(texts))
maxLen := 0
for i, txt := range texts {
enc, err := e.tokenizer.EncodeSingle(txt)
if err != nil {
return nil, err
}
encodings[i] = enc
if l := len(enc.Ids); l > maxLen {
maxLen = l
}
}
batchSize := len(texts)
inputIDs := make([]int64, batchSize*maxLen)
attentionMask := make([]int64, batchSize*maxLen)
for i, enc := range encodings {
ids := enc.Ids
offset := i * maxLen
for j, id := range ids {
inputIDs[offset+j] = int64(id)
attentionMask[offset+j] = 1
}
// Remaining positions are already zero (padding)
}
// Create tensors with shape [batchSize, maxLen]
inputTensor, _ := onnxruntime_go.NewTensor[int64](
onnxruntime_go.NewShape(int64(batchSize), int64(maxLen)),
inputIDs,
)
defer func() { _ = inputTensor.Destroy() }()
maskTensor, _ := onnxruntime_go.NewTensor[int64](
onnxruntime_go.NewShape(int64(batchSize), int64(maxLen)),
attentionMask,
)
defer func() { _ = maskTensor.Destroy() }()
outputTensor, _ := onnxruntime_go.NewEmptyTensor[float32](
onnxruntime_go.NewShape(int64(batchSize), int64(e.dims)),
)
defer func() { _ = outputTensor.Destroy() }()
err := e.session.Run(
[]onnxruntime_go.Value{inputTensor, maskTensor},
[]onnxruntime_go.Value{outputTensor},
)
if err != nil {
return nil, err
}
// Extract embeddings per batch item
data := outputTensor.GetData()
embeddings := make([][]float32, batchSize)
for i := 0; i < batchSize; i++ {
start := i * e.dims
emb := make([]float32, e.dims)
copy(emb, data[start:start+e.dims])
embeddings[i] = emb
}
return embeddings, nil
}

View File

@@ -1,6 +1,7 @@
package rag
import (
"context"
"errors"
"fmt"
"gf-lt/config"
@@ -9,16 +10,20 @@ import (
"log/slog"
"path"
"regexp"
"runtime"
"sort"
"strings"
"sync"
"time"
"github.com/neurosnap/sentences/english"
)
const ()
var (
// Status messages for TUI integration
LongJobStatusCh = make(chan string, 10) // Increased buffer size to prevent blocking
LongJobStatusCh = make(chan string, 100) // Increased buffer size for parallel batch updates
FinishedRAGStatus = "finished loading RAG file; press Enter"
LoadedFileRAGStatus = "loaded file"
ErrRAGStatus = "some error occurred; failed to transfer data to vector db"
@@ -30,30 +35,143 @@ type RAG struct {
cfg *config.Config
embedder Embedder
storage *VectorStorage
mu sync.Mutex
mu sync.RWMutex
idleMu sync.Mutex
fallbackMsg string
idleTimer *time.Timer
idleTimeout time.Duration
}
func New(l *slog.Logger, s storage.FullRepo, cfg *config.Config) *RAG {
// Initialize with API embedder by default, could be configurable later
embedder := NewAPIEmbedder(l, cfg)
// batchTask represents a single batch to be embedded
type batchTask struct {
batchIndex int
paragraphs []string
filename string
totalBatches int
}
// batchResult represents the result of embedding a batch
type batchResult struct {
batchIndex int
embeddings [][]float32
paragraphs []string
filename string
}
// sendStatusNonBlocking sends a status message without blocking
func (r *RAG) sendStatusNonBlocking(status string) {
select {
case LongJobStatusCh <- status:
default:
r.logger.Warn("LongJobStatusCh channel is full or closed, dropping status message", "message", status)
}
}
func New(l *slog.Logger, s storage.FullRepo, cfg *config.Config) (*RAG, error) {
var embedder Embedder
var fallbackMsg string
if cfg.EmbedModelPath != "" && cfg.EmbedTokenizerPath != "" {
emb, err := NewONNXEmbedder(cfg.EmbedModelPath, cfg.EmbedTokenizerPath, cfg.EmbedDims, l)
if err != nil {
l.Error("failed to create ONNX embedder, falling back to API", "error", err)
fallbackMsg = err.Error()
embedder = NewAPIEmbedder(l, cfg)
} else {
embedder = emb
l.Info("using ONNX embedder", "model", cfg.EmbedModelPath, "dims", cfg.EmbedDims)
}
} else {
embedder = NewAPIEmbedder(l, cfg)
l.Info("using API embedder", "url", cfg.EmbedURL)
}
rag := &RAG{
logger: l,
store: s,
cfg: cfg,
embedder: embedder,
storage: NewVectorStorage(l, s),
fallbackMsg: fallbackMsg,
idleTimeout: 30 * time.Second,
}
// Note: Vector tables are created via database migrations, not at runtime
return rag
return rag, nil
}
func wordCounter(sentence string) int {
return len(strings.Split(strings.TrimSpace(sentence), " "))
func createChunks(sentences []string, wordLimit, overlapWords uint32) []string {
if len(sentences) == 0 {
return nil
}
if overlapWords >= wordLimit {
overlapWords = wordLimit / 2
}
var chunks []string
i := 0
for i < len(sentences) {
var chunkWords []string
wordCount := 0
j := i
for j < len(sentences) && wordCount <= int(wordLimit) {
sentence := sentences[j]
words := strings.Fields(sentence)
chunkWords = append(chunkWords, sentence)
wordCount += len(words)
j++
// If this sentence alone exceeds limit, still include it and stop
if wordCount > int(wordLimit) {
break
}
}
if len(chunkWords) == 0 {
break
}
chunk := strings.Join(chunkWords, " ")
chunks = append(chunks, chunk)
if j >= len(sentences) {
break
}
// Move i forward by skipping overlap
if overlapWords == 0 {
i = j
continue
}
// Calculate how many sentences to skip to achieve overlapWords
overlapRemaining := int(overlapWords)
newI := i
for newI < j && overlapRemaining > 0 {
words := len(strings.Fields(sentences[newI]))
overlapRemaining -= words
if overlapRemaining >= 0 {
newI++
}
}
if newI == i {
newI = j
}
i = newI
}
return chunks
}
func sanitizeFTSQuery(query string) string {
// Remove double quotes and other problematic characters for FTS5
query = strings.ReplaceAll(query, "\"", " ")
query = strings.ReplaceAll(query, "'", " ")
query = strings.ReplaceAll(query, ";", " ")
query = strings.ReplaceAll(query, "\\", " ")
query = strings.TrimSpace(query)
if query == "" {
return "*" // match all
}
return query
}
func (r *RAG) LoadRAG(fpath string) error {
return r.LoadRAGWithContext(context.Background(), fpath)
}
func (r *RAG) LoadRAGWithContext(ctx context.Context, fpath string) error {
r.mu.Lock()
defer r.mu.Unlock()
fileText, err := ExtractText(fpath)
@@ -61,11 +179,9 @@ func (r *RAG) LoadRAG(fpath string) error {
return err
}
r.logger.Debug("rag: loaded file", "fp", fpath)
select {
case LongJobStatusCh <- LoadedFileRAGStatus:
default:
r.logger.Warn("LongJobStatusCh channel is full or closed, dropping status message", "message", LoadedFileRAGStatus)
}
// Send initial status (non-blocking with retry)
r.sendStatusNonBlocking(LoadedFileRAGStatus)
tokenizer, err := english.NewSentenceTokenizer(nil)
if err != nil {
return err
@@ -75,31 +191,9 @@ func (r *RAG) LoadRAG(fpath string) error {
for i, s := range sentences {
sents[i] = s.Text
}
// Group sentences into paragraphs based on word limit
paragraphs := []string{}
par := strings.Builder{}
for i := 0; i < len(sents); i++ {
if strings.TrimSpace(sents[i]) != "" {
if par.Len() > 0 {
par.WriteString(" ")
}
par.WriteString(sents[i])
}
if wordCounter(par.String()) > int(r.cfg.RAGWordLimit) {
paragraph := strings.TrimSpace(par.String())
if paragraph != "" {
paragraphs = append(paragraphs, paragraph)
}
par.Reset()
}
}
// Handle any remaining content in the paragraph buffer
if par.Len() > 0 {
paragraph := strings.TrimSpace(par.String())
if paragraph != "" {
paragraphs = append(paragraphs, paragraph)
}
}
// Create chunks with overlap
paragraphs := createChunks(sents, r.cfg.RAGWordLimit, r.cfg.RAGOverlapWords)
// Adjust batch size if needed
if len(paragraphs) < r.cfg.RAGBatchSize && len(paragraphs) > 0 {
r.cfg.RAGBatchSize = len(paragraphs)
@@ -107,15 +201,61 @@ func (r *RAG) LoadRAG(fpath string) error {
if len(paragraphs) == 0 {
return errors.New("no valid paragraphs found in file")
}
// Process paragraphs in batches synchronously
batchCount := 0
for i := 0; i < len(paragraphs); i += r.cfg.RAGBatchSize {
end := i + r.cfg.RAGBatchSize
totalBatches := (len(paragraphs) + r.cfg.RAGBatchSize - 1) / r.cfg.RAGBatchSize
r.logger.Debug("starting parallel embedding", "total_batches", totalBatches, "batch_size", r.cfg.RAGBatchSize)
// Determine concurrency level
concurrency := runtime.NumCPU()
if concurrency > totalBatches {
concurrency = totalBatches
}
if concurrency < 1 {
concurrency = 1
}
// If using ONNX embedder, limit concurrency to 1 due to mutex serialization
var isONNX bool
if _, isONNX = r.embedder.(*ONNXEmbedder); isONNX {
concurrency = 1
}
embedderType := "API"
if isONNX {
embedderType = "ONNX"
}
r.logger.Debug("parallel embedding setup",
"total_batches", totalBatches,
"concurrency", concurrency,
"embedder", embedderType,
"batch_size", r.cfg.RAGBatchSize)
// Create context with timeout (30 minutes) and cancellation for error handling
ctx, cancel := context.WithTimeout(ctx, 30*time.Minute)
defer cancel()
// Channels for task distribution and results
taskCh := make(chan batchTask, totalBatches)
resultCh := make(chan batchResult, totalBatches)
errorCh := make(chan error, totalBatches)
// Start worker goroutines
var wg sync.WaitGroup
for w := 0; w < concurrency; w++ {
wg.Add(1)
go r.embeddingWorker(ctx, w, taskCh, resultCh, errorCh, &wg)
}
// Close task channel after all tasks are sent (by separate goroutine)
go func() {
// Ensure task channel is closed when this goroutine exits
defer close(taskCh)
r.logger.Debug("task distributor started", "total_batches", totalBatches)
for i := 0; i < totalBatches; i++ {
start := i * r.cfg.RAGBatchSize
end := start + r.cfg.RAGBatchSize
if end > len(paragraphs) {
end = len(paragraphs)
}
batch := paragraphs[i:end]
batchCount++
batch := paragraphs[start:end]
// Filter empty paragraphs
nonEmptyBatch := make([]string, 0, len(batch))
for _, p := range batch {
@@ -123,75 +263,286 @@ func (r *RAG) LoadRAG(fpath string) error {
nonEmptyBatch = append(nonEmptyBatch, strings.TrimSpace(p))
}
}
if len(nonEmptyBatch) == 0 {
task := batchTask{
batchIndex: i,
paragraphs: nonEmptyBatch,
filename: path.Base(fpath),
totalBatches: totalBatches,
}
select {
case taskCh <- task:
r.logger.Debug("task distributor sent batch", "batch", i, "paragraphs", len(nonEmptyBatch))
case <-ctx.Done():
r.logger.Debug("task distributor cancelled", "batches_sent", i+1, "total_batches", totalBatches)
return
}
}
r.logger.Debug("task distributor finished", "batches_sent", totalBatches)
}()
// Wait for workers to finish and close result channel
go func() {
wg.Wait()
close(resultCh)
}()
// Process results in order and write to database
nextExpectedBatch := 0
resultsBuffer := make(map[int]batchResult)
filename := path.Base(fpath)
batchesProcessed := 0
for {
select {
case <-ctx.Done():
return ctx.Err()
case err := <-errorCh:
// First error from any worker, cancel everything
cancel()
r.logger.Error("embedding worker failed", "error", err)
r.sendStatusNonBlocking(ErrRAGStatus)
return fmt.Errorf("embedding failed: %w", err)
case result, ok := <-resultCh:
if !ok {
// All results processed
resultCh = nil
r.logger.Debug("result channel closed", "batches_processed", batchesProcessed, "total_batches", totalBatches)
continue
}
// Embed the batch
embeddings, err := r.embedder.EmbedSlice(nonEmptyBatch)
if err != nil {
r.logger.Error("failed to embed batch", "error", err, "batch", batchCount)
select {
case LongJobStatusCh <- ErrRAGStatus:
default:
r.logger.Warn("LongJobStatusCh channel full, dropping message")
}
return fmt.Errorf("failed to embed batch %d: %w", batchCount, err)
}
if len(embeddings) != len(nonEmptyBatch) {
err := errors.New("embedding count mismatch")
r.logger.Error("embedding mismatch", "expected", len(nonEmptyBatch), "got", len(embeddings))
// Store result in buffer
resultsBuffer[result.batchIndex] = result
// Process buffered results in order
for {
if res, exists := resultsBuffer[nextExpectedBatch]; exists {
// Write this batch to database
if err := r.writeBatchToStorage(ctx, res, filename); err != nil {
cancel()
return err
}
// Write vectors to storage
filename := path.Base(fpath)
for j, text := range nonEmptyBatch {
vector := models.VectorRow{
Embeddings: embeddings[j],
batchesProcessed++
// Send progress update
statusMsg := fmt.Sprintf("processed batch %d/%d", batchesProcessed, totalBatches)
r.sendStatusNonBlocking(statusMsg)
delete(resultsBuffer, nextExpectedBatch)
nextExpectedBatch++
} else {
break
}
}
default:
// No channels ready, check for deadlock conditions
if resultCh == nil && nextExpectedBatch < totalBatches {
// Missing batch results after result channel closed
r.logger.Error("missing batch results",
"expected", totalBatches,
"received", nextExpectedBatch,
"missing", totalBatches-nextExpectedBatch)
// Wait a short time for any delayed errors, then cancel
select {
case <-time.After(5 * time.Second):
cancel()
return fmt.Errorf("missing batch results: expected %d, got %d", totalBatches, nextExpectedBatch)
case <-ctx.Done():
return ctx.Err()
case err := <-errorCh:
cancel()
r.logger.Error("embedding worker failed after result channel closed", "error", err)
r.sendStatusNonBlocking(ErrRAGStatus)
return fmt.Errorf("embedding failed: %w", err)
}
}
// If we reach here, no deadlock yet, just busy loop prevention
time.Sleep(100 * time.Millisecond)
}
// Check if we're done
if resultCh == nil && nextExpectedBatch >= totalBatches {
r.logger.Debug("all batches processed successfully", "total", totalBatches)
break
}
}
r.logger.Debug("finished writing vectors", "batches", batchesProcessed)
r.resetIdleTimer()
r.sendStatusNonBlocking(FinishedRAGStatus)
return nil
}
// embeddingWorker processes batch embedding tasks
func (r *RAG) embeddingWorker(ctx context.Context, workerID int, taskCh <-chan batchTask, resultCh chan<- batchResult, errorCh chan<- error, wg *sync.WaitGroup) {
defer wg.Done()
r.logger.Debug("embedding worker started", "worker", workerID)
// Panic recovery to ensure worker doesn't crash silently
defer func() {
if rec := recover(); rec != nil {
r.logger.Error("embedding worker panicked", "worker", workerID, "panic", rec)
// Try to send error, but don't block if channel is full
select {
case errorCh <- fmt.Errorf("worker %d panicked: %v", workerID, rec):
default:
r.logger.Warn("error channel full, dropping panic error", "worker", workerID)
}
}
}()
for task := range taskCh {
select {
case <-ctx.Done():
r.logger.Debug("embedding worker cancelled", "worker", workerID)
return
default:
}
r.logger.Debug("worker processing batch", "worker", workerID, "batch", task.batchIndex, "paragraphs", len(task.paragraphs), "total_batches", task.totalBatches)
// Skip empty batches
if len(task.paragraphs) == 0 {
select {
case resultCh <- batchResult{
batchIndex: task.batchIndex,
embeddings: nil,
paragraphs: nil,
filename: task.filename,
}:
case <-ctx.Done():
r.logger.Debug("embedding worker cancelled while sending empty batch", "worker", workerID)
return
}
r.logger.Debug("worker sent empty batch", "worker", workerID, "batch", task.batchIndex)
continue
}
// Embed with retry for API embedder
embeddings, err := r.embedWithRetry(ctx, task.paragraphs, 3)
if err != nil {
// Try to send error, but don't block indefinitely
select {
case errorCh <- fmt.Errorf("worker %d batch %d: %w", workerID, task.batchIndex, err):
case <-ctx.Done():
r.logger.Debug("embedding worker cancelled while sending error", "worker", workerID)
}
return
}
// Send result with context awareness
select {
case resultCh <- batchResult{
batchIndex: task.batchIndex,
embeddings: embeddings,
paragraphs: task.paragraphs,
filename: task.filename,
}:
case <-ctx.Done():
r.logger.Debug("embedding worker cancelled while sending result", "worker", workerID)
return
}
r.logger.Debug("worker completed batch", "worker", workerID, "batch", task.batchIndex, "embeddings", len(embeddings))
}
r.logger.Debug("embedding worker finished", "worker", workerID)
}
// embedWithRetry attempts embedding with exponential backoff for API embedder
func (r *RAG) embedWithRetry(ctx context.Context, paragraphs []string, maxRetries int) ([][]float32, error) {
var lastErr error
for attempt := 0; attempt < maxRetries; attempt++ {
if attempt > 0 {
// Exponential backoff
backoff := time.Duration(attempt*attempt) * time.Second
if backoff > 10*time.Second {
backoff = 10 * time.Second
}
select {
case <-time.After(backoff):
case <-ctx.Done():
return nil, ctx.Err()
}
r.logger.Debug("retrying embedding", "attempt", attempt, "max_retries", maxRetries)
}
embeddings, err := r.embedder.EmbedSlice(paragraphs)
if err == nil {
// Validate embedding count
if len(embeddings) != len(paragraphs) {
return nil, fmt.Errorf("embedding count mismatch: expected %d, got %d", len(paragraphs), len(embeddings))
}
return embeddings, nil
}
lastErr = err
// Only retry for API embedder errors (network/timeout)
// For ONNX embedder, fail fast
if _, isAPI := r.embedder.(*APIEmbedder); !isAPI {
break
}
}
return nil, fmt.Errorf("embedding failed after %d attempts: %w", maxRetries, lastErr)
}
// writeBatchToStorage writes a single batch of vectors to the database
func (r *RAG) writeBatchToStorage(ctx context.Context, result batchResult, filename string) error {
if len(result.embeddings) == 0 {
// Empty batch, skip
return nil
}
// Check context before starting
select {
case <-ctx.Done():
return ctx.Err()
default:
}
// Build all vectors for batch write
vectors := make([]*models.VectorRow, 0, len(result.paragraphs))
for j, text := range result.paragraphs {
vectors = append(vectors, &models.VectorRow{
Embeddings: result.embeddings[j],
RawText: text,
Slug: fmt.Sprintf("%s_%d_%d", filename, batchCount, j),
Slug: fmt.Sprintf("%s_%d_%d", filename, result.batchIndex+1, j),
FileName: filename,
})
}
if err := r.storage.WriteVector(&vector); err != nil {
r.logger.Error("failed to write vector to DB", "error", err, "slug", vector.Slug)
select {
case LongJobStatusCh <- ErrRAGStatus:
default:
r.logger.Warn("LongJobStatusCh channel full, dropping message")
}
return fmt.Errorf("failed to write vector: %w", err)
}
}
r.logger.Debug("wrote batch to db", "batch", batchCount, "size", len(nonEmptyBatch))
// Send progress status
statusMsg := fmt.Sprintf("processed batch %d/%d", batchCount, (len(paragraphs)+r.cfg.RAGBatchSize-1)/r.cfg.RAGBatchSize)
select {
case LongJobStatusCh <- statusMsg:
default:
r.logger.Warn("LongJobStatusCh channel full, dropping message")
}
}
r.logger.Debug("finished writing vectors", "batches", batchCount)
select {
case LongJobStatusCh <- FinishedRAGStatus:
default:
r.logger.Warn("LongJobStatusCh channel is full or closed, dropping status message", "message", FinishedRAGStatus)
// Write all vectors in a single transaction
if err := r.storage.WriteVectors(vectors); err != nil {
r.logger.Error("failed to write vectors batch to DB", "error", err, "batch", result.batchIndex+1, "size", len(vectors))
r.sendStatusNonBlocking(ErrRAGStatus)
return fmt.Errorf("failed to write vectors batch: %w", err)
}
r.logger.Debug("wrote batch to db", "batch", result.batchIndex+1, "size", len(result.paragraphs))
return nil
}
func (r *RAG) LineToVector(line string) ([]float32, error) {
r.resetIdleTimer()
return r.embedder.Embed(line)
}
func (r *RAG) SearchEmb(emb *models.EmbeddingResp) ([]models.VectorRow, error) {
return r.storage.SearchClosest(emb.Embedding)
func (r *RAG) searchEmb(emb *models.EmbeddingResp, limit int) ([]models.VectorRow, error) {
r.resetIdleTimer()
return r.storage.SearchClosest(emb.Embedding, limit)
}
func (r *RAG) searchKeyword(query string, limit int) ([]models.VectorRow, error) {
r.resetIdleTimer()
sanitized := sanitizeFTSQuery(query)
return r.storage.SearchKeyword(sanitized, limit)
}
func (r *RAG) ListLoaded() ([]string, error) {
r.mu.RLock()
defer r.mu.RUnlock()
return r.storage.ListFiles()
}
func (r *RAG) RemoveFile(filename string) error {
r.mu.Lock()
defer r.mu.Unlock()
r.resetIdleTimer()
return r.storage.RemoveEmbByFileName(filename)
}
@@ -211,11 +562,14 @@ func (r *RAG) RefineQuery(query string) string {
return original
}
query = strings.ToLower(query)
words := strings.Fields(query)
if len(words) >= 3 {
for _, stopWord := range stopWords {
wordPattern := `\b` + stopWord + `\b`
re := regexp.MustCompile(wordPattern)
query = re.ReplaceAllString(query, "")
}
}
query = strings.TrimSpace(query)
if len(query) < 5 {
return original
@@ -246,7 +600,7 @@ func (r *RAG) extractImportantPhrases(query string) string {
break
}
}
if isImportant || len(word) > 3 {
if isImportant || len(word) >= 3 {
important = append(important, word)
}
}
@@ -265,6 +619,36 @@ func (r *RAG) GenerateQueryVariations(query string) []string {
if len(parts) == 0 {
return variations
}
// Get loaded filenames to filter out filename terms
filenames, err := r.storage.ListFiles()
if err == nil && len(filenames) > 0 {
// Convert to lowercase for case-insensitive matching
lowerFilenames := make([]string, len(filenames))
for i, f := range filenames {
lowerFilenames[i] = strings.ToLower(f)
}
filteredParts := make([]string, 0, len(parts))
for _, part := range parts {
partLower := strings.ToLower(part)
skip := false
for _, fn := range lowerFilenames {
if strings.Contains(fn, partLower) || strings.Contains(partLower, fn) {
skip = true
break
}
}
if !skip {
filteredParts = append(filteredParts, part)
}
}
// If filteredParts not empty and different from original, add filtered query
if len(filteredParts) > 0 && len(filteredParts) != len(parts) {
filteredQuery := strings.Join(filteredParts, " ")
if len(filteredQuery) >= 5 {
variations = append(variations, filteredQuery)
}
}
}
if len(parts) >= 2 {
trimmed := strings.Join(parts[:len(parts)-1], " ")
if len(trimmed) >= 5 {
@@ -328,9 +712,14 @@ func (r *RAG) RerankResults(results []models.VectorRow, query string) []models.V
})
unique := make([]models.VectorRow, 0)
seen := make(map[string]bool)
fileCounts := make(map[string]int)
for i := range scored {
if !seen[scored[i].row.Slug] {
if fileCounts[scored[i].row.FileName] >= 2 {
continue
}
seen[scored[i].row.Slug] = true
fileCounts[scored[i].row.FileName]++
unique = append(unique, scored[i].row)
}
}
@@ -341,6 +730,9 @@ func (r *RAG) RerankResults(results []models.VectorRow, query string) []models.V
}
func (r *RAG) SynthesizeAnswer(results []models.VectorRow, query string) (string, error) {
r.mu.RLock()
defer r.mu.RUnlock()
r.resetIdleTimer()
if len(results) == 0 {
return "No relevant information found in the vector database.", nil
}
@@ -369,7 +761,7 @@ func (r *RAG) SynthesizeAnswer(results []models.VectorRow, query string) (string
Embedding: emb,
Index: 0,
}
topResults, err := r.SearchEmb(embResp)
topResults, err := r.searchEmb(embResp, 1)
if err != nil {
r.logger.Error("failed to search for synthesis context", "error", err)
return "", err
@@ -396,9 +788,14 @@ func truncateString(s string, maxLen int) string {
}
func (r *RAG) Search(query string, limit int) ([]models.VectorRow, error) {
r.mu.RLock()
defer r.mu.RUnlock()
r.resetIdleTimer()
refined := r.RefineQuery(query)
variations := r.GenerateQueryVariations(refined)
allResults := make([]models.VectorRow, 0)
// Collect embedding search results from all variations
var embResults []models.VectorRow
seen := make(map[string]bool)
for _, q := range variations {
emb, err := r.LineToVector(q)
@@ -406,29 +803,78 @@ func (r *RAG) Search(query string, limit int) ([]models.VectorRow, error) {
r.logger.Error("failed to embed query variation", "error", err, "query", q)
continue
}
embResp := &models.EmbeddingResp{
Embedding: emb,
Index: 0,
}
results, err := r.SearchEmb(embResp)
results, err := r.searchEmb(embResp, limit*2) // Get more candidates
if err != nil {
r.logger.Error("failed to search embeddings", "error", err, "query", q)
continue
}
for _, row := range results {
if !seen[row.Slug] {
seen[row.Slug] = true
allResults = append(allResults, row)
embResults = append(embResults, row)
}
}
}
reranked := r.RerankResults(allResults, query)
if len(reranked) > limit {
reranked = reranked[:limit]
// Sort embedding results by distance (lower is better)
sort.Slice(embResults, func(i, j int) bool {
return embResults[i].Distance < embResults[j].Distance
})
// Perform keyword search
kwResults, err := r.searchKeyword(refined, limit*2)
if err != nil {
r.logger.Warn("keyword search failed, using only embeddings", "error", err)
kwResults = nil
}
// Sort keyword results by distance (already sorted by BM25 score)
// kwResults already sorted by distance (lower is better)
// Combine using Reciprocal Rank Fusion (RRF)
const rrfK = 60
type scoredRow struct {
row models.VectorRow
score float64
}
scoreMap := make(map[string]float64)
// Add embedding results
for rank, row := range embResults {
score := 1.0 / (float64(rank) + rrfK)
scoreMap[row.Slug] += score
}
// Add keyword results
for rank, row := range kwResults {
score := 1.0 / (float64(rank) + rrfK)
scoreMap[row.Slug] += score
// Ensure row exists in combined results
if _, exists := seen[row.Slug]; !exists {
embResults = append(embResults, row)
}
}
// Create slice of scored rows
scoredRows := make([]scoredRow, 0, len(embResults))
for _, row := range embResults {
score := scoreMap[row.Slug]
scoredRows = append(scoredRows, scoredRow{row: row, score: score})
}
// Sort by descending RRF score
sort.Slice(scoredRows, func(i, j int) bool {
return scoredRows[i].score > scoredRows[j].score
})
// Take top limit
if len(scoredRows) > limit {
scoredRows = scoredRows[:limit]
}
// Convert back to VectorRow
finalResults := make([]models.VectorRow, len(scoredRows))
for i, sr := range scoredRows {
finalResults[i] = sr.row
}
// Apply reranking heuristics
reranked := r.RerankResults(finalResults, query)
return reranked, nil
}
@@ -437,16 +883,58 @@ var (
ragOnce sync.Once
)
func (r *RAG) FallbackMessage() string {
return r.fallbackMsg
}
func Init(c *config.Config, l *slog.Logger, s storage.FullRepo) error {
var err error
ragOnce.Do(func() {
if c == nil || l == nil || s == nil {
return
}
ragInstance = New(l, s, c)
ragInstance, err = New(l, s, c)
})
return nil
return err
}
func GetInstance() *RAG {
return ragInstance
}
func (r *RAG) resetIdleTimer() {
r.idleMu.Lock()
defer r.idleMu.Unlock()
if r.idleTimer != nil {
r.idleTimer.Stop()
}
r.idleTimer = time.AfterFunc(r.idleTimeout, func() {
r.freeONNXMemory()
})
}
func (r *RAG) freeONNXMemory() {
r.mu.Lock()
defer r.mu.Unlock()
if onnx, ok := r.embedder.(*ONNXEmbedder); ok {
if err := onnx.Destroy(); err != nil {
r.logger.Error("failed to free ONNX memory", "error", err)
} else {
r.logger.Info("freed ONNX VRAM after idle timeout")
}
}
}
func (r *RAG) Destroy() {
r.mu.Lock()
defer r.mu.Unlock()
if r.idleTimer != nil {
r.idleTimer.Stop()
r.idleTimer = nil
}
if onnx, ok := r.embedder.(*ONNXEmbedder); ok {
if err := onnx.Destroy(); err != nil {
r.logger.Error("failed to destroy ONNX embedder", "error", err)
}
}
}

View File

@@ -1,6 +1,7 @@
package rag
import (
"database/sql"
"encoding/binary"
"fmt"
"gf-lt/models"
@@ -62,6 +63,17 @@ func (vs *VectorStorage) WriteVector(row *models.VectorRow) error {
if err != nil {
return err
}
embeddingSize := len(row.Embeddings)
// Start transaction
tx, err := vs.sqlxDB.Beginx()
if err != nil {
return err
}
defer func() {
if err != nil {
_ = tx.Rollback()
}
}()
// Serialize the embeddings to binary
serializedEmbeddings := SerializeVector(row.Embeddings)
@@ -69,10 +81,102 @@ func (vs *VectorStorage) WriteVector(row *models.VectorRow) error {
"INSERT INTO %s (embeddings, slug, raw_text, filename) VALUES (?, ?, ?, ?)",
tableName,
)
if _, err := vs.sqlxDB.Exec(query, serializedEmbeddings, row.Slug, row.RawText, row.FileName); err != nil {
if _, err := tx.Exec(query, serializedEmbeddings, row.Slug, row.RawText, row.FileName); err != nil {
vs.logger.Error("failed to write vector", "error", err, "slug", row.Slug)
return err
}
// Insert into FTS table
ftsQuery := `INSERT INTO fts_embeddings (slug, raw_text, filename, embedding_size) VALUES (?, ?, ?, ?)`
if _, err := tx.Exec(ftsQuery, row.Slug, row.RawText, row.FileName, embeddingSize); err != nil {
vs.logger.Error("failed to write to FTS table", "error", err, "slug", row.Slug)
return err
}
err = tx.Commit()
if err != nil {
vs.logger.Error("failed to commit transaction", "error", err)
return err
}
return nil
}
// WriteVectors stores multiple embedding vectors in a single transaction
func (vs *VectorStorage) WriteVectors(rows []*models.VectorRow) error {
if len(rows) == 0 {
return nil
}
// SQLite has limit of 999 parameters per statement, each row uses 4 parameters
const maxBatchSize = 200 // 200 * 4 = 800 < 999
if len(rows) > maxBatchSize {
// Process in chunks
for i := 0; i < len(rows); i += maxBatchSize {
end := i + maxBatchSize
if end > len(rows) {
end = len(rows)
}
if err := vs.WriteVectors(rows[i:end]); err != nil {
return err
}
}
return nil
}
// All rows should have same embedding size (same model)
firstSize := len(rows[0].Embeddings)
for i, row := range rows {
if len(row.Embeddings) != firstSize {
return fmt.Errorf("embedding size mismatch: row %d has size %d, expected %d", i, len(row.Embeddings), firstSize)
}
}
tableName, err := vs.getTableName(rows[0].Embeddings)
if err != nil {
return err
}
// Start transaction
tx, err := vs.sqlxDB.Beginx()
if err != nil {
return err
}
defer func() {
if err != nil {
_ = tx.Rollback()
}
}()
// Build batch insert for embeddings table
embeddingPlaceholders := make([]string, 0, len(rows))
embeddingArgs := make([]any, 0, len(rows)*4)
for _, row := range rows {
embeddingPlaceholders = append(embeddingPlaceholders, "(?, ?, ?, ?)")
embeddingArgs = append(embeddingArgs, SerializeVector(row.Embeddings), row.Slug, row.RawText, row.FileName)
}
embeddingQuery := fmt.Sprintf(
"INSERT INTO %s (embeddings, slug, raw_text, filename) VALUES %s",
tableName,
strings.Join(embeddingPlaceholders, ", "),
)
if _, err := tx.Exec(embeddingQuery, embeddingArgs...); err != nil {
vs.logger.Error("failed to write vectors batch", "error", err, "batch_size", len(rows))
return err
}
// Build batch insert for FTS table
ftsPlaceholders := make([]string, 0, len(rows))
ftsArgs := make([]any, 0, len(rows)*4)
embeddingSize := len(rows[0].Embeddings)
for _, row := range rows {
ftsPlaceholders = append(ftsPlaceholders, "(?, ?, ?, ?)")
ftsArgs = append(ftsArgs, row.Slug, row.RawText, row.FileName, embeddingSize)
}
ftsQuery := "INSERT INTO fts_embeddings (slug, raw_text, filename, embedding_size) VALUES " +
strings.Join(ftsPlaceholders, ", ")
if _, err := tx.Exec(ftsQuery, ftsArgs...); err != nil {
vs.logger.Error("failed to write FTS batch", "error", err, "batch_size", len(rows))
return err
}
err = tx.Commit()
if err != nil {
vs.logger.Error("failed to commit transaction", "error", err)
return err
}
vs.logger.Debug("wrote vectors batch", "batch_size", len(rows))
return nil
}
@@ -98,30 +202,25 @@ func (vs *VectorStorage) getTableName(emb []float32) (string, error) {
}
// SearchClosest finds vectors closest to the query vector using efficient cosine similarity calculation
func (vs *VectorStorage) SearchClosest(query []float32) ([]models.VectorRow, error) {
func (vs *VectorStorage) SearchClosest(query []float32, limit int) ([]models.VectorRow, error) {
if limit <= 0 {
limit = 10
}
tableName, err := vs.getTableName(query)
if err != nil {
return nil, err
}
// For better performance, instead of loading all vectors at once,
// we'll implement batching and potentially add L2 distance-based pre-filtering
// since cosine similarity is related to L2 distance for normalized vectors
querySQL := "SELECT embeddings, slug, raw_text, filename FROM " + tableName
rows, err := vs.sqlxDB.Query(querySQL)
if err != nil {
return nil, err
}
defer rows.Close()
// Use a min-heap or simple slice to keep track of top 3 closest vectors
type SearchResult struct {
vector models.VectorRow
distance float32
}
var topResults []SearchResult
// Process vectors one by one to avoid loading everything into memory
for rows.Next() {
var (
embeddingsBlob []byte
@@ -132,12 +231,9 @@ func (vs *VectorStorage) SearchClosest(query []float32) ([]models.VectorRow, err
vs.logger.Error("failed to scan row", "error", err)
continue
}
storedEmbeddings := DeserializeVector(embeddingsBlob)
// Calculate cosine similarity (returns value between -1 and 1, where 1 is most similar)
similarity := cosineSimilarity(query, storedEmbeddings)
distance := 1 - similarity // Convert to distance where 0 is most similar
distance := 1 - similarity
result := SearchResult{
vector: models.VectorRow{
@@ -149,20 +245,14 @@ func (vs *VectorStorage) SearchClosest(query []float32) ([]models.VectorRow, err
distance: distance,
}
// Add to top results and maintain only top 3
topResults = append(topResults, result)
// Sort and keep only top 3
sort.Slice(topResults, func(i, j int) bool {
return topResults[i].distance < topResults[j].distance
})
if len(topResults) > 3 {
topResults = topResults[:3] // Keep only closest 3
if len(topResults) > limit {
topResults = topResults[:limit]
}
}
// Convert back to VectorRow slice
results := make([]models.VectorRow, 0, len(topResults))
for _, result := range topResults {
result.vector.Distance = result.distance
@@ -171,6 +261,100 @@ func (vs *VectorStorage) SearchClosest(query []float32) ([]models.VectorRow, err
return results, nil
}
// GetVectorBySlug retrieves a vector row by its slug
func (vs *VectorStorage) GetVectorBySlug(slug string) (*models.VectorRow, error) {
embeddingSizes := []int{384, 768, 1024, 1536, 2048, 3072, 4096, 5120}
for _, size := range embeddingSizes {
table := fmt.Sprintf("embeddings_%d", size)
query := fmt.Sprintf("SELECT embeddings, slug, raw_text, filename FROM %s WHERE slug = ?", table)
row := vs.sqlxDB.QueryRow(query, slug)
var (
embeddingsBlob []byte
retrievedSlug, rawText, fileName string
)
if err := row.Scan(&embeddingsBlob, &retrievedSlug, &rawText, &fileName); err != nil {
// No row in this table, continue to next size
continue
}
storedEmbeddings := DeserializeVector(embeddingsBlob)
return &models.VectorRow{
Embeddings: storedEmbeddings,
Slug: retrievedSlug,
RawText: rawText,
FileName: fileName,
}, nil
}
return nil, fmt.Errorf("vector with slug %s not found", slug)
}
// SearchKeyword performs full-text search using FTS5
func (vs *VectorStorage) SearchKeyword(query string, limit int) ([]models.VectorRow, error) {
// Use FTS5 bm25 ranking. bm25 returns negative values where more negative is better.
// We'll order by bm25 (ascending) and limit.
ftsQuery := `SELECT slug, raw_text, filename, bm25(fts_embeddings) as score
FROM fts_embeddings
WHERE fts_embeddings MATCH ?
ORDER BY score
LIMIT ?`
// Try original query first
rows, err := vs.sqlxDB.Query(ftsQuery, query, limit)
if err != nil {
return nil, fmt.Errorf("FTS search failed: %w", err)
}
results, err := vs.scanRows(rows)
rows.Close()
if err != nil {
return nil, err
}
// If no results and query contains multiple terms, try OR fallback
if len(results) == 0 && strings.Contains(query, " ") && !strings.Contains(strings.ToUpper(query), " OR ") {
// Build OR query: term1 OR term2 OR term3
terms := strings.Fields(query)
if len(terms) > 1 {
orQuery := strings.Join(terms, " OR ")
rows, err := vs.sqlxDB.Query(ftsQuery, orQuery, limit)
if err != nil {
// Return original empty results rather than error
return results, nil
}
orResults, err := vs.scanRows(rows)
rows.Close()
if err == nil {
results = orResults
}
}
}
return results, nil
}
// scanRows converts SQL rows to VectorRow slice
func (vs *VectorStorage) scanRows(rows *sql.Rows) ([]models.VectorRow, error) {
var results []models.VectorRow
for rows.Next() {
var slug, rawText, fileName string
var score float64
if err := rows.Scan(&slug, &rawText, &fileName, &score); err != nil {
vs.logger.Error("failed to scan FTS row", "error", err)
continue
}
// Convert BM25 score to distance-like metric (lower is better)
// BM25 is negative, more negative is better. We'll normalize to positive distance.
distance := float32(-score) // Make positive (since score is negative)
if distance < 0 {
distance = 0
}
results = append(results, models.VectorRow{
Slug: slug,
RawText: rawText,
FileName: fileName,
Distance: distance,
})
}
return results, nil
}
// ListFiles returns a list of all loaded files
func (vs *VectorStorage) ListFiles() ([]string, error) {
fileLists := make([][]string, 0)
@@ -215,6 +399,10 @@ func (vs *VectorStorage) ListFiles() ([]string, error) {
// RemoveEmbByFileName removes all embeddings associated with a specific filename
func (vs *VectorStorage) RemoveEmbByFileName(filename string) error {
var errors []string
// Delete from FTS table first
if _, err := vs.sqlxDB.Exec("DELETE FROM fts_embeddings WHERE filename = ?", filename); err != nil {
errors = append(errors, err.Error())
}
embeddingSizes := []int{384, 768, 1024, 1536, 2048, 3072, 4096, 5120}
for _, size := range embeddingSizes {
table := fmt.Sprintf("embeddings_%d", size)

View File

@@ -0,0 +1,2 @@
-- Drop FTS5 virtual table
DROP TABLE IF EXISTS fts_embeddings;

View File

@@ -0,0 +1,15 @@
-- Create FTS5 virtual table for full-text search
CREATE VIRTUAL TABLE IF NOT EXISTS fts_embeddings USING fts5(
slug UNINDEXED,
raw_text,
filename UNINDEXED,
embedding_size UNINDEXED,
tokenize='porter unicode61' -- Use porter stemmer and unicode61 tokenizer
);
-- Create triggers to maintain FTS table when embeddings are inserted/deleted
-- Note: We'll handle inserts/deletes programmatically for simplicity
-- but triggers could be added here if needed.
-- Indexes for performance (FTS5 manages its own indexes)
-- No additional indexes needed for FTS5 virtual table.

View File

@@ -0,0 +1,2 @@
-- Clear FTS table (optional)
DELETE FROM fts_embeddings;

View File

@@ -0,0 +1,26 @@
-- Populate FTS table with existing embeddings
DELETE FROM fts_embeddings;
INSERT INTO fts_embeddings (slug, raw_text, filename, embedding_size)
SELECT slug, raw_text, filename, 384 FROM embeddings_384;
INSERT INTO fts_embeddings (slug, raw_text, filename, embedding_size)
SELECT slug, raw_text, filename, 768 FROM embeddings_768;
INSERT INTO fts_embeddings (slug, raw_text, filename, embedding_size)
SELECT slug, raw_text, filename, 1024 FROM embeddings_1024;
INSERT INTO fts_embeddings (slug, raw_text, filename, embedding_size)
SELECT slug, raw_text, filename, 1536 FROM embeddings_1536;
INSERT INTO fts_embeddings (slug, raw_text, filename, embedding_size)
SELECT slug, raw_text, filename, 2048 FROM embeddings_2048;
INSERT INTO fts_embeddings (slug, raw_text, filename, embedding_size)
SELECT slug, raw_text, filename, 3072 FROM embeddings_3072;
INSERT INTO fts_embeddings (slug, raw_text, filename, embedding_size)
SELECT slug, raw_text, filename, 4096 FROM embeddings_4096;
INSERT INTO fts_embeddings (slug, raw_text, filename, embedding_size)
SELECT slug, raw_text, filename, 5120 FROM embeddings_5120;

View File

@@ -102,6 +102,22 @@ func NewProviderSQL(dbPath string, logger *slog.Logger) FullRepo {
logger.Error("failed to open db connection", "error", err)
return nil
}
// Enable WAL mode for better concurrency and performance
if _, err := db.Exec("PRAGMA journal_mode = WAL;"); err != nil {
logger.Warn("failed to enable WAL mode", "error", err)
}
if _, err := db.Exec("PRAGMA synchronous = NORMAL;"); err != nil {
logger.Warn("failed to set synchronous mode", "error", err)
}
// Increase cache size for better performance
if _, err := db.Exec("PRAGMA cache_size = -2000;"); err != nil {
logger.Warn("failed to set cache size", "error", err)
}
// Log actual journal mode for debugging
var journalMode string
if err := db.QueryRow("PRAGMA journal_mode;").Scan(&journalMode); err == nil {
logger.Debug("SQLite journal mode", "mode", journalMode)
}
p := ProviderSQL{db: db, logger: logger}
if err := p.Migrate(); err != nil {
logger.Error("migration failed, app cannot start", "error", err)

View File

@@ -4,6 +4,7 @@ import (
"encoding/binary"
"fmt"
"gf-lt/models"
"sort"
"unsafe"
"github.com/jmoiron/sqlx"
@@ -11,7 +12,7 @@ import (
type VectorRepo interface {
WriteVector(*models.VectorRow) error
SearchClosest(q []float32) ([]models.VectorRow, error)
SearchClosest(q []float32, limit int) ([]models.VectorRow, error)
ListFiles() ([]string, error)
RemoveEmbByFileName(filename string) error
DB() *sqlx.DB
@@ -79,7 +80,7 @@ func (p ProviderSQL) WriteVector(row *models.VectorRow) error {
return err
}
func (p ProviderSQL) SearchClosest(q []float32) ([]models.VectorRow, error) {
func (p ProviderSQL) SearchClosest(q []float32, limit int) ([]models.VectorRow, error) {
tableName, err := fetchTableName(q)
if err != nil {
return nil, err
@@ -94,7 +95,7 @@ func (p ProviderSQL) SearchClosest(q []float32) ([]models.VectorRow, error) {
vector models.VectorRow
distance float32
}
var topResults []SearchResult
var allResults []SearchResult
for rows.Next() {
var (
embeddingsBlob []byte
@@ -119,28 +120,19 @@ func (p ProviderSQL) SearchClosest(q []float32) ([]models.VectorRow, error) {
},
distance: distance,
}
// Add to top results and maintain only top results
topResults = append(topResults, result)
// Sort and keep only top results
// We'll keep the top 3 closest vectors
if len(topResults) > 3 {
// Simple sort and truncate to maintain only 3 best matches
for i := 0; i < len(topResults); i++ {
for j := i + 1; j < len(topResults); j++ {
if topResults[i].distance > topResults[j].distance {
topResults[i], topResults[j] = topResults[j], topResults[i]
allResults = append(allResults, result)
}
// Sort by distance
sort.Slice(allResults, func(i, j int) bool {
return allResults[i].distance < allResults[j].distance
})
// Truncate to limit
if len(allResults) > limit {
allResults = allResults[:limit]
}
}
topResults = topResults[:3]
}
}
// Convert back to VectorRow slice
results := make([]models.VectorRow, len(topResults))
for i, result := range topResults {
results := make([]models.VectorRow, len(allResults))
for i, result := range allResults {
result.vector.Distance = result.distance
results[i] = result.vector
}

View File

@@ -1,7 +0,0 @@
{
"sys_prompt": "A game of cluedo. Players are {{user}}, {{char}}, {{char2}};\n\nrooms: hall, lounge, dinning room kitchen, ballroom, conservatory, billiard room, library, study;\nweapons: candlestick, dagger, lead pipe, revolver, rope, spanner;\npeople: miss Scarlett, colonel Mustard, mrs. White, reverend Green, mrs. Peacock, professor Plum;\n\nA murder happened in a mansion with 9 rooms. Victim is dr. Black.\nPlayers goal is to find out who commited a murder, in what room and with what weapon.\nWeapons, people and rooms not involved in murder are distributed between players (as cards) by tool agent.\nThe objective of the game is to deduce the details of the murder. There are six characters, six murder weapons, and nine rooms, leaving the players with 324 possibilities. As soon as a player enters a room, they may make a suggestion as to the details, naming a suspect, the room they are in, and the weapon. For example: \"I suspect Professor Plum, in the Dining Room, with the candlestick\".\nOnce a player makes a suggestion, the others are called upon to disprove it.\nBefore the player's move, tool agent will remind that players their cards. There are two types of moves: making a suggestion (suggestion_move) and disproving other player suggestion (evidence_move);\nIn this version player wins when the correct details are named in the suggestion_move.\n\n<example_game>\n{{user}}:\nlet's start a game of cluedo!\ntool: cards of {{char}} are 'LEAD PIPE', 'BALLROOM', 'CONSERVATORY', 'STUDY', 'Mrs. White'; suggestion_move;\n{{char}}:\n(putting miss Scarlet into the Hall with the Revolver) \"I suspect miss Scarlett, in the Hall, with the revolver.\"\ntool: cards of {{char2}} are 'SPANNER', 'DAGGER', 'Professor Plum', 'LIBRARY', 'Mrs. Peacock'; evidence_move;\n{{char2}}:\n\"No objections.\" (no cards matching the suspicion of {{char}})\ntool: cards of {{user}} are 'Colonel Mustard', 'Miss Scarlett', 'DINNING ROOM', 'CANDLESTICK', 'HALL'; evidence_move;\n{{user}}:\n\"I object. Miss Scarlett is innocent.\" (shows card with 'Miss Scarlett')\ntool: cards of {{char2}} are 'SPANNER', 'DAGGER', 'Professor Plum', 'LIBRARY', 'Mrs. Peacock'; suggestion_move;\n{{char2}}:\n*So it was not Miss Scarlett, good to know.*\n(moves Mrs. White to the Billiard Room) \"It might have been Mrs. White, in the Billiard Room, with the Revolver.\"\ntool: cards of {{user}} are 'Colonel Mustard', 'Miss Scarlett', 'DINNING ROOM', 'CANDLESTICK', 'HALL'; evidence_move;\n{{user}}:\n(no matching cards for the assumption of {{char2}}) \"Sounds possible to me.\"\ntool: cards of {{char}} are 'LEAD PIPE', 'BALLROOM', 'CONSERVATORY', 'STUDY', 'Mrs. White'; evidence_move;\n{{char}}:\n(shows Mrs. White card) \"No. Was not Mrs. White\"\ntool: cards of {{user}} are 'Colonel Mustard', 'Miss Scarlett', 'DINNING ROOM', 'CANDLESTICK', 'HALL'; suggestion_move;\n{{user}}:\n*So not Mrs. White...* (moves Reverend Green into the Billiard Room) \"I suspect Reverend Green, in the Billiard Room, with the Revolver.\"\ntool: Correct. It was Reverend Green in the Billiard Room, with the revolver. {{user}} wins.\n</example_game>",
"role": "CluedoPlayer",
"role2": "CluedoEnjoyer",
"filepath": "sysprompts/cluedo.json",
"first_msg": "Hey guys! Want to play cluedo?"
}

72
tui.go
View File

@@ -29,6 +29,8 @@ var (
statusLineWidget *tview.TextView
helpView *tview.TextView
flex *tview.Flex
bottomFlex *tview.Flex
notificationWidget *tview.TextView
imgView *tview.Image
defaultImage = "sysprompts/llama.png"
indexPickWindow *tview.InputField
@@ -36,6 +38,7 @@ var (
roleEditWindow *tview.InputField
shellInput *tview.InputField
confirmModal *tview.Modal
toastTimer *time.Timer
confirmPageName = "confirm"
fullscreenMode bool
positionVisible bool = true
@@ -137,8 +140,8 @@ func setShellMode(enabled bool) {
}()
}
// showToast displays a temporary message in the topright corner.
// It autohides after 3 seconds and disappears when clicked.
// showToast displays a temporary notification in the bottom-right corner.
// It auto-hides after 3 seconds.
func showToast(title, message string) {
sanitize := func(s string, maxLen int) string {
sanitized := strings.Map(func(r rune) rune {
@@ -154,6 +157,11 @@ func showToast(title, message string) {
}
title = sanitize(title, 50)
message = sanitize(message, 197)
if toastTimer != nil {
toastTimer.Stop()
}
// show blocking notification to not mess up flex
if fullscreenMode {
notification := tview.NewTextView().
SetTextAlign(tview.AlignCenter).
SetDynamicColors(true).
@@ -176,14 +184,44 @@ func showToast(title, message string) {
// Generate a unique page name (e.g., using timestamp) to allow multiple toasts.
pageName := fmt.Sprintf("toast-%d", time.Now().UnixNano())
pages.AddPage(pageName, background, true, true)
// Autodismiss after 3 seconds.
time.AfterFunc(3*time.Second, func() {
// Autodismiss after 2 seconds, since blocking is more annoying
time.AfterFunc(2*time.Second, func() {
app.QueueUpdateDraw(func() {
if pages.HasPage(pageName) {
pages.RemovePage(pageName)
}
})
})
return
}
notificationWidget.SetTitle(title)
notificationWidget.SetText(fmt.Sprintf("[yellow]%s[-]", message))
go func() {
app.QueueUpdateDraw(func() {
flex.RemoveItem(bottomFlex)
flex.RemoveItem(statusLineWidget)
bottomFlex = tview.NewFlex().SetDirection(tview.FlexColumn).
AddItem(textArea, 0, 1, true).
AddItem(notificationWidget, 40, 1, false)
flex.AddItem(bottomFlex, 0, 10, true)
if positionVisible {
flex.AddItem(statusLineWidget, 0, 2, false)
}
})
}()
toastTimer = time.AfterFunc(3*time.Second, func() {
app.QueueUpdateDraw(func() {
flex.RemoveItem(bottomFlex)
flex.RemoveItem(statusLineWidget)
bottomFlex = tview.NewFlex().SetDirection(tview.FlexColumn).
AddItem(textArea, 0, 1, true).
AddItem(notificationWidget, 0, 0, false)
flex.AddItem(bottomFlex, 0, 10, true)
if positionVisible {
flex.AddItem(statusLineWidget, 0, 2, false)
}
})
})
}
func init() {
@@ -235,7 +273,7 @@ func init() {
shellHistoryPos = -1
}
// Handle Tab key for @ file completion
if event.Key() == tcell.KeyTab {
if event.Key() == tcell.KeyTab && shellMode {
currentText := shellInput.GetText()
atIndex := strings.LastIndex(currentText, "@")
if atIndex >= 0 {
@@ -286,12 +324,26 @@ func init() {
SetDynamicColors(true).
SetRegions(true).
SetChangedFunc(func() {
// INFO:
// https://github.com/rivo/tview/wiki/Concurrency#event-handlers
// although already called by default per tview specs
// calling it explicitly makes text streaming to look more smooth
app.Draw()
})
notificationWidget = tview.NewTextView().
SetTextAlign(tview.AlignCenter).
SetDynamicColors(true).
SetRegions(true).
SetChangedFunc(func() {
})
notificationWidget.SetBorder(true).SetTitle("notification")
bottomFlex = tview.NewFlex().SetDirection(tview.FlexColumn).
AddItem(textArea, 0, 1, true).
AddItem(notificationWidget, 0, 0, false)
//
flex = tview.NewFlex().SetDirection(tview.FlexRow).
AddItem(textView, 0, 40, false).
AddItem(textArea, 0, 10, true) // Restore original height
AddItem(bottomFlex, 0, 10, true)
if positionVisible {
flex.AddItem(statusLineWidget, 0, 2, false)
}
@@ -360,10 +412,14 @@ func init() {
// y += h / 2
// return x, y, w, h
// })
notificationWidget.SetDrawFunc(func(screen tcell.Screen, x, y, w, h int) (int, int, int, int) {
y += h / 2
return x, y, w, h
})
// Initially set up flex without search bar
flex = tview.NewFlex().SetDirection(tview.FlexRow).
AddItem(textView, 0, 40, false).
AddItem(textArea, 0, 10, true) // Restore original height
AddItem(bottomFlex, 0, 10, true)
if positionVisible {
flex.AddItem(statusLineWidget, 0, 2, false)
}
@@ -1095,7 +1151,7 @@ func init() {
chatRoundChan <- &models.ChatRoundReq{Role: persona, UserMsg: msgText}
return nil
}
if event.Key() == tcell.KeyTab {
if event.Key() == tcell.KeyTab && !shellMode {
currentF := app.GetFocus()
if currentF == textArea {
currentText := textArea.GetText()