@sanity/embeddings-index-ui 1.0.0

package/src/embeddingsIndexDashboard/IndexList.tsx

@@ -0,0 +1,83 @@
+import {Box, Button, Card, Flex, Label, Stack, Text} from '@sanity/ui'
+import {IndexState} from '../api/embeddingsApi'
+import {useCallback} from 'react'
+
+export interface IndexListProps {
+  loading: boolean
+  selectedIndex?: IndexState
+  indexes: IndexState[]
+  onIndexSelected: (index: IndexState) => void
+}
+
+export function IndexList(props: IndexListProps) {
+  const {loading, selectedIndex, indexes, onIndexSelected} = props
+  return (
+    <Card tone="default" style={{opacity: loading ? 0.5 : 1}}>
+      <Stack space={2}>
+        <Card borderBottom flex={1} paddingBottom={2}>
+          <Flex>
+            <Box flex={1}>
+              <Label muted>Index name</Label>
+            </Box>
+            <Box flex={1}>
+              <Label muted>Dataset</Label>
+            </Box>
+            <Box flex={1}>
+              <Label muted>Status</Label>
+            </Box>
+            <Box flex={1}>
+              <Label muted>Progress</Label>
+            </Box>
+          </Flex>
+        </Card>
+        {indexes.length ? (
+          <Stack space={2} style={{maxHeight: 200, overflow: 'auto'}}>
+            {indexes.map((index) => (
+              <IndexRow
+                selectedIndex={selectedIndex}
+                index={index}
+                onIndexSelected={onIndexSelected}
+                key={index.indexName}
+              />
+            ))}
+          </Stack>
+        ) : (
+          <Text muted>No indexes found.</Text>
+        )}
+      </Stack>
+    </Card>
+  )
+}
+
+function IndexRow(props: {
+  selectedIndex?: IndexState
+  index: IndexState
+  onIndexSelected: (index: IndexState) => void
+}) {
+  const {selectedIndex, index, onIndexSelected} = props
+  const onSelect = useCallback(() => onIndexSelected(index), [onIndexSelected, index])
+  return (
+    <Button
+      tone={selectedIndex?.indexName === index.indexName ? 'primary' : 'default'}
+      mode={selectedIndex?.indexName === index.indexName ? 'default' : 'bleed'}
+      onClick={onSelect}
+      key={index.indexName}
+    >
+      <Flex>
+        <Box flex={1}>{index.indexName}</Box>
+        <Box flex={1}>{index.dataset}</Box>
+        <Box flex={1}>{index.status}</Box>
+        <Box flex={1}>
+          {index.startDocumentCount
+            ? Math.floor(
+                ((index.startDocumentCount - index.remainingDocumentCount) /
+                  index.startDocumentCount) *
+                  100,
+              )
+            : '?'}
+          %
+        </Box>
+      </Flex>
+    </Button>
+  )
+}

package/src/embeddingsIndexDashboard/QueryIndex.tsx

@@ -0,0 +1,107 @@
+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[] = []
+
+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}) {
+  return (
+    <Flex gap={4} align="center">
+      <Box flex={1}>
+        <DocumentPreview documentId={props.result.value.documentId} button />
+      </Box>
+      <Box>
+        <Text muted size={1}>
+          {Math.floor(props.result.score * 100)} %
+        </Text>
+      </Box>
+    </Flex>
+  )
+}

package/src/embeddingsIndexDashboard/dashboardPlugin.ts

@@ -0,0 +1,15 @@
+import {definePlugin, Tool} from 'sanity'
+import {EarthGlobeIcon} from '@sanity/icons'
+import {EmbeddingsIndexTool} from './EmbeddingsIndexTool'
+
+export const embeddingsIndexTool: Tool = {
+  name: 'embeddings-index',
+  title: 'Embeddings',
+  icon: EarthGlobeIcon,
+  component: EmbeddingsIndexTool,
+}
+
+export const embeddingsIndexDashboard = definePlugin({
+  name: '@sanity/embeddings-index-dashboard',
+  tools: [embeddingsIndexTool],
+})

package/src/embeddingsIndexDashboard/hooks.ts

