Add AI analysis of collocations feature (#96)

This change allows eligible users to get AI explanations and
generate example sentences based on the collocations present
in the sankey diagrams. Such collocations often help illustrate 
grammar rules or common speech patterns.

This is similar to TrieLingual albeit in a different diagram
form.

The prompts could likely use some tweaking, and prices for Gemini
2.5 flash seem comparable to Gemini 2.0 flash, so upgrading seems
to make sense; those changes will be done in a separate commit.

Author: mreichhoff

functions/prompts/analyze-collocation.prompt

@@ -0,0 +1,15 @@
+---
+model: vertexai/gemini-2.0-flash-001
+input:
+  schema:
+    collocation: string
+output:
+  schema: AnalyzeCollocationSchema
+---
+{{role "system"}}
+You are a helpful Chinese teacher for speakers of English who want to learn Chinese. You speak naturally, and you provide helpful sentences that illustrate how to use Chinese vocabulary.
+{{role "user"}}
+Please generate three example Chinese sentences, each with a separate English translation and pinyin, that uses the phrase "{{collocation}}":
+Each sentence must include "{{collocation}}".
+
+Please also translate "{{collocation}}" to English and provide a plain-text explanation of how such a phrase would be used.

functions/src/index.ts

@@ -9,6 +9,7 @@ import {
     imageAnalysisSchema,
     chineseSentenceGenerationSchema,
     generateChineseSentencesInputSchema,
+    analyzeCollocationSchema,
 } from "./schema";
 
 let firebaseApp: admin.app.App;
@@ -130,3 +131,29 @@ const generateChineseSentencesFlow = ai.defineFlow({
 });
 
 export const generateChineseSentences = onCallGenkit(generateChineseSentencesFlow);
+
+const AnalyzeCollocationSchema = ai.defineSchema('AnalyzeCollocationSchema', analyzeCollocationSchema);
+const analyzeCollocationPrompt = ai.prompt<z.ZodTypeAny, typeof AnalyzeCollocationSchema>('analyze-collocation');
+
+const analyzeCollocationFlow = ai.defineFlow({
+    name: "analyzeCollocation",
+    inputSchema: z.string(),
+    outputSchema: analyzeCollocationSchema,
+}, async (collocation, { context }) => {
+    if (!firebaseApp) {
+        firebaseApp = admin.initializeApp();
+    }
+    const isAuthorized = await isUserAuthorized(context);
+    if (!isAuthorized) {
+        throw new HttpsError("permission-denied", "user not authorized");
+    }
+    collocation = collocation.replaceAll(' ', '');
+    const { output } = await analyzeCollocationPrompt({ collocation });
+    if (!output) {
+        throw new HttpsError("internal", 'oh no, the model like, failed?');
+    }
+    return output;
+},
+);
+
+export const analyzeCollocation = onCallGenkit(analyzeCollocationFlow);

functions/src/schema.ts

@@ -42,11 +42,19 @@ export const generateChineseSentencesInputSchema = z.object({
     definitions: z.array(z.string()),
 });
 
+const sentenceSchema = z.array(
+    z.object({
+        chineseTextWithoutPinyin: z.string(),
+        pinyin: z.string(),
+        englishTranslation: z.string(),
+    }));
+
 export const chineseSentenceGenerationSchema = z.object({
-    sentences: z.array(
-        z.object({
-            chineseTextWithoutPinyin: z.string(),
-            pinyin: z.string(),
-            englishTranslation: z.string(),
-        })),
+    sentences: sentenceSchema,
+});
+
+export const analyzeCollocationSchema = z.object({
+    englishTranslation: z.string(),
+    plainTextExplanation: z.string(),
+    sentences: sentenceSchema,
 });

functions/tsconfig.json

@@ -9,7 +9,7 @@
     "outDir": "lib",
     "sourceMap": true,
     "strict": true,
-    "target": "es2017",
+    "target": "es2021",
     "skipLibCheck": true,
     "esModuleInterop": true,
     "moduleResolution": "nodenext",

public/css/hanzi-graph.css

@@ -342,6 +342,29 @@ https://developer.mozilla.org/en-US/docs/Web/API/File_API/Using_files_from_web_a
     margin: 0 28px 2px 4px;
 }
 
+.inline-button-container {
+    text-align: center;
+}
+
+.inline-menu-item {
+    line-height: 40px;
+    font-weight: bold;
+    cursor: pointer;
+    width: fit-content;
+    padding: 0 8px;
+    background: var(--popover-background-color);
+    border-radius: 24px;
+    border: var(--border);
+    margin: 6px;
+    display: inline-block;
+}
+
+.inline-menu-item .ai-icon {
+    position: relative;
+    top: 2px;
+    margin: 0 10px 0 4px
+}
+
 .open-button {
     opacity: 0.5;
 }

public/js/modules/data-layer.js

