9 Commits

Author SHA1 Message Date
Grail Finder
b386c1181f Fix (rag): epub load 2026-02-25 14:54:10 +03:00
Grail Finder
b8e7649e69 Enha (rag): one table to manage files and data loaded 2026-02-25 10:47:35 +03:00
Grail Finder
6664c1a0fc Dep (rag): better extractors 2026-02-25 07:51:24 +03:00
Grail Finder
e0c3fe554f Feat: rag text extractors 2026-02-25 06:51:02 +03:00
Grail Finder
40943ff4d3 Enha: spinner for tool calls 2026-02-24 21:47:57 +03:00
Grail Finder
6c03a1a277 Feat: rag tool 2026-02-24 20:24:44 +03:00
Grail Finder
27288e2aaa Enha: spinner to indicate llm response 2026-02-24 18:05:05 +03:00
Grail Finder
1c728ec7a7 Enha: close rag on success 2026-02-24 17:42:58 +03:00
Grail Finder
78059083c2 Enha (rag): singlethred 2026-02-24 15:28:18 +03:00
10 changed files with 792 additions and 342 deletions

30
bot.go
View File

@@ -826,8 +826,36 @@ func chatWatcher(ctx context.Context) {
} }
} }
// inpired by https://github.com/rivo/tview/issues/225
func showSpinner() {
spinners := []string{"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"}
var i int
botPersona := cfg.AssistantRole
if cfg.WriteNextMsgAsCompletionAgent != "" {
botPersona = cfg.WriteNextMsgAsCompletionAgent
}
for botRespMode || toolRunningMode {
time.Sleep(100 * time.Millisecond)
spin := i % len(spinners)
app.QueueUpdateDraw(func() {
if toolRunningMode {
textArea.SetTitle(spinners[spin] + " tool")
} else if botRespMode {
textArea.SetTitle(spinners[spin] + " " + botPersona)
} else {
textArea.SetTitle(spinners[spin] + " input")
}
})
i++
}
app.QueueUpdateDraw(func() {
textArea.SetTitle("input")
})
}
func chatRound(r *models.ChatRoundReq) error { func chatRound(r *models.ChatRoundReq) error {
botRespMode = true botRespMode = true
go showSpinner()
updateStatusLine() updateStatusLine()
botPersona := cfg.AssistantRole botPersona := cfg.AssistantRole
if cfg.WriteNextMsgAsCompletionAgent != "" { if cfg.WriteNextMsgAsCompletionAgent != "" {
@@ -1201,7 +1229,9 @@ func findCall(msg, toolCall string) bool {
} }
// Show tool call progress indicator before execution // Show tool call progress indicator before execution
fmt.Fprintf(textView, "\n[yellow::i][tool: %s...][-:-:-]", fc.Name) fmt.Fprintf(textView, "\n[yellow::i][tool: %s...][-:-:-]", fc.Name)
toolRunningMode = true
resp := callToolWithAgent(fc.Name, fc.Args) resp := callToolWithAgent(fc.Name, fc.Args)
toolRunningMode = false
toolMsg := string(resp) // Remove the "tool response: " prefix and %+v formatting toolMsg := string(resp) // Remove the "tool response: " prefix and %+v formatting
logger.Info("llm used a tool call", "tool_name", fc.Name, "too_args", fc.Args, "id", fc.ID, "tool_resp", toolMsg) logger.Info("llm used a tool call", "tool_name", fc.Name, "too_args", fc.Args, "id", fc.ID, "tool_resp", toolMsg)
fmt.Fprintf(textView, "%s[-:-:b](%d) <%s>: [-:-:-]\n%s\n", fmt.Fprintf(textView, "%s[-:-:b](%d) <%s>: [-:-:-]\n%s\n",

4
go.mod
View File

@@ -6,17 +6,19 @@ require (
github.com/BurntSushi/toml v1.5.0 github.com/BurntSushi/toml v1.5.0
github.com/GrailFinder/google-translate-tts v0.1.3 github.com/GrailFinder/google-translate-tts v0.1.3
github.com/GrailFinder/searchagent v0.2.0 github.com/GrailFinder/searchagent v0.2.0
github.com/PuerkitoBio/goquery v1.11.0
github.com/gdamore/tcell/v2 v2.13.2 github.com/gdamore/tcell/v2 v2.13.2
github.com/glebarez/go-sqlite v1.22.0 github.com/glebarez/go-sqlite v1.22.0
github.com/gopxl/beep/v2 v2.1.1 github.com/gopxl/beep/v2 v2.1.1
github.com/gordonklaus/portaudio v0.0.0-20250206071425-98a94950218b github.com/gordonklaus/portaudio v0.0.0-20250206071425-98a94950218b
github.com/jmoiron/sqlx v1.4.0 github.com/jmoiron/sqlx v1.4.0
github.com/ledongthuc/pdf v0.0.0-20250511090121-5959a4027728
github.com/neurosnap/sentences v1.1.2 github.com/neurosnap/sentences v1.1.2
github.com/rivo/tview v0.42.0 github.com/rivo/tview v0.42.0
github.com/yuin/goldmark v1.4.13
) )
require ( require (
github.com/PuerkitoBio/goquery v1.11.0 // indirect
github.com/andybalholm/cascadia v1.3.3 // indirect github.com/andybalholm/cascadia v1.3.3 // indirect
github.com/dustin/go-humanize v1.0.1 // indirect github.com/dustin/go-humanize v1.0.1 // indirect
github.com/ebitengine/oto/v3 v3.4.0 // indirect github.com/ebitengine/oto/v3 v3.4.0 // indirect

3
go.sum
View File

@@ -43,6 +43,8 @@ github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
github.com/jmoiron/sqlx v1.4.0 h1:1PLqN7S1UYp5t4SrVVnt4nUVNemrDAtxlulVe+Qgm3o= github.com/jmoiron/sqlx v1.4.0 h1:1PLqN7S1UYp5t4SrVVnt4nUVNemrDAtxlulVe+Qgm3o=
github.com/jmoiron/sqlx v1.4.0/go.mod h1:ZrZ7UsYB/weZdl2Bxg6jCRO9c3YHl8r3ahlKmRT4JLY= github.com/jmoiron/sqlx v1.4.0/go.mod h1:ZrZ7UsYB/weZdl2Bxg6jCRO9c3YHl8r3ahlKmRT4JLY=
github.com/ledongthuc/pdf v0.0.0-20250511090121-5959a4027728 h1:QwWKgMY28TAXaDl+ExRDqGQltzXqN/xypdKP86niVn8=
github.com/ledongthuc/pdf v0.0.0-20250511090121-5959a4027728/go.mod h1:1fEHWurg7pvf5SG6XNE5Q8UZmOwex51Mkx3SLhrW5B4=
github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw= github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw=
github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o= github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag= github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag=
@@ -67,6 +69,7 @@ github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= 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/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
github.com/yuin/goldmark v1.4.13 h1:fVcFKWvrslecOb/tg+Cc05dkeYx540o0FuFt3nUVDoE=
github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= 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= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=

98
llm.go
View File

@@ -11,7 +11,6 @@ import (
var imageAttachmentPath string // Global variable to track image attachment for next message var imageAttachmentPath string // Global variable to track image attachment for next message
var lastImg string // for ctrl+j var lastImg string // for ctrl+j
var RAGMsg = "Retrieved context for user's query:\n"
// containsToolSysMsg checks if the toolSysMsg already exists in the chat body // containsToolSysMsg checks if the toolSysMsg already exists in the chat body
func containsToolSysMsg() bool { func containsToolSysMsg() bool {
@@ -142,22 +141,6 @@ func (lcp LCPCompletion) FormMsg(msg, role string, resume bool) (io.Reader, erro
newMsg = *processMessageTag(&newMsg) newMsg = *processMessageTag(&newMsg)
chatBody.Messages = append(chatBody.Messages, newMsg) chatBody.Messages = append(chatBody.Messages, newMsg)
} }
// if rag - add as system message to avoid conflicts with tool usage
if !resume && cfg.RAGEnabled {
um := chatBody.Messages[len(chatBody.Messages)-1].Content
logger.Debug("RAG is enabled, preparing RAG context", "user_message", um)
ragResp, err := chatRagUse(um)
if err != nil {
logger.Error("failed to form a rag msg", "error", err)
return nil, err
}
logger.Debug("RAG response received", "response_len", len(ragResp),
"response_preview", ragResp[:min(len(ragResp), 100)])
// Use system role for RAG context to avoid conflicts with tool usage
ragMsg := models.RoleMsg{Role: "system", Content: RAGMsg + ragResp}
chatBody.Messages = append(chatBody.Messages, ragMsg)
logger.Debug("RAG message added to chat body", "message_count", len(chatBody.Messages))
}
// sending description of the tools and how to use them // sending description of the tools and how to use them
if cfg.ToolUse && !resume && role == cfg.UserRole && !containsToolSysMsg() { if cfg.ToolUse && !resume && role == cfg.UserRole && !containsToolSysMsg() {
chatBody.Messages = append(chatBody.Messages, models.RoleMsg{Role: cfg.ToolRole, Content: toolSysMsg}) chatBody.Messages = append(chatBody.Messages, models.RoleMsg{Role: cfg.ToolRole, Content: toolSysMsg})
@@ -301,23 +284,6 @@ func (op LCPChat) FormMsg(msg, role string, resume bool) (io.Reader, error) {
logger.Debug("LCPChat FormMsg: added message to chatBody", "role", newMsg.Role, logger.Debug("LCPChat FormMsg: added message to chatBody", "role", newMsg.Role,
"content_len", len(newMsg.Content), "message_count_after_add", len(chatBody.Messages)) "content_len", len(newMsg.Content), "message_count_after_add", len(chatBody.Messages))
} }
// if rag - add as system message to avoid conflicts with tool usage
if !resume && cfg.RAGEnabled {
um := chatBody.Messages[len(chatBody.Messages)-1].Content
logger.Debug("LCPChat: RAG is enabled, preparing RAG context", "user_message", um)
ragResp, err := chatRagUse(um)
if err != nil {
logger.Error("LCPChat: failed to form a rag msg", "error", err)
return nil, err
}
logger.Debug("LCPChat: RAG response received",
"response_len", len(ragResp), "response_preview", ragResp[:min(len(ragResp), 100)])
// Use system role for RAG context to avoid conflicts with tool usage
ragMsg := models.RoleMsg{Role: "system", Content: RAGMsg + ragResp}
chatBody.Messages = append(chatBody.Messages, ragMsg)
logger.Debug("LCPChat: RAG message added to chat body", "role", ragMsg.Role,
"rag_content_len", len(ragMsg.Content), "message_count_after_rag", len(chatBody.Messages))
}
filteredMessages, _ := filterMessagesForCurrentCharacter(chatBody.Messages) filteredMessages, _ := filterMessagesForCurrentCharacter(chatBody.Messages)
// openai /v1/chat does not support custom roles; needs to be user, assistant, system // openai /v1/chat does not support custom roles; needs to be user, assistant, system
// Add persona suffix to the last user message to indicate who the assistant should reply as // Add persona suffix to the last user message to indicate who the assistant should reply as
@@ -389,22 +355,6 @@ func (ds DeepSeekerCompletion) FormMsg(msg, role string, resume bool) (io.Reader
newMsg = *processMessageTag(&newMsg) newMsg = *processMessageTag(&newMsg)
chatBody.Messages = append(chatBody.Messages, newMsg) chatBody.Messages = append(chatBody.Messages, newMsg)
} }
// if rag - add as system message to avoid conflicts with tool usage
if !resume && cfg.RAGEnabled {
um := chatBody.Messages[len(chatBody.Messages)-1].Content
logger.Debug("DeepSeekerCompletion: RAG is enabled, preparing RAG context", "user_message", um)
ragResp, err := chatRagUse(um)
if err != nil {
logger.Error("DeepSeekerCompletion: failed to form a rag msg", "error", err)
return nil, err
}
logger.Debug("DeepSeekerCompletion: RAG response received",
"response_len", len(ragResp), "response_preview", ragResp[:min(len(ragResp), 100)])
// Use system role for RAG context to avoid conflicts with tool usage
ragMsg := models.RoleMsg{Role: "system", Content: RAGMsg + ragResp}
chatBody.Messages = append(chatBody.Messages, ragMsg)
logger.Debug("DeepSeekerCompletion: RAG message added to chat body", "message_count", len(chatBody.Messages))
}
// sending description of the tools and how to use them // sending description of the tools and how to use them
if cfg.ToolUse && !resume && role == cfg.UserRole && !containsToolSysMsg() { if cfg.ToolUse && !resume && role == cfg.UserRole && !containsToolSysMsg() {
chatBody.Messages = append(chatBody.Messages, models.RoleMsg{Role: cfg.ToolRole, Content: toolSysMsg}) chatBody.Messages = append(chatBody.Messages, models.RoleMsg{Role: cfg.ToolRole, Content: toolSysMsg})
@@ -474,22 +424,6 @@ func (ds DeepSeekerChat) FormMsg(msg, role string, resume bool) (io.Reader, erro
newMsg = *processMessageTag(&newMsg) newMsg = *processMessageTag(&newMsg)
chatBody.Messages = append(chatBody.Messages, newMsg) chatBody.Messages = append(chatBody.Messages, newMsg)
} }
// if rag - add as system message to avoid conflicts with tool usage
if !resume && cfg.RAGEnabled {
um := chatBody.Messages[len(chatBody.Messages)-1].Content
logger.Debug("RAG is enabled, preparing RAG context", "user_message", um)
ragResp, err := chatRagUse(um)
if err != nil {
logger.Error("failed to form a rag msg", "error", err)
return nil, err
}
logger.Debug("RAG response received", "response_len", len(ragResp),
"response_preview", ragResp[:min(len(ragResp), 100)])
// Use system role for RAG context to avoid conflicts with tool usage
ragMsg := models.RoleMsg{Role: "system", Content: RAGMsg + ragResp}
chatBody.Messages = append(chatBody.Messages, ragMsg)
logger.Debug("RAG message added to chat body", "message_count", len(chatBody.Messages))
}
// Create copy of chat body with standardized user role // Create copy of chat body with standardized user role
filteredMessages, _ := filterMessagesForCurrentCharacter(chatBody.Messages) filteredMessages, _ := filterMessagesForCurrentCharacter(chatBody.Messages)
// Add persona suffix to the last user message to indicate who the assistant should reply as // Add persona suffix to the last user message to indicate who the assistant should reply as
@@ -552,22 +486,6 @@ func (or OpenRouterCompletion) FormMsg(msg, role string, resume bool) (io.Reader
newMsg = *processMessageTag(&newMsg) newMsg = *processMessageTag(&newMsg)
chatBody.Messages = append(chatBody.Messages, newMsg) chatBody.Messages = append(chatBody.Messages, newMsg)
} }
// if rag - add as system message to avoid conflicts with tool usage
if !resume && cfg.RAGEnabled {
um := chatBody.Messages[len(chatBody.Messages)-1].Content
logger.Debug("RAG is enabled, preparing RAG context", "user_message", um)
ragResp, err := chatRagUse(um)
if err != nil {
logger.Error("failed to form a rag msg", "error", err)
return nil, err
}
logger.Debug("RAG response received", "response_len",
len(ragResp), "response_preview", ragResp[:min(len(ragResp), 100)])
// Use system role for RAG context to avoid conflicts with tool usage
ragMsg := models.RoleMsg{Role: "system", Content: RAGMsg + ragResp}
chatBody.Messages = append(chatBody.Messages, ragMsg)
logger.Debug("RAG message added to chat body", "message_count", len(chatBody.Messages))
}
// sending description of the tools and how to use them // sending description of the tools and how to use them
if cfg.ToolUse && !resume && role == cfg.UserRole && !containsToolSysMsg() { if cfg.ToolUse && !resume && role == cfg.UserRole && !containsToolSysMsg() {
chatBody.Messages = append(chatBody.Messages, models.RoleMsg{Role: cfg.ToolRole, Content: toolSysMsg}) chatBody.Messages = append(chatBody.Messages, models.RoleMsg{Role: cfg.ToolRole, Content: toolSysMsg})
@@ -670,22 +588,6 @@ func (or OpenRouterChat) FormMsg(msg, role string, resume bool) (io.Reader, erro
newMsg = *processMessageTag(&newMsg) newMsg = *processMessageTag(&newMsg)
chatBody.Messages = append(chatBody.Messages, newMsg) chatBody.Messages = append(chatBody.Messages, newMsg)
} }
// if rag - add as system message to avoid conflicts with tool usage
if !resume && cfg.RAGEnabled {
um := chatBody.Messages[len(chatBody.Messages)-1].Content
logger.Debug("RAG is enabled, preparing RAG context", "user_message", um)
ragResp, err := chatRagUse(um)
if err != nil {
logger.Error("failed to form a rag msg", "error", err)
return nil, err
}
logger.Debug("RAG response received", "response_len", len(ragResp),
"response_preview", ragResp[:min(len(ragResp), 100)])
// Use system role for RAG context to avoid conflicts with tool usage
ragMsg := models.RoleMsg{Role: "system", Content: RAGMsg + ragResp}
chatBody.Messages = append(chatBody.Messages, ragMsg)
logger.Debug("RAG message added to chat body", "message_count", len(chatBody.Messages))
}
// Create copy of chat body with standardized user role // Create copy of chat body with standardized user role
filteredMessages, _ := filterMessagesForCurrentCharacter(chatBody.Messages) filteredMessages, _ := filterMessagesForCurrentCharacter(chatBody.Messages)
// Add persona suffix to the last user message to indicate who the assistant should reply as // Add persona suffix to the last user message to indicate who the assistant should reply as

View File

@@ -7,6 +7,7 @@ import (
var ( var (
boolColors = map[bool]string{true: "green", false: "red"} boolColors = map[bool]string{true: "green", false: "red"}
botRespMode = false botRespMode = false
toolRunningMode = false
editMode = false editMode = false
roleEditMode = false roleEditMode = false
injectRole = true injectRole = true

184
rag/extractors.go Normal file
View File

@@ -0,0 +1,184 @@
package rag
import (
"archive/zip"
"bytes"
"errors"
"fmt"
"io"
"os"
"os/exec"
"path"
"strings"
"github.com/PuerkitoBio/goquery"
"github.com/ledongthuc/pdf"
"github.com/yuin/goldmark"
"github.com/yuin/goldmark/extension"
"github.com/yuin/goldmark/parser"
"github.com/yuin/goldmark/renderer/html"
)
func ExtractText(fpath string) (string, error) {
ext := strings.ToLower(path.Ext(fpath))
switch ext {
case ".txt":
return extractTextFromFile(fpath)
case ".md", ".markdown":
return extractTextFromMarkdown(fpath)
case ".html", ".htm":
return extractTextFromHtmlFile(fpath)
case ".epub":
return extractTextFromEpub(fpath)
case ".pdf":
return extractTextFromPdf(fpath)
default:
return "", fmt.Errorf("unsupported file format: %s", ext)
}
}
func extractTextFromFile(fpath string) (string, error) {
data, err := os.ReadFile(fpath)
if err != nil {
return "", err
}
return string(data), nil
}
func extractTextFromHtmlFile(fpath string) (string, error) {
data, err := os.ReadFile(fpath)
if err != nil {
return "", err
}
return extractTextFromHtmlContent(data)
}
// non utf-8 encoding?
func extractTextFromHtmlContent(data []byte) (string, error) {
doc, err := goquery.NewDocumentFromReader(bytes.NewReader(data))
if err != nil {
return "", err
}
// Remove script and style tags
doc.Find("script, style, noscript").Each(func(i int, s *goquery.Selection) {
s.Remove()
})
// Get text and clean it
text := doc.Text()
// Collapse all whitespace (newlines, tabs, multiple spaces) into single spaces
cleaned := strings.Join(strings.Fields(text), " ")
return cleaned, nil
}
func extractTextFromMarkdown(fpath string) (string, error) {
data, err := os.ReadFile(fpath)
if err != nil {
return "", err
}
// Convert markdown to HTML
md := goldmark.New(
goldmark.WithExtensions(extension.GFM),
goldmark.WithParserOptions(parser.WithAutoHeadingID()),
goldmark.WithRendererOptions(html.WithUnsafe()), // allow raw HTML if needed
)
var buf bytes.Buffer
if err := md.Convert(data, &buf); err != nil {
return "", err
}
// Now extract text from the resulting HTML (using goquery or similar)
return extractTextFromHtmlContent(buf.Bytes())
}
func extractTextFromEpub(fpath string) (string, error) {
r, err := zip.OpenReader(fpath)
if err != nil {
return "", fmt.Errorf("failed to open epub: %w", err)
}
defer r.Close()
var sb strings.Builder
for _, f := range r.File {
ext := strings.ToLower(path.Ext(f.Name))
if ext != ".xhtml" && ext != ".html" && ext != ".htm" && ext != ".xml" {
continue
}
// Skip manifest, toc, ncx files - they don't contain book content
nameLower := strings.ToLower(f.Name)
if strings.Contains(nameLower, "toc") || strings.Contains(nameLower, "nav") ||
strings.Contains(nameLower, "manifest") || strings.Contains(nameLower, ".opf") ||
strings.HasSuffix(nameLower, ".ncx") {
continue
}
rc, err := f.Open()
if err != nil {
continue
}
if sb.Len() > 0 {
sb.WriteString("\n\n")
}
sb.WriteString(f.Name)
sb.WriteString("\n")
buf, readErr := io.ReadAll(rc)
rc.Close()
if readErr == nil {
sb.WriteString(stripHTML(string(buf)))
}
}
if sb.Len() == 0 {
return "", errors.New("no content extracted from epub")
}
return sb.String(), nil
}
func stripHTML(html string) string {
var sb strings.Builder
inTag := false
for _, r := range html {
switch r {
case '<':
inTag = true
case '>':
inTag = false
default:
if !inTag {
sb.WriteRune(r)
}
}
}
return sb.String()
}
func extractTextFromPdf(fpath string) (string, error) {
_, err := exec.LookPath("pdftotext")
if err == nil {
out, err := exec.Command("pdftotext", "-layout", fpath, "-").Output()
if err == nil && len(out) > 0 {
return string(out), nil
}
}
return extractTextFromPdfPureGo(fpath)
}
func extractTextFromPdfPureGo(fpath string) (string, error) {
df, r, err := pdf.Open(fpath)
if err != nil {
return "", fmt.Errorf("failed to open pdf: %w", err)
}
defer df.Close()
textReader, err := r.GetPlainText()
if err != nil {
return "", fmt.Errorf("failed to extract text from pdf: %w", err)
}
var buf bytes.Buffer
_, err = io.Copy(&buf, textReader)
if err != nil {
return "", fmt.Errorf("failed to read pdf text: %w", err)
}
return buf.String(), nil
}

View File

@@ -7,8 +7,9 @@ import (
"gf-lt/models" "gf-lt/models"
"gf-lt/storage" "gf-lt/storage"
"log/slog" "log/slog"
"os"
"path" "path"
"regexp"
"sort"
"strings" "strings"
"sync" "sync"
@@ -29,6 +30,7 @@ type RAG struct {
cfg *config.Config cfg *config.Config
embedder Embedder embedder Embedder
storage *VectorStorage storage *VectorStorage
mu sync.Mutex
} }
func New(l *slog.Logger, s storage.FullRepo, cfg *config.Config) *RAG { func New(l *slog.Logger, s storage.FullRepo, cfg *config.Config) *RAG {
@@ -53,7 +55,9 @@ func wordCounter(sentence string) int {
} }
func (r *RAG) LoadRAG(fpath string) error { func (r *RAG) LoadRAG(fpath string) error {
data, err := os.ReadFile(fpath) r.mu.Lock()
defer r.mu.Unlock()
fileText, err := ExtractText(fpath)
if err != nil { if err != nil {
return err return err
} }
@@ -62,10 +66,7 @@ func (r *RAG) LoadRAG(fpath string) error {
case LongJobStatusCh <- LoadedFileRAGStatus: case LongJobStatusCh <- LoadedFileRAGStatus:
default: default:
r.logger.Warn("LongJobStatusCh channel is full or closed, dropping status message", "message", LoadedFileRAGStatus) r.logger.Warn("LongJobStatusCh channel is full or closed, dropping status message", "message", LoadedFileRAGStatus)
// Channel is full or closed, ignore the message to prevent panic
} }
fileText := string(data)
tokenizer, err := english.NewSentenceTokenizer(nil) tokenizer, err := english.NewSentenceTokenizer(nil)
if err != nil { if err != nil {
return err return err
@@ -75,19 +76,16 @@ func (r *RAG) LoadRAG(fpath string) error {
for i, s := range sentences { for i, s := range sentences {
sents[i] = s.Text sents[i] = s.Text
} }
// Group sentences into paragraphs based on word limit // Group sentences into paragraphs based on word limit
paragraphs := []string{} paragraphs := []string{}
par := strings.Builder{} par := strings.Builder{}
for i := 0; i < len(sents); i++ { for i := 0; i < len(sents); i++ {
// Only add sentences that aren't empty
if strings.TrimSpace(sents[i]) != "" { if strings.TrimSpace(sents[i]) != "" {
if par.Len() > 0 { if par.Len() > 0 {
par.WriteString(" ") // Add space between sentences par.WriteString(" ")
} }
par.WriteString(sents[i]) par.WriteString(sents[i])
} }
if wordCounter(par.String()) > int(r.cfg.RAGWordLimit) { if wordCounter(par.String()) > int(r.cfg.RAGWordLimit) {
paragraph := strings.TrimSpace(par.String()) paragraph := strings.TrimSpace(par.String())
if paragraph != "" { if paragraph != "" {
@@ -96,7 +94,6 @@ func (r *RAG) LoadRAG(fpath string) error {
par.Reset() par.Reset()
} }
} }
// Handle any remaining content in the paragraph buffer // Handle any remaining content in the paragraph buffer
if par.Len() > 0 { if par.Len() > 0 {
paragraph := strings.TrimSpace(par.String()) paragraph := strings.TrimSpace(par.String())
@@ -104,217 +101,82 @@ func (r *RAG) LoadRAG(fpath string) error {
paragraphs = append(paragraphs, paragraph) paragraphs = append(paragraphs, paragraph)
} }
} }
// Adjust batch size if needed // Adjust batch size if needed
if len(paragraphs) < r.cfg.RAGBatchSize && len(paragraphs) > 0 { if len(paragraphs) < r.cfg.RAGBatchSize && len(paragraphs) > 0 {
r.cfg.RAGBatchSize = len(paragraphs) r.cfg.RAGBatchSize = len(paragraphs)
} }
if len(paragraphs) == 0 { if len(paragraphs) == 0 {
return errors.New("no valid paragraphs found in file") return errors.New("no valid paragraphs found in file")
} }
// Process paragraphs in batches synchronously
var ( batchCount := 0
maxChSize = 100 for i := 0; i < len(paragraphs); i += r.cfg.RAGBatchSize {
left = 0 end := i + r.cfg.RAGBatchSize
right = r.cfg.RAGBatchSize if end > len(paragraphs) {
batchCh = make(chan map[int][]string, maxChSize) end = len(paragraphs)
vectorCh = make(chan []models.VectorRow, maxChSize)
errCh = make(chan error, 1)
doneCh = make(chan struct{})
wg = new(sync.WaitGroup)
)
defer close(doneCh)
defer close(errCh)
defer close(batchCh)
// Fill input channel with batches
ctn := 0
totalParagraphs := len(paragraphs)
for {
if right > totalParagraphs {
batchCh <- map[int][]string{left: paragraphs[left:]}
break
} }
batchCh <- map[int][]string{left: paragraphs[left:right]} batch := paragraphs[i:end]
left, right = right, right+r.cfg.RAGBatchSize batchCount++
ctn++ // Filter empty paragraphs
nonEmptyBatch := make([]string, 0, len(batch))
for _, p := range batch {
if strings.TrimSpace(p) != "" {
nonEmptyBatch = append(nonEmptyBatch, strings.TrimSpace(p))
} }
finishedBatchesMsg := fmt.Sprintf("finished batching batches#: %d; paragraphs: %d; sentences: %d\n", ctn+1, len(paragraphs), len(sents))
r.logger.Debug(finishedBatchesMsg)
select {
case LongJobStatusCh <- finishedBatchesMsg:
default:
r.logger.Warn("LongJobStatusCh channel is full or closed, dropping status message", "message", finishedBatchesMsg)
// Channel is full or closed, ignore the message to prevent panic
} }
if len(nonEmptyBatch) == 0 {
// Start worker goroutines with WaitGroup continue
wg.Add(int(r.cfg.RAGWorkers))
for w := 0; w < int(r.cfg.RAGWorkers); w++ {
go func(workerID int) {
defer wg.Done()
r.batchToVectorAsync(workerID, batchCh, vectorCh, errCh, doneCh, path.Base(fpath))
}(w)
} }
// Embed the batch
// Close batchCh to signal workers no more data is coming embeddings, err := r.embedder.EmbedSlice(nonEmptyBatch)
close(batchCh)
// Wait for all workers to finish, then close vectorCh
go func() {
wg.Wait()
close(vectorCh)
}()
// Check for errors from workers - this will block until an error occurs or all workers finish
select {
case err := <-errCh:
if err != nil { if err != nil {
r.logger.Error("error during RAG processing", "error", err) 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))
return err return err
} }
default: // Write vectors to storage
// No immediate error, continue filename := path.Base(fpath)
for j, text := range nonEmptyBatch {
vector := models.VectorRow{
Embeddings: embeddings[j],
RawText: text,
Slug: fmt.Sprintf("%s_%d_%d", filename, batchCount, j),
FileName: filename,
} }
// Write vectors to storage - this will block until vectorCh is closed
return r.writeVectors(vectorCh, errCh)
}
func (r *RAG) writeVectors(vectorCh chan []models.VectorRow, errCh chan error) error {
// Use a select to handle both vectorCh and errCh
for {
select {
case err := <-errCh:
if err != nil {
r.logger.Error("error during RAG processing in writeVectors", "error", err)
return err
}
case batch, ok := <-vectorCh:
if !ok {
r.logger.Debug("vector channel closed, finished writing vectors")
select {
case LongJobStatusCh <- FinishedRAGStatus:
default:
r.logger.Warn("LongJobStatusCh channel is full or closed, dropping status message", "message", FinishedRAGStatus)
}
return nil
}
for _, vector := range batch {
if err := r.storage.WriteVector(&vector); err != nil { if err := r.storage.WriteVector(&vector); err != nil {
r.logger.Error("failed to write vector to DB", "error", err, "slug", vector.Slug) r.logger.Error("failed to write vector to DB", "error", err, "slug", vector.Slug)
select { select {
case LongJobStatusCh <- ErrRAGStatus: case LongJobStatusCh <- ErrRAGStatus:
default: default:
r.logger.Warn("LongJobStatusCh channel is full or closed, dropping status message", "message", ErrRAGStatus) r.logger.Warn("LongJobStatusCh channel full, dropping message")
} }
return err return fmt.Errorf("failed to write vector: %w", err)
} }
} }
r.logger.Debug("wrote batch to db", "size", len(batch)) 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)
}
func (r *RAG) batchToVectorAsync(id int, inputCh <-chan map[int][]string,
vectorCh chan<- []models.VectorRow, errCh chan error, doneCh <-chan struct{}, filename string) {
var err error
defer func() {
if err != nil {
select {
case errCh <- err:
default:
r.logger.Warn("errCh channel full or closed, skipping error propagation", "worker", id, "error", err)
}
}
}()
for {
select {
case <-doneCh:
r.logger.Debug("worker received done signal", "worker", id)
return
case linesMap, ok := <-inputCh:
if !ok {
r.logger.Debug("input channel closed, worker exiting", "worker", id)
return
}
for leftI, lines := range linesMap {
select {
case <-doneCh:
return
default:
}
if err := r.fetchEmb(lines, errCh, vectorCh, fmt.Sprintf("%s_%d", filename, leftI), filename); err != nil {
r.logger.Error("error fetching embeddings", "error", err, "worker", id)
return
}
}
r.logger.Debug("processed batch", "worker#", id)
statusMsg := fmt.Sprintf("converted to vector; worker#: %d", id)
select { select {
case LongJobStatusCh <- statusMsg: case LongJobStatusCh <- statusMsg:
default: default:
r.logger.Warn("LongJobStatusCh channel full or closed, dropping status message", "message", statusMsg) 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)
} }
}
func (r *RAG) fetchEmb(lines []string, errCh chan error, vectorCh chan<- []models.VectorRow, slug, filename string) error {
// Filter out empty lines before sending to embedder
nonEmptyLines := make([]string, 0, len(lines))
for _, line := range lines {
trimmed := strings.TrimSpace(line)
if trimmed != "" {
nonEmptyLines = append(nonEmptyLines, trimmed)
}
}
// Skip if no non-empty lines
if len(nonEmptyLines) == 0 {
// Send empty result but don't error
vectorCh <- []models.VectorRow{}
return nil
}
embeddings, err := r.embedder.EmbedSlice(nonEmptyLines)
if err != nil {
r.logger.Error("failed to embed lines", "err", err.Error())
errCh <- err
return err
}
if len(embeddings) == 0 {
err := errors.New("no embeddings returned")
r.logger.Error("empty embeddings")
errCh <- err
return err
}
if len(embeddings) != len(nonEmptyLines) {
err := errors.New("mismatch between number of lines and embeddings returned")
r.logger.Error("embedding mismatch", "err", err.Error())
errCh <- err
return err
}
// Create a VectorRow for each line in the batch
vectors := make([]models.VectorRow, len(nonEmptyLines))
for i, line := range nonEmptyLines {
vectors[i] = models.VectorRow{
Embeddings: embeddings[i],
RawText: line,
Slug: fmt.Sprintf("%s_%d", slug, i),
FileName: filename,
}
}
vectorCh <- vectors
return nil return nil
} }
@@ -333,3 +195,309 @@ func (r *RAG) ListLoaded() ([]string, error) {
func (r *RAG) RemoveFile(filename string) error { func (r *RAG) RemoveFile(filename string) error {
return r.storage.RemoveEmbByFileName(filename) return r.storage.RemoveEmbByFileName(filename)
} }
var (
queryRefinementPattern = regexp.MustCompile(`(?i)(based on my (vector db|vector db|vector database|rags?|past (conversations?|chat|messages?))|from my (files?|documents?|data|information|memory)|search (in|my) (vector db|database|rags?)|rag search for)`)
importantKeywords = []string{"project", "architecture", "code", "file", "chat", "conversation", "topic", "summary", "details", "history", "previous", "my", "user", "me"}
stopWords = []string{"the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with", "by", "from", "up", "down", "left", "right"}
)
func (r *RAG) RefineQuery(query string) string {
original := query
query = strings.TrimSpace(query)
if len(query) == 0 {
return original
}
if len(query) <= 3 {
return original
}
query = strings.ToLower(query)
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
}
if queryRefinementPattern.MatchString(original) {
cleaned := queryRefinementPattern.ReplaceAllString(original, "")
cleaned = strings.TrimSpace(cleaned)
if len(cleaned) >= 5 {
return cleaned
}
}
query = r.extractImportantPhrases(query)
if len(query) < 5 {
return original
}
return query
}
func (r *RAG) extractImportantPhrases(query string) string {
words := strings.Fields(query)
var important []string
for _, word := range words {
word = strings.Trim(word, ".,!?;:'\"()[]{}")
isImportant := false
for _, kw := range importantKeywords {
if strings.Contains(strings.ToLower(word), kw) {
isImportant = true
break
}
}
if isImportant || len(word) > 3 {
important = append(important, word)
}
}
if len(important) == 0 {
return query
}
return strings.Join(important, " ")
}
func (r *RAG) GenerateQueryVariations(query string) []string {
variations := []string{query}
if len(query) < 5 {
return variations
}
parts := strings.Fields(query)
if len(parts) == 0 {
return variations
}
if len(parts) >= 2 {
trimmed := strings.Join(parts[:len(parts)-1], " ")
if len(trimmed) >= 5 {
variations = append(variations, trimmed)
}
}
if len(parts) >= 2 {
trimmed := strings.Join(parts[1:], " ")
if len(trimmed) >= 5 {
variations = append(variations, trimmed)
}
}
if !strings.HasSuffix(query, " explanation") {
variations = append(variations, query+" explanation")
}
if !strings.HasPrefix(query, "what is ") {
variations = append(variations, "what is "+query)
}
if !strings.HasSuffix(query, " details") {
variations = append(variations, query+" details")
}
if !strings.HasSuffix(query, " summary") {
variations = append(variations, query+" summary")
}
return variations
}
func (r *RAG) RerankResults(results []models.VectorRow, query string) []models.VectorRow {
type scoredResult struct {
row models.VectorRow
distance float32
}
scored := make([]scoredResult, 0, len(results))
for i := range results {
row := results[i]
score := float32(0)
rawTextLower := strings.ToLower(row.RawText)
queryLower := strings.ToLower(query)
if strings.Contains(rawTextLower, queryLower) {
score += 10
}
queryWords := strings.Fields(queryLower)
matchCount := 0
for _, word := range queryWords {
if len(word) > 2 && strings.Contains(rawTextLower, word) {
matchCount++
}
}
if len(queryWords) > 0 {
score += float32(matchCount) / float32(len(queryWords)) * 5
}
if row.FileName == "chat" || strings.Contains(strings.ToLower(row.FileName), "conversation") {
score += 3
}
distance := row.Distance - score/100
scored = append(scored, scoredResult{row: row, distance: distance})
}
sort.Slice(scored, func(i, j int) bool {
return scored[i].distance < scored[j].distance
})
unique := make([]models.VectorRow, 0)
seen := make(map[string]bool)
for i := range scored {
if !seen[scored[i].row.Slug] {
seen[scored[i].row.Slug] = true
unique = append(unique, scored[i].row)
}
}
if len(unique) > 10 {
unique = unique[:10]
}
return unique
}
func (r *RAG) SynthesizeAnswer(results []models.VectorRow, query string) (string, error) {
if len(results) == 0 {
return "No relevant information found in the vector database.", nil
}
var contextBuilder strings.Builder
contextBuilder.WriteString("User Query: ")
contextBuilder.WriteString(query)
contextBuilder.WriteString("\n\nRetrieved Context:\n")
for i, row := range results {
contextBuilder.WriteString(fmt.Sprintf("[Source %d: %s]\n", i+1, row.FileName))
contextBuilder.WriteString(row.RawText)
contextBuilder.WriteString("\n\n")
}
contextBuilder.WriteString("Instructions: ")
contextBuilder.WriteString("Based on the retrieved context above, provide a concise, coherent answer to the user's query. ")
contextBuilder.WriteString("Extract only the most relevant information. ")
contextBuilder.WriteString("If no relevant information is found, state that clearly. ")
contextBuilder.WriteString("Cite sources by filename when relevant. ")
contextBuilder.WriteString("Do not include unnecessary preamble or explanations.")
synthesisPrompt := contextBuilder.String()
emb, err := r.LineToVector(synthesisPrompt)
if err != nil {
r.logger.Error("failed to embed synthesis prompt", "error", err)
return "", err
}
embResp := &models.EmbeddingResp{
Embedding: emb,
Index: 0,
}
topResults, err := r.SearchEmb(embResp)
if err != nil {
r.logger.Error("failed to search for synthesis context", "error", err)
return "", err
}
if len(topResults) > 0 && topResults[0].RawText != synthesisPrompt {
return topResults[0].RawText, nil
}
var finalAnswer strings.Builder
finalAnswer.WriteString("Based on the retrieved context:\n\n")
for i, row := range results {
if i >= 5 {
break
}
finalAnswer.WriteString(fmt.Sprintf("- From %s: %s\n", row.FileName, truncateString(row.RawText, 200)))
}
return finalAnswer.String(), nil
}
func truncateString(s string, maxLen int) string {
if len(s) <= maxLen {
return s
}
return s[:maxLen] + "..."
}
func (r *RAG) Search(query string, limit int) ([]models.VectorRow, error) {
refined := r.RefineQuery(query)
variations := r.GenerateQueryVariations(refined)
allResults := make([]models.VectorRow, 0)
seen := make(map[string]bool)
for _, q := range variations {
emb, err := r.LineToVector(q)
if err != nil {
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)
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)
}
}
}
reranked := r.RerankResults(allResults, query)
if len(reranked) > limit {
reranked = reranked[:limit]
}
return reranked, nil
}
var (
ragInstance *RAG
ragOnce sync.Once
)
func Init(c *config.Config, l *slog.Logger, s storage.FullRepo) error {
ragOnce.Do(func() {
if c == nil || l == nil || s == nil {
return
}
ragInstance = New(l, s, c)
})
return nil
}
func GetInstance() *RAG {
return ragInstance
}

126
tables.go
View File

@@ -247,9 +247,49 @@ func formatSize(size int64) string {
return fmt.Sprintf("%.1f%s", s, units[i]) return fmt.Sprintf("%.1f%s", s, units[i])
} }
func makeRAGTable(fileList []string) *tview.Flex { type ragFileInfo struct {
actions := []string{"load", "delete"} name string
rows, cols := len(fileList), len(actions)+2 inRAGDir bool
isLoaded bool
fullPath string
}
func makeRAGTable(fileList []string, loadedFiles []string) *tview.Flex {
// Build set of loaded files for quick lookup
loadedSet := make(map[string]bool)
for _, f := range loadedFiles {
loadedSet[f] = true
}
// Build merged list: files from ragdir + orphaned files from DB
ragFiles := make([]ragFileInfo, 0, len(fileList)+len(loadedFiles))
seen := make(map[string]bool)
// Add files from ragdir
for _, f := range fileList {
ragFiles = append(ragFiles, ragFileInfo{
name: f,
inRAGDir: true,
isLoaded: loadedSet[f],
fullPath: path.Join(cfg.RAGDir, f),
})
seen[f] = true
}
// Add orphaned files (in DB but not in ragdir)
for _, f := range loadedFiles {
if !seen[f] {
ragFiles = append(ragFiles, ragFileInfo{
name: f,
inRAGDir: false,
isLoaded: true,
fullPath: "",
})
}
}
rows := len(ragFiles)
cols := 4 // File Name | Preview | Action | Delete
fileTable := tview.NewTable(). fileTable := tview.NewTable().
SetBorders(true) SetBorders(true)
longStatusView := tview.NewTextView() longStatusView := tview.NewTextView()
@@ -273,7 +313,7 @@ func makeRAGTable(fileList []string) *tview.Flex {
SetAlign(tview.AlignCenter). SetAlign(tview.AlignCenter).
SetSelectable(false)) SetSelectable(false))
fileTable.SetCell(0, 2, fileTable.SetCell(0, 2,
tview.NewTableCell("Load"). tview.NewTableCell("Load/Unload").
SetTextColor(tcell.ColorWhite). SetTextColor(tcell.ColorWhite).
SetAlign(tview.AlignCenter). SetAlign(tview.AlignCenter).
SetSelectable(false)) SetSelectable(false))
@@ -284,18 +324,29 @@ func makeRAGTable(fileList []string) *tview.Flex {
SetSelectable(false)) SetSelectable(false))
// Add the file rows starting from row 1 // Add the file rows starting from row 1
for r := 0; r < rows; r++ { for r := 0; r < rows; r++ {
f := ragFiles[r]
for c := 0; c < cols; c++ { for c := 0; c < cols; c++ {
color := tcell.ColorWhite color := tcell.ColorWhite
switch { switch {
case c == 0: case c == 0:
displayName := f.name
if !f.inRAGDir {
displayName = f.name + " (orphaned)"
}
fileTable.SetCell(r+1, c, fileTable.SetCell(r+1, c,
tview.NewTableCell(fileList[r]). tview.NewTableCell(displayName).
SetTextColor(color). SetTextColor(color).
SetAlign(tview.AlignCenter). SetAlign(tview.AlignCenter).
SetSelectable(false)) SetSelectable(false))
case c == 1: case c == 1:
fpath := path.Join(cfg.RAGDir, fileList[r]) if !f.inRAGDir {
if fi, err := os.Stat(fpath); err == nil { // Orphaned file - no preview available
fileTable.SetCell(r+1, c,
tview.NewTableCell("not in ragdir").
SetTextColor(tcell.ColorYellow).
SetAlign(tview.AlignCenter).
SetSelectable(false))
} else if fi, err := os.Stat(f.fullPath); err == nil {
size := fi.Size() size := fi.Size()
modTime := fi.ModTime() modTime := fi.ModTime()
preview := fmt.Sprintf("%s | %s", formatSize(size), modTime.Format("2006-01-02 15:04")) preview := fmt.Sprintf("%s | %s", formatSize(size), modTime.Format("2006-01-02 15:04"))
@@ -312,11 +363,27 @@ func makeRAGTable(fileList []string) *tview.Flex {
SetSelectable(false)) SetSelectable(false))
} }
case c == 2: case c == 2:
actionText := "load"
if f.isLoaded {
actionText = "unload"
}
if !f.inRAGDir {
// Orphaned file - can only unload
actionText = "unload"
}
fileTable.SetCell(r+1, c, fileTable.SetCell(r+1, c,
tview.NewTableCell("load"). tview.NewTableCell(actionText).
SetTextColor(color). SetTextColor(color).
SetAlign(tview.AlignCenter)) SetAlign(tview.AlignCenter))
default: case c == 3:
if !f.inRAGDir {
// Orphaned file - cannot delete from ragdir (not there)
fileTable.SetCell(r+1, c,
tview.NewTableCell("-").
SetTextColor(tcell.ColorDarkGray).
SetAlign(tview.AlignCenter).
SetSelectable(false))
} else {
fileTable.SetCell(r+1, c, fileTable.SetCell(r+1, c,
tview.NewTableCell("delete"). tview.NewTableCell("delete").
SetTextColor(color). SetTextColor(color).
@@ -324,6 +391,7 @@ func makeRAGTable(fileList []string) *tview.Flex {
} }
} }
} }
}
errCh := make(chan error, 1) // why? errCh := make(chan error, 1) // why?
go func() { go func() {
defer pages.RemovePage(RAGPage) defer pages.RemovePage(RAGPage)
@@ -376,25 +444,51 @@ func makeRAGTable(fileList []string) *tview.Flex {
pages.RemovePage(RAGPage) pages.RemovePage(RAGPage)
return return
} }
// For file rows, get the filename (row index - 1 because of the exit row at index 0) // For file rows, get the file info (row index - 1 because of the exit row at index 0)
fpath := fileList[row-1] // -1 to account for the exit row at index 0 f := ragFiles[row-1]
// notification := fmt.Sprintf("chat: %s; action: %s", fpath, tc.Text) // Handle "-" case (orphaned file with no delete option)
if tc.Text == "-" {
pages.RemovePage(RAGPage)
return
}
switch tc.Text { switch tc.Text {
case "load": case "load":
fpath = path.Join(cfg.RAGDir, fpath) fpath := path.Join(cfg.RAGDir, f.name)
longStatusView.SetText("clicked load") longStatusView.SetText("clicked load")
go func() { go func() {
if err := ragger.LoadRAG(fpath); err != nil { if err := ragger.LoadRAG(fpath); err != nil {
logger.Error("failed to embed file", "chat", fpath, "error", err) logger.Error("failed to embed file", "chat", fpath, "error", err)
_ = notifyUser("RAG", "failed to embed file; error: "+err.Error()) _ = notifyUser("RAG", "failed to embed file; error: "+err.Error())
errCh <- err app.QueueUpdate(func() {
// pages.RemovePage(RAGPage) pages.RemovePage(RAGPage)
})
return return
} }
_ = notifyUser("RAG", "file loaded successfully")
app.QueueUpdate(func() {
pages.RemovePage(RAGPage)
})
}()
return
case "unload":
longStatusView.SetText("clicked unload")
go func() {
if err := ragger.RemoveFile(f.name); err != nil {
logger.Error("failed to unload file from RAG", "filename", f.name, "error", err)
_ = notifyUser("RAG", "failed to unload file; error: "+err.Error())
app.QueueUpdate(func() {
pages.RemovePage(RAGPage)
})
return
}
_ = notifyUser("RAG", "file unloaded successfully")
app.QueueUpdate(func() {
pages.RemovePage(RAGPage)
})
}() }()
return return
case "delete": case "delete":
fpath = path.Join(cfg.RAGDir, fpath) fpath := path.Join(cfg.RAGDir, f.name)
if err := os.Remove(fpath); err != nil { if err := os.Remove(fpath); err != nil {
logger.Error("failed to delete file", "filename", fpath, "error", err) logger.Error("failed to delete file", "filename", fpath, "error", err)
return return

View File

@@ -16,6 +16,7 @@ import (
"sync" "sync"
"time" "time"
"gf-lt/rag"
"github.com/GrailFinder/searchagent/searcher" "github.com/GrailFinder/searchagent/searcher"
) )
@@ -58,9 +59,9 @@ Your current tools:
"when_to_use": "when asked to search the web for information; returns clean summary without html,css and other web elements; limit is optional (default 3)" "when_to_use": "when asked to search the web for information; returns clean summary without html,css and other web elements; limit is optional (default 3)"
}, },
{ {
"name":"websearch_raw", "name":"rag_search",
"args": ["query", "limit"], "args": ["query", "limit"],
"when_to_use": "when asked to search the web for information; returns raw data as is without processing; limit is optional (default 3)" "when_to_use": "when asked to search the local document database for information; performs query refinement, semantic search, reranking, and synthesis; returns clean summary with sources; limit is optional (default 3)"
}, },
{ {
"name":"read_url", "name":"read_url",
@@ -146,6 +147,7 @@ under the topic: Adam's number is stored:
After that you are free to respond to the user. After that you are free to respond to the user.
` `
webSearchSysPrompt = `Summarize the web search results, extracting key information and presenting a concise answer. Provide sources and URLs where relevant.` webSearchSysPrompt = `Summarize the web search results, extracting key information and presenting a concise answer. Provide sources and URLs where relevant.`
ragSearchSysPrompt = `Synthesize the document search results, extracting key information and presenting a concise answer. Provide sources and document IDs where relevant.`
readURLSysPrompt = `Extract and summarize the content from the webpage. Provide key information, main points, and any relevant details.` readURLSysPrompt = `Extract and summarize the content from the webpage. Provide key information, main points, and any relevant details.`
summarySysPrompt = `Please provide a concise summary of the following conversation. Focus on key points, decisions, and actions. Provide only the summary, no additional commentary.` summarySysPrompt = `Please provide a concise summary of the following conversation. Focus on key points, decisions, and actions. Provide only the summary, no additional commentary.`
basicCard = &models.CharCard{ basicCard = &models.CharCard{
@@ -170,6 +172,10 @@ func init() {
panic("failed to init seachagent; error: " + err.Error()) panic("failed to init seachagent; error: " + err.Error())
} }
WebSearcher = sa WebSearcher = sa
if err := rag.Init(cfg, logger, store); err != nil {
logger.Warn("failed to init rag; rag_search tool will not be available", "error", err)
}
} }
// getWebAgentClient returns a singleton AgentClient for web agents. // getWebAgentClient returns a singleton AgentClient for web agents.
@@ -196,6 +202,8 @@ func getWebAgentClient() *agent.AgentClient {
func registerWebAgents() { func registerWebAgents() {
webAgentsOnce.Do(func() { webAgentsOnce.Do(func() {
client := getWebAgentClient() client := getWebAgentClient()
// Register rag_search agent
agent.Register("rag_search", agent.NewWebAgentB(client, ragSearchSysPrompt))
// Register websearch agent // Register websearch agent
agent.Register("websearch", agent.NewWebAgentB(client, webSearchSysPrompt)) agent.Register("websearch", agent.NewWebAgentB(client, webSearchSysPrompt))
// Register read_url agent // Register read_url agent
@@ -239,6 +247,48 @@ func websearch(args map[string]string) []byte {
return data return data
} }
// rag search (searches local document database)
func ragsearch(args map[string]string) []byte {
query, ok := args["query"]
if !ok || query == "" {
msg := "query not provided to rag_search tool"
logger.Error(msg)
return []byte(msg)
}
limitS, ok := args["limit"]
if !ok || limitS == "" {
limitS = "3"
}
limit, err := strconv.Atoi(limitS)
if err != nil || limit == 0 {
logger.Warn("ragsearch limit; passed bad value; setting to default (3)",
"limit_arg", limitS, "error", err)
limit = 3
}
ragInstance := rag.GetInstance()
if ragInstance == nil {
msg := "rag not initialized; rag_search tool is not available"
logger.Error(msg)
return []byte(msg)
}
results, err := ragInstance.Search(query, limit)
if err != nil {
msg := "rag search failed; error: " + err.Error()
logger.Error(msg)
return []byte(msg)
}
data, err := json.Marshal(results)
if err != nil {
msg := "failed to marshal rag search result; error: " + err.Error()
logger.Error(msg)
return []byte(msg)
}
return data
}
// web search raw (returns raw data without processing) // web search raw (returns raw data without processing)
func websearchRaw(args map[string]string) []byte { func websearchRaw(args map[string]string) []byte {
// make http request return bytes // make http request return bytes
@@ -997,6 +1047,7 @@ var fnMap = map[string]fnSig{
"recall": recall, "recall": recall,
"recall_topics": recallTopics, "recall_topics": recallTopics,
"memorise": memorise, "memorise": memorise,
"rag_search": ragsearch,
"websearch": websearch, "websearch": websearch,
"websearch_raw": websearchRaw, "websearch_raw": websearchRaw,
"read_url": readURL, "read_url": readURL,
@@ -1033,6 +1084,28 @@ func callToolWithAgent(name string, args map[string]string) []byte {
// openai style def // openai style def
var baseTools = []models.Tool{ var baseTools = []models.Tool{
// rag_search
models.Tool{
Type: "function",
Function: models.ToolFunc{
Name: "rag_search",
Description: "Search local document database given query, limit of sources (default 3). Performs query refinement, semantic search, reranking, and synthesis.",
Parameters: models.ToolFuncParams{
Type: "object",
Required: []string{"query", "limit"},
Properties: map[string]models.ToolArgProps{
"query": models.ToolArgProps{
Type: "string",
Description: "search query",
},
"limit": models.ToolArgProps{
Type: "string",
Description: "limit of the document results",
},
},
},
},
},
// websearch // websearch
models.Tool{ models.Tool{
Type: "function", Type: "function",

19
tui.go
View File

@@ -913,6 +913,7 @@ func init() {
return nil return nil
} }
} }
// Get files from ragdir
fileList := []string{} fileList := []string{}
for _, f := range files { for _, f := range files {
if f.IsDir() { if f.IsDir() {
@@ -920,22 +921,14 @@ func init() {
} }
fileList = append(fileList, f.Name()) fileList = append(fileList, f.Name())
} }
chatRAGTable := makeRAGTable(fileList) // Get loaded files from vector DB
pages.AddPage(RAGPage, chatRAGTable, true, true) loadedFiles, err := ragger.ListLoaded()
return nil
}
if event.Key() == tcell.KeyCtrlY { // Use Ctrl+Y to list loaded RAG files
// List files already loaded into the RAG system
fileList, err := ragger.ListLoaded()
if err != nil { if err != nil {
logger.Error("failed to list loaded RAG files", "error", err) logger.Error("failed to list loaded RAG files", "error", err)
if notifyerr := notifyUser("failed to list RAG files", err.Error()); notifyerr != nil { loadedFiles = []string{} // Continue with empty list on error
logger.Error("failed to send notification", "error", notifyerr)
} }
return nil chatRAGTable := makeRAGTable(fileList, loadedFiles)
} pages.AddPage(RAGPage, chatRAGTable, true, true)
chatLoadedRAGTable := makeLoadedRAGTable(fileList)
pages.AddPage(RAGLoadedPage, chatLoadedRAGTable, true, true)
return nil return nil
} }
if event.Key() == tcell.KeyRune && event.Modifiers() == tcell.ModAlt && event.Rune() == '1' { if event.Key() == tcell.KeyRune && event.Modifiers() == tcell.ModAlt && event.Rune() == '1' {