157 lines
3.6 KiB
Go
157 lines
3.6 KiB
Go
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
|
|
}
|