@@ -619,6 +619,21 @@ async function generateChineseSentences(word, definitions) {
     return tokenizedSentences;
 }
 
+async function analyzeCollocation(collocation) {
+    const functions = getFunctions();
+    // connectFunctionsEmulator(functions, "127.0.0.1", 5001);
+    const analyzeCollocation = httpsCallable(functions, 'analyzeCollocation');
+    const aiData = await analyzeCollocation(collocation);
+    // same explanation of goofiness as `generateChineseSentences`. In the collocation case, we both generate sentences
+    // and get an explanation back. Tokenize the sentences, then return the whole thing.
+    const tokenizedSentences = await new Promise((resolve, reject) => {
+        document.dispatchEvent(new CustomEvent('sentence-generation-response', { detail: { aiData, collocation, resolve, reject } }));
+    });
+    // i know i shouldn't
+    aiData.data['sentences'] = tokenizedSentences;
+    return aiData;
+}
+
 function isAiEligible() {
     return aiEligible;
 }
@@ -631,4 +646,4 @@ async function analyzeImage(base64ImageContents) {
     return result;
 }
 
-export { writeExploreState, readExploreState, writeOptionState, readOptionState, registerCallback, saveStudyList, addCard, inStudyList, countWordsWithoutCards, getStudyList, isFlashCardUser, removeFromStudyList, findOtherCards, updateCard, recordEvent, getStudyResults, explainChineseSentence, translateEnglish, analyzeImage, generateChineseSentences, isAiEligible, hasCardWithWord, initialize, studyResult, dataTypes, cardTypes }
+export { writeExploreState, readExploreState, writeOptionState, readOptionState, registerCallback, saveStudyList, addCard, inStudyList, countWordsWithoutCards, getStudyList, isFlashCardUser, removeFromStudyList, findOtherCards, updateCard, recordEvent, getStudyResults, explainChineseSentence, translateEnglish, analyzeImage, generateChineseSentences, analyzeCollocation, isAiEligible, hasCardWithWord, initialize, studyResult, dataTypes, cardTypes }

public/js/modules/dom.js

@@ -35,4 +35,12 @@ function writeSeoMetaTags(urlState, graphDisplay) {
     }
 }
 
-export { hanziBox, notFoundElement, walkThrough, examplesList, searchControl, writeSeoMetaTags }
+function createLoadingDots() {
+    const loadingDots = document.createElement('div');
+    loadingDots.classList.add('loading-dots');
+    // there uh....there's four dots, ok?
+    loadingDots.innerHTML = '<div></div><div></div><div></div><div></div>';
+    return loadingDots;
+}
+
+export { hanziBox, notFoundElement, walkThrough, examplesList, searchControl, createLoadingDots, writeSeoMetaTags }

public/js/modules/explore.js

@@ -1,5 +1,5 @@
 import { writeExploreState, addCard, inStudyList, isFlashCardUser, explainChineseSentence, generateChineseSentences, isAiEligible, countWordsWithoutCards, hasCardWithWord, registerCallback, dataTypes } from "./data-layer.js";
-import { hanziBox, notFoundElement, walkThrough, examplesList } from "./dom.js";
+import { hanziBox, notFoundElement, walkThrough, examplesList, createLoadingDots } from "./dom.js";
 import { getActiveGraph, getPartition } from "./options.js";
 import { renderCoverageGraph } from "./coverage-graph"
 import { diagramKeys, switchDiagramView, showNotification } from "./ui-orchestrator.js";
@@ -949,6 +949,19 @@ function renderExplanation(explanation, container) {
     container.appendChild(explanationContainer);
 }
 
