KB: add webui sections

example/webui/rag.go

@@ -0,0 +1,156 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strings"
+	"sync"
+
+	"jaytaylor.com/html2text"
+
+	"github.com/mudler/local-agent-framework/llm"
+	sitemap "github.com/oxffaa/gopher-parse-sitemap"
+)
+
+type InMemoryDatabase struct {
+	sync.Mutex
+	Database []string
+	path     string
+}
+
+func loadDB(path string) ([]string, error) {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return nil, err
+	}
+
+	poolData := []string{}
+	err = json.Unmarshal(data, &poolData)
+	return poolData, err
+}
+
+func NewInMemoryDB(knowledgebase string) (*InMemoryDatabase, error) {
+	// if file exists, try to load an existing pool.
+	// if file does not exist, create a new pool.
+
+	poolfile := filepath.Join(knowledgebase, "knowledgebase.json")
+
+	if _, err := os.Stat(poolfile); err != nil {
+		// file does not exist, return a new pool
+		return &InMemoryDatabase{
+			Database: []string{},
+			path:     poolfile,
+		}, nil
+	}
+
+	poolData, err := loadDB(poolfile)
+	if err != nil {
+		return nil, err
+	}
+	return &InMemoryDatabase{
+		Database: poolData,
+		path:     poolfile,
+	}, nil
+}
+
+func (db *InMemoryDatabase) SaveToStore(apiKey string, apiURL string) error {
+	for _, d := range db.Database {
+		lai := llm.NewClient(apiKey, apiURL+"/v1")
+		laiStore := llm.NewStoreClient(apiURL, apiKey)
+
+		err := llm.StoreStringEmbeddingInVectorDB(laiStore, lai, d)
+		if err != nil {
+			return fmt.Errorf("Error storing in the KB: %w", err)
+		}
+	}
+
+	return nil
+}
+func (db *InMemoryDatabase) AddEntry(entry string) error {
+	db.Lock()
+	defer db.Unlock()
+	db.Database = append(db.Database, entry)
+	return nil
+}
+
+func (db *InMemoryDatabase) SaveDB() error {
+	db.Lock()
+	defer db.Unlock()
+	data, err := json.Marshal(db.Database)
+	if err != nil {
+		return err
+	}
+
+	err = os.WriteFile(db.path, data, 0644)
+	return err
+}
+
+func getWebPage(url string) (string, error) {
+	resp, err := http.Get(url)
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return "", err
+	}
+	return html2text.FromString(string(body), html2text.Options{PrettyTables: true})
+}
+
+func Sitemap(url string) (res []string, err error) {
+	err = sitemap.ParseFromSite(url, func(e sitemap.Entry) error {
+		fmt.Println("Sitemap page: " + e.GetLocation())
+		content, err := getWebPage(e.GetLocation())
+		if err == nil {
+			res = append(res, content)
+		}
+		return nil
+	})
+	return
+}
+
+// splitParagraphIntoChunks takes a paragraph and a maxChunkSize as input,
+// and returns a slice of strings where each string is a chunk of the paragraph
+// that is at most maxChunkSize long, ensuring that words are not split.
+func splitParagraphIntoChunks(paragraph string, maxChunkSize int) []string {
+	// Check if the paragraph length is less than or equal to maxChunkSize.
+	// If so, return the paragraph as the only chunk.
+	if len(paragraph) <= maxChunkSize {
+		return []string{paragraph}
+	}
+
+	var chunks []string
+	var currentChunk strings.Builder
+
+	words := strings.Fields(paragraph) // Splits the paragraph into words.
+
+	for _, word := range words {
+		// Check if adding the next word would exceed the maxChunkSize.
+		// If so, add the currentChunk to the chunks slice and start a new chunk.
+		if currentChunk.Len()+len(word) > maxChunkSize {
+			chunks = append(chunks, currentChunk.String())
+			currentChunk.Reset()
+		}
+
+		// Add a space before the word if it's not the beginning of a new chunk.
+		if currentChunk.Len() > 0 {
+			currentChunk.WriteString(" ")
+		}
+
+		// Add the word to the current chunk.
+		currentChunk.WriteString(word)
+	}
+
+	// Add the last chunk if it's not empty.
+	if currentChunk.Len() > 0 {
+		chunks = append(chunks, currentChunk.String())
+	}
+
+	return chunks
+}