Support search by entity ID (#679)

* support enter id in the search box

Signed-off-by: shanghaikid <jiangruiyi@gmail.com>

* WIP

Signed-off-by: shanghaikid <jiangruiyi@gmail.com>

* finish

Signed-off-by: shanghaikid <jiangruiyi@gmail.com>

---------

Signed-off-by: shanghaikid <jiangruiyi@gmail.com>

client/src/i18n/cn/search.ts

@@ -28,6 +28,7 @@ const searchTrans = {
   outputFields: '输出字段',
   consistency: '一致性',
   graphNodeHoverTip: '双击以查看更多',
+  inputVectorTip: '向量或实体ID',
 };
 
 export default searchTrans;

client/src/i18n/en/search.ts

@@ -28,6 +28,7 @@ const searchTrans = {
   outputFields: 'Outputs',
   consistency: 'Consistency',
   graphNodeHoverTip: 'Double click to explore more',
+  inputVectorTip: 'Vector or entity id',
 };
 
 export default searchTrans;

client/src/pages/databases/collections/search/Search.tsx

@@ -428,6 +428,7 @@ const Search = (props: CollectionDataProps) => {
                     <VectorInputBox
                       searchParams={s}
                       onChange={onVectorInputChange}
+                      collection={collection}
                     />
 
                     <Typography className="text">

client/src/pages/databases/collections/search/VectorInputBox.tsx

@@ -1,12 +1,14 @@
 import { useRef, useEffect, useState } from 'react';
+import { useTranslation } from 'react-i18next';
 import { EditorState, Compartment } from '@codemirror/state';
-import { EditorView, keymap } from '@codemirror/view';
+import { EditorView, keymap, placeholder } from '@codemirror/view';
 import { insertTab } from '@codemirror/commands';
 import { indentUnit } from '@codemirror/language';
 import { minimalSetup } from 'codemirror';
 import { javascript } from '@codemirror/lang-javascript';
 import { linter, Diagnostic } from '@codemirror/lint';
-import { FieldObject } from '@server/types';
+import { CollectionFullObject, FieldObject } from '@server/types';
+import { CollectionService } from '@/http';
 import { DataTypeStringEnum } from '@/consts';
 import { SearchSingleParams } from '../../types';
 import { isSparseVector, transformObjStrToJSONStr } from '@/utils';
@@ -106,13 +108,17 @@ const Validator = {
 export type VectorInputBoxProps = {
   onChange: (anns_field: string, value: string) => void;
   searchParams: SearchSingleParams;
+  collection: CollectionFullObject;
 };
 
+let queryTimeout: NodeJS.Timeout;
+
 export default function VectorInputBox(props: VectorInputBoxProps) {
   const theme = useTheme();
+  const { t: searchTrans } = useTranslation('search');
 
   // props
-  const { searchParams, onChange } = props;
+  const { searchParams, onChange, collection } = props;
   const { field, data } = searchParams;
 
   // UI states
@@ -127,6 +133,7 @@ export default function VectorInputBox(props: VectorInputBoxProps) {
   const onChangeRef = useRef(onChange);
   const dataRef = useRef(data);
   const fieldRef = useRef(field);
+  const searchParamsRef = useRef(searchParams);
 
   const themeCompartment = new Compartment();
 
@@ -138,6 +145,7 @@ export default function VectorInputBox(props: VectorInputBoxProps) {
     dataRef.current = data;
     onChangeRef.current = onChange;
     fieldRef.current = field;
+    searchParamsRef.current = searchParams;
 
     if (editor.current) {
       // only data replace should trigger this, otherwise, let cm handle the state
@@ -145,13 +153,52 @@ export default function VectorInputBox(props: VectorInputBoxProps) {
         editor.current.dispatch({
           changes: {
             from: 0,
-            to: editor.current.state.doc.length,
+            to: data.length + 1,
             insert: data,
           },
         });
       }
     }
-  }, [JSON.stringify(searchParams)]);
+  }, [JSON.stringify(searchParams), onChange]);
+
+  const getVectorById = (text: string) => {
+    if (queryTimeout) {
+      clearTimeout(queryTimeout);
+    }
+    // only search for text that doesn't have space, comma, or brackets or curly brackets
+    if (!text.trim().match(/[\s,{}]/)) {
+      const isVarChar =
+        collection.schema.primaryField.data_type === DataTypeStringEnum.VarChar;
+
+      if (!isVarChar && isNaN(Number(text))) {
+        return;
+      }
+
+      queryTimeout = setTimeout(() => {
+        try {
+          CollectionService.queryData(collection.collection_name, {
+            expr: isVarChar
+              ? `${collection.schema.primaryField.name} == '${text}'`
+              : `${collection.schema.primaryField.name} == ${text}`,
+            output_fields: [searchParamsRef.current.anns_field],
+          })
+            .then(res => {
+              if (res.data && res.data.length === 1) {
+                onChangeRef.current(
+                  searchParamsRef.current.anns_field,
+                  JSON.stringify(
+                    res.data[0][searchParamsRef.current.anns_field]
+                  )
+                );
+              }
+            })
+            .catch(e => {
+              console.log(0, e);
+            });
+        } catch (e) {}
+      }, 300);
+    }
+  };
 
   // create editor
   useEffect(() => {
@@ -161,6 +208,7 @@ export default function VectorInputBox(props: VectorInputBoxProps) {
         extensions: [
           minimalSetup,
           javascript(),
+          placeholder('vector or entity id'),
           linter(view => {
             const text = view.state.doc.toString();
 
@@ -214,10 +262,16 @@ export default function VectorInputBox(props: VectorInputBoxProps) {
           EditorView.lineWrapping,
           EditorView.updateListener.of(update => {
             if (update.docChanged) {
+              if (queryTimeout) {
+                clearTimeout(queryTimeout);
+              }
               const text = update.state.doc.toString();
+
               const { valid } = validator(text, fieldRef.current);
               if (valid || text === '') {
                 onChangeRef.current(searchParams.anns_field, text);
+              } else {
+                getVectorById(text);
               }
             }
             if (update.focusChanged) {
@@ -234,19 +288,28 @@ export default function VectorInputBox(props: VectorInputBoxProps) {
 
       editor.current = view;
 
+      // focus editor, the cursor will be at the end of the text
+      const endPos = editor.current.state.doc.length;
+      editor.current.dispatch({
+        selection: { anchor: endPos },
+      });
+      editor.current.focus(); // 聚焦到编辑器
+
       return () => {
         view.destroy();
         editor.current = undefined;
       };
     }
-  }, [JSON.stringify(field)]);
+  }, [JSON.stringify(field), getVectorById]);
 
   useEffect(() => {
     // dispatch dark mode change to editor
     if (editor.current) {
       editor.current.dispatch({
         effects: themeCompartment.reconfigure(
-          themeCompartment.of(theme.palette.mode === 'light' ? githubLight : githubDark)
+          themeCompartment.of(
+            theme.palette.mode === 'light' ? githubLight : githubDark
+          )
         ),
       });
     }
@@ -256,9 +319,6 @@ export default function VectorInputBox(props: VectorInputBoxProps) {
     <div
       className={`${classes.vectorInputBox} ${isFocused ? 'focused' : ''}`}
       ref={editorEl}
-      onClick={() => {
-        if (editor.current) editor.current.focus();
-      }}
     ></div>
   );
 }