+// TODO combine with renderExplanation (or like, just use a framework)
+function renderStandaloneTranslation(translation, term, container) {
+    const translationHeader = document.createElement('h3');
+    translationHeader.innerText = `Translation`;
+
+    const translationContainer = document.createElement('div');
+    translationContainer.classList.add('emphasized-but-not-that-emphasized');
+    translationContainer.innerText = `"${term}" translates as: "${translation}".`;
+
+    container.appendChild(translationHeader);
+    container.appendChild(translationContainer);
+}
+
 function renderAiExplanationResponse(words, response, container) {
     // TODO: error states
     const data = response.data;
@@ -974,14 +987,6 @@ function renderAiExplanationResponse(words, response, container) {
     container.appendChild(grammarContainer);
 }
 
-function createLoadingDots() {
-    const loadingDots = document.createElement('div');
-    loadingDots.classList.add('loading-dots');
-    // there uh....there's four dots, ok?
-    loadingDots.innerHTML = '<div></div><div></div><div></div><div></div>';
-    return loadingDots;
-}
-
 let setupExamples = function (words, type, skipState, allowExplain, aiData) {
     currentExamples = {};
     currentDefinitions = {};
@@ -1162,6 +1167,22 @@ let initialize = function () {
         // resizing the screen throws off positioning of the menu...just hide it
         menuPopover.hidePopover();
     });
+    // TODO this does not belong here...move the AI stuff to its own module, and common rendering functions
+    // to another.
+    document.addEventListener('collocation-analysis-response', function (event) {
+        const explanationContainer = document.createElement('div');
+        const translationContainer = document.createElement('div');
+        explanationContainer.classList.add('ai-explanation');
+        renderStandaloneTranslation(event.detail.aiData.data.englishTranslation, event.detail.collocation, translationContainer);
+        explanationContainer.appendChild(translationContainer);
+        const explanation = event.detail.aiData.data.plainTextExplanation;
+        event.detail.aiContainer.appendChild(explanationContainer);
+        renderExplanation(explanation, explanationContainer);
+        const exampleContainer = document.createElement('ul');
+        exampleContainer.classList.add('examples');
+        event.detail.aiContainer.appendChild(exampleContainer);
+        setupExampleElements(null, event.detail.aiData.data.sentences, exampleContainer, 'ai');
+    });
     voice = getVoice();
     fetchStats();
     registerCallback(dataTypes.studyList, function () {

public/js/modules/flow-diagram.js

@@ -3,6 +3,8 @@ import { sankey, sankeyLinkHorizontal, sankeyCenter, sankeyJustify, sankeyRight,
 import { map, schemeTableau10, union, scaleOrdinal, format as d3format, create } from "d3";
 import { getPartition, getActiveGraph } from "./options";
 import { faqTypes, showFaq } from "./faq";
+import { isAiEligible, analyzeCollocation } from "./data-layer";
+import { createLoadingDots } from "./dom";
 
 function addToTrie(trie, collocation, count, term, maxDepth) {
     let words = collocation.split(' ');
@@ -132,7 +134,7 @@ function renderCollocationData(term, collocations, nextSibling, container) {
     let description = document.createElement('p');
     description.className = 'collocations-detail';
     // TODO: assumption of ranks being present ok for now, but should be switched (well, a couple refactors would be good there)
-    description.innerHTML = `When you see <span class="emphasized freq${getFrequencyLevel(wordSet[term], getActiveGraph().ranks)}">${term}</span>, it's often used with:`;
+    description.innerHTML = `<span class="emphasized freq${getFrequencyLevel(wordSet[term], getActiveGraph().ranks)}">${term}</span> is often used with:`;
     for (const collocation of sorted) {
         let collocationsContainer = document.createElement('p');
         collocationsContainer.classList.add('collocation');
@@ -158,6 +160,10 @@ async function renderUsageDiagram(term, container) {
     // TODO(refactor): consolidate explanation classes
     explanation.classList.add('flow-explanation');
     container.appendChild(explanation);
+    // will be empty unless the user is eligible
+    const aiContainer = document.createElement('div');
+    aiContainer.classList.add('inline-button-container');
+    container.appendChild(aiContainer);
     explanation.innerText = 'Loading...';
     let count = 0;
     let loadingIndicator = setInterval(function () {
@@ -203,6 +209,32 @@ async function renderUsageDiagram(term, container) {
         nodeAlign: 'center',
         linkTitle: d => `${elements.labels[d.source.id]} ${elements.labels[d.target.id]}: ${d.value}`,
         linkClickHandler: (d, i) => {
+            if (isAiEligible()) {
+                aiContainer.innerHTML = '';
+                aiContainer.classList.remove('ai-explanation-container');
+                aiContainer.classList.add('inline-button-container');
+                const collocationsAtClickedNode = elements.collocations[i.id];
+                for (const collocation of collocationsAtClickedNode) {
+                    const collocationSentencesContainer = document.createElement('div');
+                    collocationSentencesContainer.classList.add('inline-menu-item');
+                    const aiIcon = document.createElement('span');
+                    aiIcon.classList.add('ai-icon');
+                    const sentenceButton = document.createElement('span');
+                    sentenceButton.innerText = `Analyze "${collocation}"`;
+                    collocationSentencesContainer.appendChild(aiIcon);
+                    collocationSentencesContainer.appendChild(sentenceButton);
+                    collocationSentencesContainer.addEventListener('click', async function () {
+                        aiContainer.innerHTML = '';
+                        aiContainer.classList.remove('inline-button-container');
+                        aiContainer.appendChild(createLoadingDots());
+                        const aiData = await analyzeCollocation(collocation);
+                        aiContainer.innerHTML = '';
+                        aiContainer.classList.add('ai-explanation-container');
+                        document.dispatchEvent(new CustomEvent('collocation-analysis-response', { detail: { aiData, term, collocation, aiContainer } }));
+                    });
+                    aiContainer.appendChild(collocationSentencesContainer);
+                }
+            }
             getCollocations(elements.labels[i.id]);
             switchDiagramView(diagramKeys.main);
             document.dispatchEvent(new CustomEvent('graph-update', { detail: elements.labels[i.id] }));