Mason Payne
e0430c62bd
set up the database with full text search and a simple prompt template engine
70 lines
1.5 KiB
Go
70 lines
1.5 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"ctxGPT/database"
|
|
"fmt"
|
|
"text/template"
|
|
|
|
"github.com/pkoukk/tiktoken-go"
|
|
)
|
|
|
|
const encodingName = "gpt-4"
|
|
|
|
func main() {
|
|
|
|
db, err := database.NewDB()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
defer db.Close()
|
|
value, err := db.Get(context.Background(), "context1")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
fmt.Println(value)
|
|
|
|
err = db.Save(context.Background(), "context2", "value2")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
// to get text out of PDF, DOC, DOCX, XML, HTML, RTF, ODT pages documents and images to plain text
|
|
// use https://github.com/sajari/docconv
|
|
|
|
summarizeConvoPrompt, err := BuildPrompt("summarize.tmpl", struct{ WordLimit int }{WordLimit: 100})
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
fmt.Println(summarizeConvoPrompt)
|
|
tokenCount, err := GetTokenCount(summarizeConvoPrompt)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
fmt.Println(tokenCount)
|
|
}
|
|
|
|
func BuildPrompt(name string, in interface{}) (string, error) {
|
|
fileLocation := "./prompts/" + name
|
|
tmpl, err := template.New(name).ParseFiles(fileLocation)
|
|
if err != nil {
|
|
return "", fmt.Errorf("error parsing template: %w", err)
|
|
}
|
|
b := bytes.Buffer{}
|
|
err = tmpl.Execute(&b, in)
|
|
if err != nil {
|
|
return "", fmt.Errorf("error executing template: %w", err)
|
|
}
|
|
return b.String(), nil
|
|
}
|
|
|
|
func GetTokenCount(input string) (int, error) {
|
|
tke, err := tiktoken.EncodingForModel(encodingName) // cached in "TIKTOKEN_CACHE_DIR"
|
|
if err != nil {
|
|
return 0, fmt.Errorf("error getting encoding: %w", err)
|
|
}
|
|
token := tke.Encode(input, nil, nil)
|
|
return len(token), nil
|
|
}
|