ui: support enable index

pingcap · Apr 16, 2024 · ab7bfe6 · ab7bfe6
1 parent f68726a
commit ab7bfe6
Show file tree

Hide file tree

Showing 13 changed files with 137 additions and 39 deletions.
diff --git a/src/app/(main)/(admin)/indexes/[id]/page.tsx b/src/app/(main)/(admin)/indexes/[id]/page.tsx
@@ -4,11 +4,18 @@ import { useIndex } from '@/app/(main)/(admin)/indexes/[id]/context';
 import { SummaryStatsBar } from '@/app/(main)/(admin)/indexes/[id]/SummaryStatsBar';
 import { Form, FormControl, FormDescription, FormItem, FormLabel } from '@/components/ui/form';
 import { Switch } from '@/components/ui/switch';
+import { useOperation } from '@/components/use-operation';
+import { handleErrors } from '@/lib/fetch';
+import { withToast } from '@/lib/toast';
+import { useRouter } from 'next/navigation';
 import { useForm } from 'react-hook-form';
 
 export default function Page () {
+  const router = useRouter();
   const index = useIndex();
   const form = useForm();
+  const { operating: enabling, operate: enable } = useOperation(enableIndex);
+
   return (
     <>
       <article className="prose prose-sm prose-neutral dark:prose-invert">
@@ -19,22 +26,33 @@ export default function Page () {
           This is configuration for document index <code>{index.name}</code>(using {index.config.provider}).
         </p>
         <p>
-          You need to setup all configurations in tabs above. Once your configuration is completed,
-          turn on the switch below. You <b>can&#39;t</b> change configurations after index is marked `configured`.
+          You need to setup all configurations in tabs above. Once enable your index (turn on the switch below),
+          you <b>can&#39;t</b> change your index configurations.
         </p>
         <div className="mt-8">
           <Form {...form}>
             <FormItem className="flex flex-row items-center justify-between rounded-lg border p-4">
               <div className="space-y-0.5">
                 <FormLabel className="text-base">
-                  Configured
+                  Enabled
                 </FormLabel>
                 <FormDescription className="text-yellow-500">
                   You can&#39;t revert this operation
                 </FormDescription>
               </div>
               <FormControl>
-                <Switch checked={!!index.configured} disabled={!!index.configured} />
+                <Switch
+                  checked={!!index.configured || enabling}
+                  disabled={!!index.configured || enabling}
+                  onCheckedChange={checked => {
+                    if (checked) {
+                      enable(index.name)
+                        .then(() => {
+                          router.refresh();
+                        });
+                    }
+                  }}
+                />
               </FormControl>
             </FormItem>
           </Form>
@@ -43,9 +61,16 @@ export default function Page () {
           <h2>
             Index Stats
           </h2>
-          <SummaryStatsBar index={index} className='my-2' />
+          <SummaryStatsBar index={index} className="my-2" />
         </>}
       </article>
     </>
   );
-}
+}
+
+const enableIndex = withToast((name: string) =>
+  fetch(`/api/v1/indexes/${name}/enable`, {
+    method: 'POST',
+  })
+    .then(handleErrors),
+);
diff --git a/src/app/api/v1/indexes/[name]/config/[key]/route.ts b/src/app/api/v1/indexes/[name]/config/[key]/route.ts
@@ -6,7 +6,7 @@ import z from 'zod';
 export const GET = defineHandler({
   params: z.object({
     name: z.string(),
-    key: z.enum(['parser', 'metadata_extractors', 'llm', 'embedding']),
+    key: z.enum(['parser', 'reader', 'metadata_extractors', 'llm', 'embedding']),
   }),
 }, async ({
   params: {
@@ -23,22 +23,21 @@ export const GET = defineHandler({
 
 export const PUT = defineHandler({
   params: z.object({
-    id: z.coerce.number().int(),
-    key: z.enum(['parser', 'metadata_extractors', 'llm']),
+    name: z.string(),
+    key: z.enum(['parser', 'reader', 'metadata_extractors', 'llm', 'embedding']),
   }),
   body: z.any(),
 }, async ({
-  params: { id, key },
+  params: { name, key },
   body,
 }) => {
-  const index = await getIndex(id);
+  const index = await getIndexByName(name);
   if (!index) {
     notFound();
   }
-
   // TODO: validate config schema
 
-  await updateIndexConfig(id, `.${key}`, body);
+  await updateIndexConfig(index.id, `.${key}`, body);
 });
 
 export const dynamic = 'force-dynamic';
diff --git a/src/app/api/v1/indexes/[name]/config/reader/[provider]/route.ts b/src/app/api/v1/indexes/[name]/config/reader/[provider]/route.ts
@@ -24,15 +24,16 @@ export const GET = defineHandler({ params }, async ({
 
 export const PUT = defineHandler({
   params: z.object({
-    id: z.coerce.number().int(),
+    name: z.string(),
     provider: z.string(),
   }),
   body: z.any(),
 }, async ({
-  params: { id, provider },
+  params: { name, provider },
   body,
 }) => {
-  if (!await getIndex(id)) {
+  const index = await getIndexByName(name);
+  if (!index) {
     notFound();
   }
   const component = await baseRegistry.getComponent(provider);
@@ -43,7 +44,7 @@ export const PUT = defineHandler({
 
   await component.optionsSchema.parse(body);
 
-  await updateIndexConfig(id, `.reader.${provider}`, body);
+  await updateIndexConfig(index.id, `.reader.${provider}`, body);
 });
 
 export const dynamic = 'force-dynamic';
diff --git a/src/app/api/v1/indexes/[name]/enable/route.ts b/src/app/api/v1/indexes/[name]/enable/route.ts
@@ -0,0 +1,19 @@
+import { enableIndex, getIndexByName } from '@/core/repositories/index_';
+import { defineHandler } from '@/lib/next/handler';
+import { notFound } from 'next/navigation';
+import { z } from 'zod';
+
+export const POST = defineHandler({
+  auth: 'admin',
+  params: z.object({
+    name: z.string(),
+  }),
+}, async ({ params }) => {
+  const index = await getIndexByName(params.name);
+  if (!index) {
+    notFound();
+  }
+  await enableIndex(index.id);
+});
+
+export const dynamic = 'force-dynamic';
diff --git a/src/client/operations/index_.ts b/src/client/operations/index_.ts
@@ -17,8 +17,8 @@ export const createIndex = withToast(
   },
 );
 
-export const updateIndexConfigPart = withToast(<K extends keyof IndexConfig> (id: number, key: K, value: IndexConfig[K]) =>
-  fetch(`/api/v1/indexes/${id}/config/${key}`, {
+export const updateIndexConfigPart = withToast(<K extends keyof IndexConfig> (name: string, key: K, value: IndexConfig[K]) =>
+  fetch(`/api/v1/indexes/${name}/config/${key}`, {
     method: 'put',
     body: JSON.stringify(value),
   }).then(handleErrors),

diff --git a/src/components/llamaindex/config/EmbeddingConfigViewer.tsx b/src/components/llamaindex/config/EmbeddingConfigViewer.tsx
@@ -12,6 +12,8 @@ const schema = z.object({
   provider: z.string(),
   config: z.object({
     model: z.string(),
+    vectorColumn: z.literal('embedding').default('embedding'),
+    vectorDimension: z.coerce.number(),
   })
 });
 
@@ -25,8 +27,7 @@ export function EmbeddingConfigViewer ({ index }: { index: Index }) {
   });
 
   const handleSubmit = form.handleSubmit((value) => {
-    // updateEmbedding(value);
-    console.log(value);
+    updateEmbedding(value);
   });
 
   useEffect(() => {
@@ -64,6 +65,19 @@ export function EmbeddingConfigViewer ({ index }: { index: Index }) {
             </FormItem>
           )}
         />
+        <FormField
+          control={form.control}
+          name="config.vectorDimension"
+          render={({ field }) => (
+            <FormItem>
+              <FormLabel>Vector Dimensions</FormLabel>
+              <FormControl>
+                <Input {...field} value={field.value ?? ''} />
+              </FormControl>
+              <FormMessage />
+            </FormItem>
+          )}
+        />
         <Button className="mt-4" disabled={disabled}>Submit</Button>
       </Form>
     </form>

diff --git a/src/components/llamaindex/config/LLMConfigViewer.tsx b/src/components/llamaindex/config/LLMConfigViewer.tsx
@@ -26,8 +26,7 @@ export function LlmConfigViewer ({ index }: { index: Index }) {
   });
 
   const handleSubmit = form.handleSubmit((value) => {
-    // updateLlm(value);
-    console.log(value);
+    updateLlm(value);
   });
 
   useEffect(() => {

diff --git a/src/components/llamaindex/config/ParserConfigViewer.tsx b/src/components/llamaindex/config/ParserConfigViewer.tsx
@@ -29,8 +29,7 @@ export function ParserConfigViewer ({ index }: { index: Index }) {
   });
 
   const handleSubmit = form.handleSubmit((value) => {
-    // updateParser(value);
-    console.log(value);
+    updateParser(value);
   });
 
   useEffect(() => {

diff --git a/src/components/llamaindex/config/ReaderConfigViewer.tsx b/src/components/llamaindex/config/ReaderConfigViewer.tsx
@@ -12,10 +12,12 @@ import { z } from 'zod';
 
 const schema = z.object({
   'rag_$loader_$html': z.object({
-    url: z.string(),
-    excludeSelectors: htmlSelectorArray(),
-    selectors: htmlSelectorArray(),
-  }).array().describe('HTML Loader Configuration'),
+    contentExtraction: z.object({
+      url: z.string(),
+      excludeSelectors: htmlSelectorArray(),
+      selectors: htmlSelectorArray(),
+    }).array().describe('HTML Loader Rules')
+  }).describe('HTML Loader Configuration'),
 });
 
 function toSafeName (record: Record<string, any> | undefined) {
@@ -44,8 +46,10 @@ export function ReaderConfigViewer ({ index }: { index: Index }) {
   });
 
   const handleSubmit = form.handleSubmit((value) => {
-    // updateReader(value);
     value = fromSafeName(value);
+    updateReader(value);
+  }, e => {
+    console.log(e)
   });
 
   useEffect(() => {
@@ -58,7 +62,7 @@ export function ReaderConfigViewer ({ index }: { index: Index }) {
     <form onSubmit={handleSubmit}>
       <Form {...form}>
         <Accordion type="single" defaultValue="rag_$loader_$html.contentExtraction">
-          <AutoFormArray name="rag_$loader_$html.contentExtraction" path={['rag_$loader_$html', 'contentExtraction']} item={schema.shape.rag_$loader_$html} form={form} />
+          <AutoFormArray name="rag_$loader_$html.contentExtraction" path={['rag_$loader_$html', 'contentExtraction']} item={schema.shape.rag_$loader_$html.shape.contentExtraction} form={form} />
         </Accordion>
         <Button className="mt-4" disabled={disabled}>Submit</Button>
       </Form>

diff --git a/src/components/llamaindex/config/use-index-config-part.ts b/src/components/llamaindex/config/use-index-config-part.ts
@@ -5,7 +5,7 @@ import { useCallback, useState } from 'react';
 import useSWR from 'swr';
 
 export function useIndexConfigPart<K extends keyof IndexConfig> (index: Index, part: K) {
-  const { data, isLoading, mutate } = useSWR<IndexConfig[K]>(['get', `/api/v1/indexes/${index.id}/config/${part}`], {
+  const { data, isLoading, mutate } = useSWR<IndexConfig[K]>(['get', `/api/v1/indexes/${index.name}/config/${part}`], {
     fetcher,
     revalidateOnMount: false,
     fallbackData: index.config[part],
@@ -15,7 +15,7 @@ export function useIndexConfigPart<K extends keyof IndexConfig> (index: Index, p
 
   const update = useCallback((value: IndexConfig[K]) => {
     setIsUpdating(true);
-    updateIndexConfigPart(index.id, part, value)
+    updateIndexConfigPart(index.name, part, value)
       .then(() => {
         try {
           index.config[part] = value;

diff --git a/src/components/semantic-search.tsx b/src/components/semantic-search.tsx
@@ -73,7 +73,6 @@ function InternalSearchBox () {
         .then(res => res.json())
         .then((res: SearchResult) => {
           const set = new Set<string>();
-          console.log(set);
           res.relevantChunks = res.relevantChunks.filter(res => {
             if (set.has(res.source_uri)) {
               return false;

diff --git a/src/components/use-operation.ts b/src/components/use-operation.ts
@@ -0,0 +1,42 @@
+import { useEffect, useRef, useState } from 'react';
+
+export function useOperation<T extends any[], R> (operation: (...args: T) => Promise<R>) {
+  const [operating, setOperating] = useState(false);
+  const [result, setResult] = useState<R>();
+  const [error, setError] = useState<unknown>();
+
+  const operationRef = useRef(operation);
+  useEffect(() => {
+    operationRef.current = operation;
+  });
+
+  const operatingRef = useRef(operating);
+  operatingRef.current = operating;
+
+  const operate = (...args: T) => {
+    if (operatingRef.current) {
+      return Promise.reject<R>('Operation is running');
+    }
+    setOperating(true);
+    setResult(undefined);
+    setError(undefined);
+    return operationRef.current(...args)
+      .then(result => {
+        setResult(result);
+        return result;
+      }, error => {
+        setError(error);
+        return Promise.reject(error);
+      })
+      .finally(() => {
+        setOperating(false);
+      });
+  };
+
+  return {
+    operate,
+    operating,
+    error,
+    result,
+  };
+}
diff --git a/src/core/repositories/index_.ts b/src/core/repositories/index_.ts
@@ -102,7 +102,7 @@ export async function enableIndex (id: number) {
   }
 
   await sql`
-      CREATE TABLE IF NOT EXISTS ${index.config.provider}_document_chunk_node_${index.name}
+      CREATE TABLE IF NOT EXISTS ${sql.raw(`${index.config.provider}_document_chunk_node_${index.name}`)}
       (
           -- Text Node Common Fields
           id          BINARY(16)   NOT NULL,
@@ -128,12 +128,9 @@ export async function enableIndex (id: number) {
 }
 
 export async function updateIndexConfig<Path extends KeyPaths<IndexConfig>> (id: number, configPath: Path, value: Extract<IndexConfig, Path>) {
-  if (configPath.startsWith('.embedding')) {
-    throw new Error('cannot modify embedding config');
-  }
   const { numUpdatedRows } = await getDb()
     .updateTable('index')
-    .set('config', eb => eb.fn<string>('json_set', ['index.config', eb.val('$' + configPath), eb.val(JSON.stringify(value))]))
+    .set('config', eb => eb.fn<string>('json_set', ['index.config', eb.val('$' + configPath), eb.fn('JSON_EXTRACT', [eb.val(JSON.stringify(value)), eb.val('$')])]))
     .set('last_modified_at', new Date())
     .where('id', '=', id)
     .where('configured', '=', 0)