@sanity/embeddings-index-ui 1.0.2 → 1.1.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@sanity/embeddings-index-ui",
-  "version": "1.0.2",
+  "version": "1.1.0",
   "description": "Various Sanity Studio plugins for integrating with the embeddings index API",
   "keywords": [
     "sanity",

package/src/api/isEnabled.tsx

@@ -0,0 +1,53 @@
+import {useClient} from 'sanity'
+import {createContext, PropsWithChildren, useContext, useEffect, useState} from 'react'
+import {Card, Text} from '@sanity/ui'
+
+const featureName = 'embeddingsIndexApi'
+
+export type FeatureStatus = 'enabled' | 'disabled' | 'loading'
+export const FeatureEnabledContext = createContext<FeatureStatus>('loading')
+
+export function useIsFeatureEnabled() {
+  const client = useClient({apiVersion: '2023-09-01'})
+  const [status, setStatus] = useState<FeatureStatus>('loading')
+
+  useEffect(() => {
+    client
+      .request<string | boolean>({
+        method: 'GET',
+        url: `/projects/${client.config().projectId}/features/${featureName}`,
+      })
+      .then((isEnabled) => {
+        setStatus(isEnabled === 'true' || isEnabled === true ? 'enabled' : 'disabled')
+      })
+      .catch((err) => {
+        console.error(err)
+        setStatus('disabled')
+      })
+  }, [client])
+
+  return status
+}
+
+export function FeatureEnabledProvider(props: PropsWithChildren<{}>) {
+  const status = useIsFeatureEnabled()
+  return (
+    <FeatureEnabledContext.Provider value={status}>{props.children}</FeatureEnabledContext.Provider>
+  )
+}
+
+export function useIsFeatureEnabledContext(): FeatureStatus {
+  return useContext(FeatureEnabledContext)
+}
+
+export function FeatureDisabledNotice() {
+  return (
+    <Card tone="primary" border padding={2}>
+      <Text size={1}>
+        Embeddings index APIs are only available on the{' '}
+        <a href="https://sanity.io/pricing">Team tier and above</a>. Please upgrade to enable
+        access.
+      </Text>
+    </Card>
+  )
+}

package/src/embeddingsIndexDashboard/EmbeddingsIndexTool.tsx

@@ -6,13 +6,23 @@ import {EditIndexDialog} from './IndexEditor'
 import {IndexList} from './IndexList'
 import {IndexInfo} from './IndexInfo'
 import {useApiClient} from '../api/embeddingsApiHooks'
+import {FeatureDisabledNotice, useIsFeatureEnabled} from '../api/isEnabled'
 
 export function EmbeddingsIndexTool() {
+  const featureState = useIsFeatureEnabled()
   return (
     <Card>
       <Flex justify="center" flex={1}>
         <Card flex={1} style={{maxWidth: 1200}} padding={5}>
-          <Indexes />
+          {featureState == 'loading' ? (
+            <Box padding={2}>
+              <Spinner />
+            </Box>
+          ) : null}
+
+          {featureState == 'disabled' ? <FeatureDisabledNotice /> : null}
+
+          {featureState == 'enabled' ? <Indexes /> : null}
         </Card>
       </Flex>
     </Card>

package/src/embeddingsIndexDashboard/IndexEditor.tsx

@@ -141,19 +141,21 @@ export function IndexEditor(props: {
         />
         <IndexFormInput label="Dataset" index={index} prop="dataset" onChange={setIndex} readOnly />
         <IndexFormInput
-          label="Projection"
-          placeholder={defaultIndex.projection}
+          label="Filter"
+          description="Must be a valid GROQ filter"
+          placeholder={defaultIndex.filter}
           index={index}
-          prop="projection"
+          prop="filter"
           onChange={setIndex}
           readOnly={readOnly}
           type="textarea"
         />
         <IndexFormInput
-          label="Filter"
-          placeholder={defaultIndex.filter}
+          label="Projection"
+          description="Must be a valid GROQ projection, starting { and ending with }"
+          placeholder={defaultIndex.projection}
           index={index}
-          prop="filter"
+          prop="projection"
           onChange={setIndex}
           readOnly={readOnly}
           type="textarea"

package/src/embeddingsIndexDashboard/IndexFormInput.tsx

@@ -1,11 +1,12 @@
 import {NamedIndex} from '../api/embeddingsApi'
 import {Dispatch, FormEvent, SetStateAction, useCallback, useId} from 'react'
-import {Label, Stack, TextArea, TextInput} from '@sanity/ui'
+import {Box, Label, Stack, TextArea, TextInput, Text} from '@sanity/ui'
 
 export interface IndexFormInputProps {
   index: Partial<NamedIndex>
   prop: keyof NamedIndex
   label: string
+  description?: string
   onChange: Dispatch<SetStateAction<Partial<NamedIndex>>>
   readOnly: boolean
   placeholder?: string
@@ -13,7 +14,7 @@ export interface IndexFormInputProps {
 }
 
 export function IndexFormInput(props: IndexFormInputProps) {
-  const {label, index, prop, onChange, readOnly, placeholder, type} = props
+  const {label, description, index, prop, onChange, readOnly, placeholder, type} = props
   const handleChange = useCallback(
     (propValue: string) => onChange((current) => ({...current, [prop]: propValue})),
     [onChange, prop],
@@ -21,6 +22,7 @@ export function IndexFormInput(props: IndexFormInputProps) {
   return (
     <FormInput
       label={label}
+      description={description}
       onChange={handleChange}
       value={index[prop] ?? ''}
       readOnly={readOnly}
@@ -32,6 +34,7 @@ export function IndexFormInput(props: IndexFormInputProps) {
 
 interface FormInputProps {
   label: string
+  description?: string
   onChange: (value: string) => void
   value: string
   readOnly: boolean
@@ -40,7 +43,7 @@ interface FormInputProps {
 }
 
 function FormInput(props: FormInputProps) {
-  const {label, onChange, value, readOnly, placeholder, type = 'text'} = props
+  const {label, description, onChange, value, readOnly, placeholder, type = 'text'} = props
   const id = useId()
   const handleChange = useCallback(
     (e: FormEvent<HTMLInputElement | HTMLTextAreaElement>) => onChange(e.currentTarget.value),
@@ -51,6 +54,13 @@ function FormInput(props: FormInputProps) {
       <Label muted htmlFor={id}>
         <label htmlFor={id}>{label}</label>
       </Label>
+      {description && (
+        <Box>
+          <Text size={1} muted>
+            {description}
+          </Text>
+        </Box>
+      )}
       {type === 'text' ? (
         <TextInput
           id={id}

package/src/embeddingsIndexDashboard/IndexList.tsx

@@ -14,7 +14,7 @@ export function IndexList(props: IndexListProps) {
   return (
     <Card tone="default" style={{opacity: loading ? 0.5 : 1}}>
       <Stack space={2}>
-        <Card borderBottom flex={1} paddingBottom={2}>
+        <Card borderBottom flex={1} paddingBottom={2} padding={3}>
           <Flex>
             <Box flex={1}>
               <Label muted>Index name</Label>
@@ -59,12 +59,15 @@ function IndexRow(props: {
   return (
     <Button
       tone={selectedIndex?.indexName === index.indexName ? 'primary' : 'default'}
-      mode={selectedIndex?.indexName === index.indexName ? 'default' : 'bleed'}
+      mode={selectedIndex?.indexName === index.indexName ? 'default' : 'ghost'}
       onClick={onSelect}
       key={index.indexName}
+      padding={3}
     >
       <Flex>
-        <Box flex={1}>{index.indexName}</Box>
+        <Box flex={1}>
+          <strong>{index.indexName}</strong>
+        </Box>
         <Box flex={1}>{index.dataset}</Box>
         <Box flex={1}>{index.status}</Box>
         <Box flex={1}>

package/src/embeddingsIndexDashboard/QueryIndex.tsx

@@ -1,108 +1,13 @@
-import {SearchIcon} from '@sanity/icons'
-import {Box, Card, Flex, Spinner, Stack, Text, TextInput} from '@sanity/ui'
-import {FormEvent, useCallback, useState, KeyboardEvent} from 'react'
-import {DocumentPreview} from '../preview/DocumentPreview'
-import {queryIndex, QueryResult} from '../api/embeddingsApi'
-import {useApiClient} from '../api/embeddingsApiHooks'
-
-const NO_RESULTS: QueryResult[] = []
+import {useCallback, useMemo} from 'react'
+import {SemanticSearchAutocomplete} from '../referenceInput/SemanticSearchAutocomplete'
+import {EmbeddingsIndexConfig} from '../schemas/typeDefExtensions'
 
 export function QueryIndex(props: {indexName: string}) {
   const {indexName} = props
-  const [query, setQuery] = useState('')
-  const [searching, setSearching] = useState(false)
-  const [results, setResults] = useState(NO_RESULTS)
-
-  const client = useApiClient()
-
-  const search = useCallback(
-    (queryString: string) => {
-      setSearching(true)
-      return queryIndex(
-        {
-          query: queryString,
-          indexName,
-          maxResults: 5,
-        },
-        client,
-      )
-        .then(setResults)
-        .finally(() => setSearching(false))
-    },
-    [client, indexName],
-  )
-
-  const onInputChange = useCallback((e: FormEvent<HTMLInputElement>) => {
-    setQuery(e.currentTarget.value)
-  }, [])
-
-  const onKeyDown = useCallback(
-    (e: KeyboardEvent<HTMLInputElement>) => {
-      if (e.key === 'Enter') {
-        search(query).catch(console.error)
-      }
-    },
-    [search, query],
-  )
-
-  return (
-    <Stack space={3} flex={1}>
-      <Flex flex={1}>
-        <Card flex={1}>
-          <TextInput
-            iconRight={
-              searching ? (
-                <Box style={{marginTop: 5}}>
-                  <Spinner />
-                </Box>
-              ) : (
-                SearchIcon
-              )
-            }
-            placeholder={'Find documents'}
-            value={query}
-            disabled={searching}
-            onChange={onInputChange}
-            onKeyDown={onKeyDown}
-          />
-        </Card>
-      </Flex>
-      <Flex gap={4} style={{opacity: searching ? 0.5 : 1}}>
-        <Box flex={1}>
-          <ResultList results={results} query={query} />
-        </Box>
-      </Flex>
-    </Stack>
-  )
-}
-
-export function ResultList(props: {results: QueryResult[]; query: string}) {
-  const {results, query} = props
-
-  return (
-    <Stack space={4} height="fill">
-      <Stack space={2}>
-        {results.map((r) => (
-          <ResultEntry result={r} key={r.value.documentId} />
-        ))}
-        {!results.length && query ? 'No results.' : null}
-      </Stack>
-    </Stack>
-  )
-}
-
-function ResultEntry(props: {result: QueryResult}) {
-  const value = props.result.value
-  return (
-    <Flex gap={4} align="center">
-      <Box flex={1}>
-        <DocumentPreview documentId={value.documentId} schemaTypeName={value.type} button />
-      </Box>
-      <Box>
-        <Text muted size={1}>
-          {Math.floor(props.result.score * 100)} %
-        </Text>
-      </Box>
-    </Flex>
+  const getEmpty = useCallback(() => 'anything', [])
+  const indexConfig: EmbeddingsIndexConfig = useMemo(
+    () => ({indexName, maxResults: 8}),
+    [indexName],
   )
+  return <SemanticSearchAutocomplete getEmptySearchValue={getEmpty} indexConfig={indexConfig} />
 }

package/src/referenceInput/SemanticSearchAutocomplete.tsx

@@ -0,0 +1,205 @@
+/* eslint-disable max-nested-callbacks */
+import {Autocomplete, AutocompleteOpenButtonProps, Box, Button, Flex, Text} from '@sanity/ui'
+import {DocumentPreview} from '../preview/DocumentPreview'
+import {
+  FocusEventHandler,
+  forwardRef,
+  useCallback,
+  useEffect,
+  useId,
+  useMemo,
+  useRef,
+  useState,
+} from 'react'
+import {EmbeddingsIndexConfig} from '../schemas/typeDefExtensions'
+import {queryIndex, QueryResult} from '../api/embeddingsApi'
+import {typed} from 'sanity'
+import {useApiClient} from '../api/embeddingsApiHooks'
+
+export interface SemanticSearchAutocompleteProps {
+  indexConfig: EmbeddingsIndexConfig
+  getEmptySearchValue: () => string
+  typeFilter?: string[]
+  filterResult?: (hit: QueryResult) => boolean
+  onChange?: (value: string) => void
+  onFocus?: FocusEventHandler<HTMLInputElement>
+  onBlur?: FocusEventHandler<HTMLInputElement>
+  readOnly?: boolean
+}
+
+interface Option {
+  result: QueryResult
+  value: string
+}
+
+const NO_RESULTS_VALUE = '' as const
+
+interface NoResultOption {
+  value: string
+}
+
+const NO_OPTIONS: NoResultOption[] = []
+const NO_FILTER = () => true
+
+export const SemanticSearchAutocomplete = forwardRef(function SemanticSearchAutocomplete(
+  props: SemanticSearchAutocompleteProps,
+  ref: any,
+) {
+  const {
+    indexConfig,
+    filterResult,
+    getEmptySearchValue,
+    readOnly,
+    onFocus,
+    onBlur,
+    onChange,
+    typeFilter,
+  } = props
+  const id = useId()
+  const [query, setQuery] = useState('')
+  const queryRef = useRef(query)
+  const debouncedQuery = useDebouncedValue(query, 300)
+  const prevDebouncedQuery = useRef(debouncedQuery)
+
+  const [searching, setSearching] = useState(false)
+  const [options, setOptions] = useState<Option[] | NoResultOption[]>(NO_OPTIONS)
+
+  const client = useApiClient()
+
+  const runIndexQuery = useCallback(
+    (queryString: string) => {
+      setSearching(true)
+      const indexName = indexConfig?.indexName
+      const maxResults = indexConfig?.maxResults
+
+      if (!indexName) {
+        throw new Error(`Reference option embeddingsIndex.indexName is required, but was missing`)
+      }
+
+      queryIndex(
+        {
+          query: queryString.trim().length ? queryString : getEmptySearchValue() ?? '',
+          indexName,
+          maxResults,
+          filter: {
+            type: typeFilter,
+          },
+        },
+        client,
+      )
+        .then((result: QueryResult[]) => {
+          if (queryRef.current === queryString) {
+            setSearching(false)
+            setOptions([])
+            const resultOptions = result
+              .filter((hit) => (filterResult ? filterResult(hit) : true))
+              .map((r) => typed<Option>({result: r, value: r.value.documentId}))
+            if (resultOptions.length) {
+              setOptions(resultOptions)
+            } else {
+              setOptions([{value: NO_RESULTS_VALUE}])
+            }
+          }
+        })
+        .catch((e) => {
+          if (queryRef.current === queryString) {
+            setSearching(false)
+          }
+          throw e
+        })
+    },
+    [client, indexConfig, getEmptySearchValue, filterResult, typeFilter],
+  )
+
+  useEffect(() => {
+    if (prevDebouncedQuery.current !== debouncedQuery) {
+      runIndexQuery(debouncedQuery)
+    }
+    prevDebouncedQuery.current = debouncedQuery
+  }, [debouncedQuery, runIndexQuery])
+
+  const openButtonConfig: AutocompleteOpenButtonProps = useMemo(
+    () => ({onClick: () => runIndexQuery(queryRef.current)}),
+    [runIndexQuery, queryRef],
+  )
+
+  const handleQueryChange = useCallback(
+    (newValue: string | null) => {
+      const newQuery = newValue ?? ''
+      queryRef.current = newQuery
+      setQuery(newQuery)
+    },
+    [setQuery],
+  )
+
+  const handleChange = useCallback(
+    (value: string) => {
+      if (value === NO_RESULTS_VALUE) {
+        setOptions(NO_OPTIONS)
+        return
+      }
+      onChange?.(value)
+    },
+    [onChange],
+  )
+
+  return (
+    <Autocomplete
+      id={id}
+      ref={ref}
+      data-testid="semantic-autocomplete"
+      placeholder="Type to search..."
+      openButton={openButtonConfig}
+      onFocus={onFocus}
+      onChange={handleChange}
+      loading={searching}
+      onBlur={onBlur}
+      readOnly={readOnly}
+      filterOption={NO_FILTER}
+      onQueryChange={handleQueryChange}
+      options={options}
+      renderOption={AutocompleteOption}
+    />
+  )
+})
+
+function AutocompleteOption(props: Option | NoResultOption) {
+  if ('result' in props) {
+    const value = props.result.value
+    return (
+      <Button mode="bleed" padding={1} style={{width: '100%'}}>
+        <Flex gap={2} align="center">
+          <Box flex={1}>
+            <DocumentPreview documentId={value.documentId} schemaTypeName={value.type} />
+          </Box>
+          <Box padding={2}>
+            <Text size={1} muted title={'Relevance'}>
+              {Math.floor(props.result.score * 100)}%
+            </Text>
+          </Box>
+        </Flex>
+      </Button>
+    )
+  }
+
+  return (
+    <Button mode="bleed" padding={1} style={{width: '100%'}} disabled>
+      <Flex gap={2} align="center">
+        No results.
+      </Flex>
+    </Button>
+  )
+}
+
+function useDebouncedValue<T>(value: T, ms: number) {
+  const [debouncedValue, setDebouncedValue] = useState(value)
+
+  useEffect(() => {
+    const timeoutId = setTimeout(() => {
+      setDebouncedValue(value)
+    }, ms)
+    return () => clearTimeout(timeoutId)
+  }, [value, ms])
+
+  return debouncedValue
+}