@@ -0,0 +1,36 @@
+import {ObjectSchemaType, Schema} from 'sanity'
+import {useMemo} from 'react'
+import {isType} from '../utils/types'
+
+const defaultProjection = '{...}'
+
+// eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
+export function useDefaultIndex(schema: Schema, dataset: string) {
+  const defaultFilter = useMemo(
+    () =>
+      `_type in [${schema
+        .getTypeNames()
+        .map((n) => schema.get(n))
+        .filter((schemaType): schemaType is ObjectSchemaType =>
+          Boolean(schemaType && isType(schemaType, 'document')),
+        )
+        .filter(
+          (documentType) =>
+            !documentType.name.startsWith('sanity.') &&
+            !documentType.name.startsWith('assist.') &&
+            documentType.name !== 'document',
+        )
+        .map((documentType) => `"${documentType.name}"`)
+        .join(',\n')}]`,
+    [schema],
+  )
+
+  return useMemo(
+    () => ({
+      dataset,
+      projection: defaultProjection,
+      filter: defaultFilter,
+    }),
+    [defaultFilter, dataset],
+  )
+}

package/src/index.ts

@@ -0,0 +1,9 @@
+import './schemas/typeDefExtensions'
+import {embeddingsIndexReferenceInput} from './referenceInput/referencePlugin'
+import {embeddingsIndexDashboard} from './embeddingsIndexDashboard/dashboardPlugin'
+
+export {embeddingsIndexReferenceInput}
+
+export {embeddingsIndexDashboard}
+
+export * from './api/embeddingsApi'

package/src/preview/DocumentPreview.tsx

@@ -0,0 +1,135 @@
+import {CSSProperties, useEffect, useMemo, useState} from 'react'
+import {useMemoObservable} from 'react-rx'
+import {
+  DefaultPreview,
+  getPreviewStateObservable,
+  getPreviewValueWithFallback,
+  SanityDefaultPreview,
+  SanityDocument,
+  SchemaType,
+  useClient,
+  useDocumentPreviewStore,
+  useSchema,
+} from 'sanity'
+import {useIntentLink} from 'sanity/router'
+import {Box, Button, ButtonProps, Card} from '@sanity/ui'
+import {ErrorOutlineIcon} from '@sanity/icons'
+
+interface ResultPreviewProps {
+  documentId: string
+  button?: boolean
+  style?: CSSProperties
+}
+
+export function DocumentPreview({
+  documentId,
+  style,
+  ...buttonProps
+}: ResultPreviewProps & ButtonProps) {
+  const client = useClient({apiVersion: '2023-06-06'})
+  const [type, setType] = useState<string | undefined>(undefined)
+  const schema = useSchema()
+  const schemaType = type ? schema.get(type) : undefined
+
+  useEffect(() => {
+    client.fetch<string>(`* [_id==$id]._type`, {id: documentId}).then(setType)
+  }, [documentId, client])
+
+  if (!type) {
+    return (
+      <Card style={{minHeight: '36px'}}>
+        <DefaultPreview
+          withShadow={false}
+          withBorder={false}
+          title={'Loading...'}
+          schemaType={schemaType}
+          isPlaceholder
+        />
+      </Card>
+    )
+  }
+
+  if (!schemaType) {
+    return (
+      <Card>
+        <DefaultPreview
+          withShadow={false}
+          withBorder={false}
+          media={() => <ErrorOutlineIcon />}
+          title={
+            <>
+              Unknown type <code>{type ?? 'N/A'}</code> for {documentId}
+            </>
+          }
+        />
+      </Card>
+    )
+  }
+
+  return (
+    <DocumentPreviewInner
+      documentId={documentId}
+      schemaType={schemaType}
+      style={style}
+      {...buttonProps}
+    />
+  )
+}
+
+function DocumentPreviewInner({
+  documentId,
+  schemaType,
+  style,
+  button,
+}: ResultPreviewProps & {schemaType: SchemaType} & ButtonProps) {
+  const documentPreviewStore = useDocumentPreviewStore()
+
+  // NOTE: this emits sync so can never be null
+  // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+  const {draft, published, isLoading} = useMemoObservable(
+    () => getPreviewStateObservable(documentPreviewStore, schemaType, documentId, ''),
+    [documentId, documentPreviewStore, schemaType],
+  )!
+
+  const sanityDocument = useMemo(() => {
+    return {
+      _id: documentId,
+      _type: schemaType?.name,
+    } as SanityDocument
+  }, [documentId, schemaType?.name])
+
+  const {onClick: onIntentClick, href} = useIntentLink({
+    intent: 'edit',
+    params: {
+      id: documentId,
+      type: schemaType?.name,
+    },
+  })
+
+  const preview = (
+    <SanityDefaultPreview
+      {...getPreviewValueWithFallback({
+        draft,
+        published,
+        value: sanityDocument,
+      })}
+      isPlaceholder={isLoading ?? true}
+      layout="default"
+      icon={schemaType?.icon}
+    />
+  )
+  if (button) {
+    return (
+      <Button
+        as={'a'}
+        href={href}
+        onClick={onIntentClick}
+        mode="ghost"
+        style={{width: '100%', ...style}}
+      >
+        {preview}
+      </Button>
+    )
+  }
+  return <Box style={{width: '100%'}}>{preview}</Box>
+}

