From 5d7ace205abfa6a6261f71883474d68b306d51db Mon Sep 17 00:00:00 2001 From: Ian Gulliver Date: Wed, 27 Nov 2024 22:29:03 -0600 Subject: [PATCH] OpenAI-based short link suggestions --- main.go | 54 ++++++++++++++++++++++++++- openai.go | 94 +++++++++++++++++++++++++++++++++++++++++++++++ static/index.html | 91 ++++++++++++++++++++++++++++++++------------- 3 files changed, 211 insertions(+), 28 deletions(-) create mode 100644 openai.go diff --git a/main.go b/main.go index 7c791dc..bd1f919 100644 --- a/main.go +++ b/main.go @@ -7,6 +7,7 @@ import ( "log" "net/http" "os" + "strings" "time" _ "github.com/lib/pq" @@ -18,12 +19,17 @@ type ShortLinks struct { mux *http.ServeMux db *sql.DB r *rand.Rand + oai *oaiClient } -type response struct { +type setResponse struct { Short string `json:"short"` } +type suggestResponse struct { + Shorts []string `json:"shorts"` +} + func NewShortLinks(db *sql.DB) (*ShortLinks, error) { tmpl := template.New("index.html") @@ -32,16 +38,23 @@ func NewShortLinks(db *sql.DB) (*ShortLinks, error) { return nil, fmt.Errorf("static/index.html: %w", err) } + oai, err := newOAIClientFromEnv() + if err != nil { + return nil, fmt.Errorf("newOAIClientFromEnv: %w", err) + } + sl := &ShortLinks{ tmpl: tmpl, mux: http.NewServeMux(), db: db, r: rand.New(rand.NewSource(uint64(time.Now().UnixNano()))), + oai: oai, } sl.mux.HandleFunc("GET /{$}", sl.serveRoot) sl.mux.HandleFunc("GET /{short}", sl.serveShort) sl.mux.HandleFunc("POST /{$}", sl.serveSet) + sl.mux.HandleFunc("QUERY /{$}", sl.serveSuggest) return sl, nil } @@ -117,11 +130,48 @@ DO UPDATE SET long = $2; return } - sendJSON(w, response{ + sendJSON(w, setResponse{ Short: short, }) } +func (sl *ShortLinks) serveSuggest(w http.ResponseWriter, r *http.Request) { + err := r.ParseForm() + if err != nil { + sendError(w, http.StatusBadRequest, "Parse form: %s", err) + return + } + + log.Printf("%s %s %s", r.RemoteAddr, r.URL.Path, r.Form.Encode()) + + if !r.Form.Has("short") { + sendError(w, http.StatusBadRequest, "short= param required") + return + } + + user := strings.Join(r.Form["short"], "\n") + + comp, err := sl.oai.completeChat( + "You are an assistant helping a user choose useful short names for a URL shortener. The request contains a list recents names chosen by the user, separated by newlines, with the most recent names first. Respond with only a list of possible suggestions for additional short names, separated by newlines. Suggestions may include conceptual variations of the names provided, plural/singular variations, hyphenation variations, or other variations that are likely to be useful. Your bar for suggestions should be relatively high; responding with a short list of high quality suggestions is preferred.", + user, + ) + if err != nil { + sendError(w, http.StatusInternalServerError, "oai.completeChat: %s", err) + return + } + + shorts := []string{} + for _, short := range strings.Split(comp, "\n") { + if short != "" { + shorts = append(shorts, strings.TrimSpace(short)) + } + } + + sendJSON(w, suggestResponse{ + Shorts: shorts, + }) +} + func (sl *ShortLinks) genShort() (string, error) { for chars := 3; chars <= 10; chars++ { b := make([]byte, chars) diff --git a/openai.go b/openai.go new file mode 100644 index 0000000..a450b7b --- /dev/null +++ b/openai.go @@ -0,0 +1,94 @@ +package main + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "os" +) + +type oaiClient struct { + c *http.Client + apiKey string +} + +type oaiRequest struct { + Model string `json:"model"` + Messages []oaiMessage `json:"messages"` +} + +type oaiMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +type oaiResponse struct { + Choices []oaiChoice `json:"choices"` +} + +type oaiChoice struct { + Message oaiMessage `json:"message"` +} + +func newOAIClient(apiKey string) *oaiClient { + return &oaiClient{ + c: &http.Client{}, + apiKey: apiKey, + } +} + +func newOAIClientFromEnv() (*oaiClient, error) { + apiKey := os.Getenv("OPENAI_API_KEY") + if apiKey == "" { + return nil, fmt.Errorf("OPENAI_API_KEY is not set") + } + + return newOAIClient(apiKey), nil +} + +func (oai *oaiClient) completeChat(system, user string) (string, error) { + buf := &bytes.Buffer{} + err := json.NewEncoder(buf).Encode(&oaiRequest{ + Model: "gpt-4o", + Messages: []oaiMessage{ + {Role: "system", Content: system}, + {Role: "user", Content: user}, + }, + }) + + if err != nil { + return "", err + } + + req, err := http.NewRequest("POST", "https://api.openai.com/v1/chat/completions", buf) + if err != nil { + return "", err + } + + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", oai.apiKey)) + req.Header.Set("Content-Type", "application/json") + + resp, err := oai.c.Do(req) + if err != nil { + return "", err + } + + defer resp.Body.Close() + + if resp.StatusCode != 200 { + body, _ := io.ReadAll(resp.Body) + return "", fmt.Errorf("%s", string(body)) + } + + dec := json.NewDecoder(resp.Body) + var res oaiResponse + + err = dec.Decode(&res) + if err != nil { + return "", err + } + + return res.Choices[0].Message.Content, nil +} diff --git a/static/index.html b/static/index.html index d84fd8c..d55c51b 100644 --- a/static/index.html +++ b/static/index.html @@ -76,7 +76,7 @@ function error(err1, err2) { document.getElementById('err').show(); } -async function set() { +async function setFromInputs() { const short = document.getElementById('short').value; const long = document.getElementById('long').value; @@ -88,6 +88,10 @@ async function set() { document.getElementById('short-icon').setAttribute('name', 'check-square-fill'); document.getElementById('long-icon').setAttribute('name', 'check-square-fill'); + set(short, long); +} + +function set(short, long) { if (short != '') { setShortItem(short, 'check-square-fill'); } @@ -96,23 +100,42 @@ async function set() { params.set('short', short); params.set('long', long); - const resp = await fetch(`./?${params.toString()}`, { + fetch(`./?${params.toString()}`, { method: 'POST', + }).then(async (resp) => { + if (resp.status !== 200) { + error('Failed to set', (await resp.json()).message); + return; + } + + const data = await resp.json(); + const newShort = data.short; + + setShortItem(newShort, 'check-square'); + + if (document.getElementById('short').value == short && document.getElementById('long').value == long) { + document.getElementById('short-icon').setAttribute('name', 'check-square'); + document.getElementById('long-icon').setAttribute('name', 'check-square'); + await navigator.clipboard.writeText(`${window.location.origin}/${newShort}`); + } }); - if (resp.status !== 200) { - error('Failed to set', (await resp.json()).message); - return; + const suggestParams = new URLSearchParams(); + for (const elem of document.getElementById('tree').children) { + const icon = elem.getElementsByTagName('sl-icon')[0]; + if (icon.getAttribute('name') == 'check-square-fill' || + icon.getAttribute('name') == 'check-square') { + suggestParams.append('short', elem.textContent); + } } - const newShort = (await resp.json()).short; - - if (document.getElementById('short').value == short && document.getElementById('long').value == long) { - document.getElementById('short-icon').setAttribute('name', 'check-square'); - document.getElementById('long-icon').setAttribute('name', 'check-square'); - setShortItem(newShort, 'check-square'); - await navigator.clipboard.writeText(`${window.location.origin}/${newShort}`); - } + fetch(`./?${suggestParams.toString()}`, { + method: 'QUERY', + }).then(async (resp) => { + for (const short of (await resp.json()).shorts) { + appendShortItem(short, long); + } + }); } function setShortItem(short, icon) { @@ -125,18 +148,34 @@ function setShortItem(short, icon) { } const item = document.createElement('sl-tree-item'); - const url = `${window.location.origin}/${short}`; item.appendChild(document.createElement('sl-icon')).setAttribute('name', icon); item.appendChild(document.createTextNode(short)); - - const copy = document.createElement('sl-copy-button'); - copy.setAttribute('value', url); - copy.style.color = 'var(--sl-color-neutral-300)'; - item.appendChild(copy); + item.addEventListener('click', () => { + navigator.clipboard.writeText(`${window.location.origin}/${short}`); + }); tree.insertBefore(item, tree.firstChild); } +function appendShortItem(short, long) { + const tree = document.getElementById('tree'); + + for (const item of tree.children) { + if (item.textContent == short) { + return; + } + } + + const item = document.createElement('sl-tree-item'); + item.appendChild(document.createElement('sl-icon')).setAttribute('name', 'square'); + item.appendChild(document.createTextNode(short)); + item.addEventListener('click', () => { + set(short, long); + }); + + tree.appendChild(item); +} + document.addEventListener('DOMContentLoaded', async () => { await Promise.all([ customElements.whenDefined('sl-input'), @@ -155,18 +194,18 @@ document.addEventListener('DOMContentLoaded', async () => { document.getElementById('short').addEventListener('keydown', (e) => { if (e.key === 'Enter') { - set(); + setFromInputs(); } }); document.getElementById('short').addEventListener('paste', () => { if (document.getElementById('long').value != '') { - setTimeout(() => set(), 0); + setTimeout(() => setFromInputs(), 0); } }); document.getElementById('short-icon').addEventListener('click', () => { - set(); + setFromInputs(); }); @@ -177,7 +216,7 @@ document.addEventListener('DOMContentLoaded', async () => { document.getElementById('long').addEventListener('keydown', (e) => { if (e.key === 'Enter') { - set(); + setFromInputs(); } else { document.getElementById('tree').replaceChildren(); } @@ -185,17 +224,17 @@ document.addEventListener('DOMContentLoaded', async () => { document.getElementById('long').addEventListener('paste', () => { if (document.getElementById('short').value != '') { - setTimeout(() => set(), 0); + setTimeout(() => setFromInputs(), 0); } }); document.getElementById('long-icon').addEventListener('click', () => { - set(); + setFromInputs(); }); document.getElementById('set').addEventListener('click', () => { - set(); + setFromInputs(); });