OpenAI-based short link suggestions

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)

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
+}

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 newShort = (await resp.json()).short;
+        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');
-        setShortItem(newShort, 'check-square');
             await navigator.clipboard.writeText(`${window.location.origin}/${newShort}`);
         }
+    });
+
+    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);
+        }
+    }
+
+    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));
-
+    item.addEventListener('click', () => {
-    const copy = document.createElement('sl-copy-button');
+        navigator.clipboard.writeText(`${window.location.origin}/${short}`);
-    copy.setAttribute('value', url);
+    });
-    copy.style.color = 'var(--sl-color-neutral-300)';
-    item.appendChild(copy);
 
     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();
     });