package/src/referenceInput/SemanticSearchReferenceInput.tsx

@@ -0,0 +1,227 @@
+import {
+  ObjectInputProps,
+  ReferenceSchemaType,
+  set,
+  setIfMissing,
+  typed,
+  unset,
+  useClient,
+} from 'sanity'
+import {Autocomplete, Box, Button, Flex, Text, AutocompleteOpenButtonProps} from '@sanity/ui'
+import {EarthGlobeIcon, LinkIcon} from '@sanity/icons'
+import {useCallback, useEffect, useId, useMemo, useRef, useState} from 'react'
+import {DocumentPreview} from '../preview/DocumentPreview'
+import {useDocumentPane} from 'sanity/desk'
+import {queryIndex, QueryResult} from '../api/embeddingsApi'
+import {publicId} from '../utils/id'
+
+interface Option {
+  result: QueryResult
+  value: string
+}
+
+const NO_OPTIONS: Option[] = []
+const NO_FILTER = () => true
+
+export function SemanticSearchReferenceInput(props: ObjectInputProps) {
+  const defaultEnabled =
+    (props.schemaType as ReferenceSchemaType)?.options?.embeddingsIndex?.searchMode === 'embeddings'
+
+  const [semantic, setSemantic] = useState<boolean>(defaultEnabled)
+  const toggleSemantic = useCallback(() => setSemantic((current) => !current), [])
+
+  return (
+    <Flex gap={2} flex={1} style={{width: '100%'}}>
+      <Box flex={1} style={{maxHeight: 36, overflow: 'hidden'}}>
+        {semantic ? <SemanticSearchInput {...props} /> : props.renderDefault(props)}
+      </Box>
+      <Button
+        icon={semantic ? EarthGlobeIcon : LinkIcon}
+        onClick={toggleSemantic}
+        mode="bleed"
+        title={
+          semantic ? 'Switch to standard reference search' : 'Switch to semantic reference search'
+        }
+      />
+    </Flex>
+  )
+}
+
+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
+}
+
+function useApiClient() {
+  const client = useClient({apiVersion: 'vX'})
+  return useMemo(() => client, [client])
+}
+
+function SemanticSearchInput(props: ObjectInputProps) {
+  const {onPathFocus, onChange, readOnly, schemaType, value} = props
+
+  const {value: currentDocument} = useDocumentPane()
+  const docRef = useRef(currentDocument)
+  const autocompleteRef = useRef<HTMLInputElement>(null)
+
+  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(NO_OPTIONS)
+
+  const client = useApiClient()
+
+  useEffect(() => {
+    docRef.current = currentDocument
+  }, [currentDocument])
+
+  useEffect(() => {
+    // if this component is rendered, and there is a value, replace was selected
+    if (value?._ref) {
+      autocompleteRef.current?.focus()
+    }
+    // intentional empty deps
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [])
+
+  const handleFocus = useCallback(() => onPathFocus(['_ref']), [onPathFocus])
+  const handleBlur = useCallback(() => onPathFocus([]), [onPathFocus])
+
+  const runIndexQuery = useCallback(
+    (queryString: string) => {
+      setSearching(true)
+      const refSchema = schemaType as ReferenceSchemaType
+      const indexName = refSchema.options?.embeddingsIndex?.indexName
+      const maxResults = refSchema.options?.embeddingsIndex?.maxResults
+      const typeFilter = refSchema.to.map((ref) => ref.name)
+
+      if (!indexName) {
+        throw new Error(
+          `Reference option embeddingsIndex.indexName is required, but was missing in type ${refSchema.name}`,
+        )
+      }
+
+      queryIndex(
+        {
+          query: queryString,
+          indexName,
+          maxResults,
+          filter: {
+            type: typeFilter,
+          },
+        },
+        client,
+      )
+        .then((result: QueryResult[]) => {
+          if (queryRef.current === queryString) {
+            setSearching(false)
+            setOptions(
+              result
+                .filter((r) => r.value.documentId !== publicId(docRef.current._id))
+                .map((r) => typed<Option>({result: r, value: r.value.documentId})),
+            )
+          }
+        })
+        .catch((e) => {
+          if (queryRef.current === queryString) {
+            setSearching(false)
+          }
+          throw e
+        })
+    },
+    [client, schemaType],
+  )
+
+  useEffect(() => {
+    if (prevDebouncedQuery.current !== debouncedQuery) {
+      runIndexQuery(debouncedQuery)
+    }
+    prevDebouncedQuery.current = debouncedQuery
+  }, [debouncedQuery, runIndexQuery])
+
+  const handleChange = useCallback(
+    (nextId: string) => {
+      if (!nextId) {
+        onChange(unset())
+        onPathFocus([])
+        return
+      }
+
+      const patches = [
+        setIfMissing({}),
+        set(schemaType.name, ['_type']),
+        set(publicId(nextId), ['_ref']),
+        unset(['_weak']),
+        unset(['_strengthenOnPublish']),
+      ]
+
+      onChange(patches)
+      // Move focus away from _ref and one level up
+      onPathFocus([])
+    },
+    [onChange, onPathFocus, schemaType.name],
+  )
+
+  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],
+  )
+
+  return (
+    <Autocomplete
+      id={id}
+      ref={autocompleteRef}
+      data-testid="semantic-autocomplete"
+      placeholder="Type to search"
+      openButton={openButtonConfig}
+      onFocus={handleFocus}
+      onChange={handleChange}
+      loading={searching}
+      onBlur={handleBlur}
+      readOnly={readOnly}
+      filterOption={NO_FILTER}
+      onQueryChange={handleQueryChange}
+      options={options}
+      renderOption={AutocompleteOption}
+    />
+  )
+}
+
+function AutocompleteOption(props: Option) {
+  return (
+    <Button mode="bleed" padding={1} style={{width: '100%'}}>
+      <Flex gap={2} align="center">
+        <Box flex={1}>
+          <DocumentPreview documentId={props.result.value.documentId} />
+        </Box>
+        <Box padding={2}>
+          <Text size={1} muted title={'Relevance'}>
+            {Math.floor(props.result.score * 100)}%
+          </Text>
+        </Box>
+      </Flex>
+    </Button>
+  )
+}

package/src/referenceInput/referencePlugin.tsx

@@ -0,0 +1,21 @@
+import {definePlugin, isObjectInputProps, ObjectInputProps, ReferenceSchemaType} from 'sanity'
+import {SemanticSearchReferenceInput} from './SemanticSearchReferenceInput'
+import {isType} from '../utils/types'
+
+export const embeddingsIndexReferenceInput = definePlugin({
+  name: '@sanity/embeddings-index-reference-input',
+  form: {
+    components: {
+      input: (props) => {
+        if (
+          isObjectInputProps(props) &&
+          isType(props.schemaType, 'reference') &&
+          (props.schemaType as ReferenceSchemaType).options?.embeddingsIndex?.indexName
+        ) {
+          return <SemanticSearchReferenceInput {...(props as ObjectInputProps)} />
+        }
+        return props.renderDefault(props)
+      },
+    },
+  },
+})

package/src/schemas/typeDefExtensions.ts

@@ -0,0 +1,19 @@
+import 'sanity'
+/* eslint-disable no-unused-vars */
+declare module 'sanity' {
+  interface ReferenceBaseOptions {
+    embeddingsIndex?: {
+      /**
+       * Name of the index
+       */
+      indexName: string
+      maxResults?: number
+      /**
+       * Determines if which search mode is enabled by default for the reference field.
+       * Default is the studio default search, while 'embeddings' enables
+       * Defaults to 'default' behaviour
+       */
+      searchMode?: 'embeddings' | 'default'
+    }
+  }
+}

package/src/utils/id.ts

@@ -0,0 +1,3 @@
+export function publicId(id: string): string {
+  return id.replace('drafts.', '')
+}

package/src/utils/types.ts

@@ -0,0 +1,11 @@
+import {SchemaType} from 'sanity'
+
+export function isType(schemaType: SchemaType, typeName: string): boolean {
+  if (schemaType.name === typeName) {
+    return true
+  }
+  if (!schemaType.type) {
+    return false
+  }
+  return isType(schemaType.type, typeName)
+}

package/v2-incompatible.js

@@ -0,0 +1,11 @@
+const {showIncompatiblePluginDialog} = require('@sanity/incompatible-plugin')
+const {name, version, sanityExchangeUrl} = require('./package.json')
+
+export default showIncompatiblePluginDialog({
+  name: name,
+  versions: {
+    v3: version,
+    v2: undefined,
+  },
+  sanityExchangeUrl,
+})