├── src
    ├── vite-env.d.ts
    ├── index.css
    ├── App.css
    ├── Loading.tsx
    ├── main.tsx
    ├── useComparison.ts
    ├── Prompt.tsx
    ├── EnvCheck.tsx
    ├── Compare.tsx
    ├── App.tsx
    ├── Sources.tsx
    ├── AddSource.tsx
    ├── Chunks.tsx
    ├── assets
    │   └── react.svg
    └── Search.tsx
├── screenshot.png
├── public
    └── favicon.ico
├── postcss.config.js
├── vite.config.ts
├── tsconfig.node.json
├── convex
    ├── lib
    │   ├── utils.ts
    │   └── embeddings.ts
    ├── _generated
    │   ├── api.js
    │   ├── api.d.ts
    │   ├── dataModel.d.ts
    │   ├── server.js
    │   └── server.d.ts
    ├── tsconfig.json
    ├── schema.ts
    ├── README.md
    ├── comparisons.ts
    ├── searches.ts
    └── sources.ts
├── .gitignore
├── index.html
├── .eslintrc.mjs
├── tsconfig.json
├── tailwind.config.js
├── package.json
├── scripts
    ├── addURL.py
    └── addFiles.ts
└── README.md


/src/vite-env.d.ts:
--------------------------------------------------------------------------------
1 | /// <reference types="vite/client" />
2 | 


--------------------------------------------------------------------------------
/src/index.css:
--------------------------------------------------------------------------------
1 | @tailwind base;
2 | @tailwind components;
3 | @tailwind utilities;
4 | 


--------------------------------------------------------------------------------
/screenshot.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/ianmacartney/embeddings-in-convex/HEAD/screenshot.png


--------------------------------------------------------------------------------
/src/App.css:
--------------------------------------------------------------------------------
1 | #root {
2 |   max-width: 1280px;
3 |   margin: 0 auto;
4 |   padding: 2rem;
5 | }
6 | 


--------------------------------------------------------------------------------
/public/favicon.ico:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/ianmacartney/embeddings-in-convex/HEAD/public/favicon.ico


--------------------------------------------------------------------------------
/postcss.config.js:
--------------------------------------------------------------------------------
1 | module.exports = {
2 |   plugins: {
3 |     tailwindcss: {},
4 |     autoprefixer: {},
5 |   },
6 | };
7 | 


--------------------------------------------------------------------------------
/vite.config.ts:
--------------------------------------------------------------------------------
1 | import { defineConfig } from 'vite'
2 | import react from '@vitejs/plugin-react'
3 | 
4 | // https://vitejs.dev/config/
5 | export default defineConfig({
6 |   plugins: [react()],
7 | })
8 | 


--------------------------------------------------------------------------------
/src/Loading.tsx:
--------------------------------------------------------------------------------
1 | import { Progress } from "@rewind-ui/core";
2 | 
3 | export function Loading() {
4 |   return (
5 |     <Progress animated={true} showValue={false} value={100} striped={true} />
6 |   );
7 | }
8 | 


--------------------------------------------------------------------------------
/tsconfig.node.json:
--------------------------------------------------------------------------------
 1 | {
 2 |   "compilerOptions": {
 3 |     "composite": true,
 4 |     "skipLibCheck": true,
 5 |     "module": "ESNext",
 6 |     "moduleResolution": "bundler",
 7 |     "allowSyntheticDefaultImports": true
 8 |   },
 9 |   "include": ["vite.config.ts"]
10 | }
11 | 


--------------------------------------------------------------------------------
/convex/lib/utils.ts:
--------------------------------------------------------------------------------
1 | /**
2 |  * Filters out null elements from an array.
3 |  * @param list List of elements that might be null.
4 |  * @returns List of elements with nulls removed.
5 |  */
6 | export function pruneNull<T>(list: (T | null)[]): T[] {
7 |   return list.filter((i) => i !== null) as T[];
8 | }
9 | 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
 1 | # Logs
 2 | logs
 3 | *.log
 4 | npm-debug.log*
 5 | yarn-debug.log*
 6 | yarn-error.log*
 7 | pnpm-debug.log*
 8 | lerna-debug.log*
 9 | 
10 | node_modules
11 | dist
12 | dist-ssr
13 | *.local
14 | 
15 | # Editor directories and files
16 | .vscode/*
17 | !.vscode/extensions.json
18 | .idea
19 | .DS_Store
20 | *.suo
21 | *.ntvs*
22 | *.njsproj
23 | *.sln
24 | *.sw?
25 | 


--------------------------------------------------------------------------------
/index.html:
--------------------------------------------------------------------------------
 1 | <!DOCTYPE html>
 2 | <html lang="en">
 3 |   <head>
 4 |     <meta charset="UTF-8" />
 5 |     <link rel="icon" type="image/icon" href="/favicon.ico" />
 6 |     <meta name="viewport" content="width=device-width, initial-scale=1.0" />
 7 |     <title>Embeddings Playground: OpenAI + Convex</title>
 8 |   </head>
 9 |   <body>
10 |     <div id="root"></div>
11 |     <script type="module" src="/src/main.tsx"></script>
12 |   </body>
13 | </html>
14 | 


--------------------------------------------------------------------------------
/src/main.tsx:
--------------------------------------------------------------------------------
 1 | import React from "react";
 2 | import ReactDOM from "react-dom/client";
 3 | import App from "./App.tsx";
 4 | import "./index.css";
 5 | import { ConvexProvider, ConvexReactClient } from "convex/react";
 6 | 
 7 | const convex = new ConvexReactClient(import.meta.env.VITE_CONVEX_URL);
 8 | 
 9 | ReactDOM.createRoot(document.getElementById("root") as HTMLElement).render(
10 |   <React.StrictMode>
11 |     <ConvexProvider client={convex}>
12 |       <App />
13 |     </ConvexProvider>
14 |   </React.StrictMode>
15 | );
16 | 


--------------------------------------------------------------------------------
/.eslintrc.mjs:
--------------------------------------------------------------------------------
 1 | module.exports = {
 2 |   env: { browser: true, es2020: true },
 3 |   extends: [
 4 |     "eslint:recommended",
 5 |     "plugin:@typescript-eslint/recommended",
 6 |     "plugin:react-hooks/recommended",
 7 |   ],
 8 |   parser: "@typescript-eslint/parser",
 9 |   parserOptions: { ecmaVersion: "latest", sourceType: "module" },
10 |   plugins: ["react-refresh"],
11 |   rules: {
12 |     "react-refresh/only-export-components": "warn",
13 |     "@typescript-eslint/no-non-null-assertion": "off",
14 |     "@typescript-eslint/no-floating-promises": "error",
15 |   },
16 | };
17 | 


--------------------------------------------------------------------------------
/convex/_generated/api.js:
--------------------------------------------------------------------------------
 1 | /* prettier-ignore-start */
 2 | 
 3 | /* eslint-disable */
 4 | /**
 5 |  * Generated `api` utility.
 6 |  *
 7 |  * THIS CODE IS AUTOMATICALLY GENERATED.
 8 |  *
 9 |  * To regenerate, run `npx convex dev`.
10 |  * @module
11 |  */
12 | 
13 | import { anyApi } from "convex/server";
14 | 
15 | /**
16 |  * A utility for referencing Convex functions in your app's API.
17 |  *
18 |  * Usage:
19 |  * ```js
20 |  * const myFunctionReference = api.myModule.myFunction;
21 |  * ```
22 |  */
23 | export const api = anyApi;
24 | export const internal = anyApi;
25 | 
26 | /* prettier-ignore-end */
27 | 


--------------------------------------------------------------------------------
/tsconfig.json:
--------------------------------------------------------------------------------
 1 | {
 2 |   "compilerOptions": {
 3 |     "target": "ESNext",
 4 |     "lib": ["DOM", "DOM.Iterable", "ESNext"],
 5 |     "module": "ESNext",
 6 |     "skipLibCheck": true,
 7 | 
 8 |     /* Bundler mode */
 9 |     "moduleResolution": "bundler",
10 |     "allowImportingTsExtensions": true,
11 |     "resolveJsonModule": true,
12 |     "isolatedModules": true,
13 |     "noEmit": true,
14 |     "jsx": "react-jsx",
15 | 
16 |     /* Linting */
17 |     "strict": true,
18 |     "noUnusedLocals": true,
19 |     "noUnusedParameters": true,
20 |     "noFallthroughCasesInSwitch": true
21 |   },
22 |   "include": ["src"],
23 |   "references": [{ "path": "./tsconfig.node.json" }]
24 | }
25 | 


--------------------------------------------------------------------------------
/tailwind.config.js:
--------------------------------------------------------------------------------
 1 | /** @type {import('tailwindcss').Config} */
 2 | module.exports = {
 3 |   content: [
 4 |     "./index.html",
 5 |     "./src/**/*.{html,jsx,tsx}",
 6 |     // you can either add all styles
 7 |     "./node_modules/@rewind-ui/core/dist/theme/styles/*.js",
 8 |     // OR you can add only the styles you need
 9 |     // "./node_modules/@rewind-ui/core/dist/theme/styles/Button.styles.js",
10 |     // "./node_modules/@rewind-ui/core/dist/theme/styles/Text.styles.js",
11 |   ],
12 |   theme: {
13 |     extend: {},
14 |   },
15 |   plugins: [
16 |     require("@tailwindcss/typography"),
17 |     require("tailwind-scrollbar")({ nocompatible: true }),
18 |     require("@tailwindcss/forms")({
19 |       strategy: "class", // only generate classes
20 |     }),
21 |   ],
22 | };
23 | 


--------------------------------------------------------------------------------
/convex/tsconfig.json:
--------------------------------------------------------------------------------
 1 | {
 2 |   /* This TypeScript project config describes the environment that
 3 |    * Convex functions run in and is used to typecheck them.
 4 |    * You can modify it, but some settings required to use Convex.
 5 |    */
 6 |   "compilerOptions": {
 7 |     /* These settings are not required by Convex and can be modified. */
 8 |     "allowJs": true,
 9 |     "strict": true,
10 | 
11 |     /* These compiler options are required by Convex */
12 |     "target": "ESNext",
13 |     "lib": ["ES2021", "dom"],
14 |     "forceConsistentCasingInFileNames": true,
15 |     "allowSyntheticDefaultImports": true,
16 |     "module": "ESNext",
17 |     "moduleResolution": "Node",
18 |     "isolatedModules": true,
19 |     "noEmit": true
20 |   },
21 |   "include": ["./**/*"],
22 |   "exclude": ["./_generated"]
23 | }
24 | 


--------------------------------------------------------------------------------
/src/useComparison.ts:
--------------------------------------------------------------------------------
 1 | import { useMutation } from "convex/react";
 2 | import { RefObject, useCallback, useState } from "react";
 3 | import { Id } from "../convex/_generated/dataModel";
 4 | import { api } from "../convex/_generated/api";
 5 | 
 6 | export function useComparison(ref: RefObject<HTMLButtonElement>) {
 7 |   const [target, setTarget] = useState<Target>();
 8 |   const upsertComparison = useMutation(api.comparisons.upsert);
 9 |   const compare = useCallback(
10 |     (chunkId?: Id<"chunks">) => {
11 |       if (chunkId) {
12 |         upsertComparison({ target: chunkId, count: 10 }).then((id) =>
13 |           setTarget({ chunkId, comparisonId: id })
14 |         );
15 |         ref.current?.click();
16 |       } else {
17 |         setTarget(undefined);
18 |       }
19 |     },
20 |     [upsertComparison, ref]
21 |   );
22 | 
23 |   return [target, compare] as const;
24 | }
25 | export type Target = { chunkId: Id<"chunks">; comparisonId: Id<"comparisons"> };
26 | export type CompareFn = (id?: Id<"chunks">) => void;
27 | 


--------------------------------------------------------------------------------
/src/Prompt.tsx:
--------------------------------------------------------------------------------
 1 | // import {
 2 | //   useAction,
 3 | //   useMutation,
 4 | //   usePaginatedQuery,
 5 | //   useQuery,
 6 | // } from "convex/react";
 7 | import {
 8 |   Accordion,
 9 |   // Button, Table, Text, InputGroup
10 | } from "@rewind-ui/core";
11 | // import { Id } from "../convex/_generated/dataModel";
12 | // import { api } from "../convex/_generated/api";
13 | // import { Loading } from "./Loading";
14 | // import { Chunks } from "./Chunks";
15 | export function Prompt() {
16 |   return (
17 |     <>
18 |       <Accordion>
19 |         <Accordion.Item anchor="prompt">
20 |           <Accordion.Header>Prompt</Accordion.Header>
21 |           <Accordion.Body></Accordion.Body>
22 |         </Accordion.Item>
23 |         <Accordion.Item anchor="sources">
24 |           <Accordion.Header>Sources</Accordion.Header>
25 |           <Accordion.Body></Accordion.Body>
26 |         </Accordion.Item>
27 |         <Accordion.Item anchor="template">
28 |           <Accordion.Header>Template</Accordion.Header>
29 |           <Accordion.Body></Accordion.Body>
30 |         </Accordion.Item>
31 |         <Accordion.Item anchor="history">
32 |           <Accordion.Header>History</Accordion.Header>
33 |           <Accordion.Body></Accordion.Body>
34 |         </Accordion.Item>
35 |       </Accordion>
36 |     </>
37 |   );
38 | }
39 | 


--------------------------------------------------------------------------------
/src/EnvCheck.tsx:
--------------------------------------------------------------------------------
 1 | import { ReactNode } from "react";
 2 | import { useQuery } from "convex/react";
 3 | import { Alert } from "@rewind-ui/core";
 4 | import { api } from "../convex/_generated/api";
 5 | 
 6 | export function EnvCheck({ children }: { children: ReactNode }) {
 7 |   const embeddingsEnv = useQuery(api.lib.embeddings.envCheck) ?? {};
 8 |   const envCheck = { ...embeddingsEnv };
 9 |   const missingEnv = Object.entries(envCheck).reduce<string[]>(
10 |     (badOnes, [key, value]) => (value ? badOnes : [...badOnes, key]),
11 |     []
12 |   );
13 |   if (missingEnv.length) {
14 |     return (
15 |       <Alert color="red">
16 |         You are missing the following Environment Variables:
17 |         <ol>
18 |           {missingEnv.map((missing) => (
19 |             <li key={missing}>{missing}</li>
20 |           ))}
21 |         </ol>
22 |         Please enter them on{" "}
23 |         <a className="text-blue-500" href="https://dashboard.convex.dev">
24 |           the Convex Dashboard
25 |         </a>
26 |         . See{" "}
27 |         <a
28 |           className="text-blue-500"
29 |           href="https://docs.convex.dev/production/environment-variables"
30 |         >
31 |           the docs
32 |         </a>{" "}
33 |         for details.
34 |       </Alert>
35 |     );
36 |   }
37 |   return <>{children}</>;
38 | }
39 | 


--------------------------------------------------------------------------------
/convex/_generated/api.d.ts:
--------------------------------------------------------------------------------
 1 | /* prettier-ignore-start */
 2 | 
 3 | /* eslint-disable */
 4 | /**
 5 |  * Generated `api` utility.
 6 |  *
 7 |  * THIS CODE IS AUTOMATICALLY GENERATED.
 8 |  *
 9 |  * To regenerate, run `npx convex dev`.
10 |  * @module
11 |  */
12 | 
13 | import type {
14 |   ApiFromModules,
15 |   FilterApi,
16 |   FunctionReference,
17 | } from "convex/server";
18 | import type * as comparisons from "../comparisons.js";
19 | import type * as lib_embeddings from "../lib/embeddings.js";
20 | import type * as lib_utils from "../lib/utils.js";
21 | import type * as searches from "../searches.js";
22 | import type * as sources from "../sources.js";
23 | 
24 | /**
25 |  * A utility for referencing Convex functions in your app's API.
26 |  *
27 |  * Usage:
28 |  * ```js
29 |  * const myFunctionReference = api.myModule.myFunction;
30 |  * ```
31 |  */
32 | declare const fullApi: ApiFromModules<{
33 |   comparisons: typeof comparisons;
34 |   "lib/embeddings": typeof lib_embeddings;
35 |   "lib/utils": typeof lib_utils;
36 |   searches: typeof searches;
37 |   sources: typeof sources;
38 | }>;
39 | export declare const api: FilterApi<
40 |   typeof fullApi,
41 |   FunctionReference<any, "public">
42 | >;
43 | export declare const internal: FilterApi<
44 |   typeof fullApi,
45 |   FunctionReference<any, "internal">
46 | >;
47 | 
48 | /* prettier-ignore-end */
49 | 


--------------------------------------------------------------------------------
/package.json:
--------------------------------------------------------------------------------
 1 | {
 2 |   "name": "embeddings-in-convex",
 3 |   "private": true,
 4 |   "version": "0.0.0",
 5 |   "scripts": {
 6 |     "dev": "npm-run-all --parallel dev:backend dev:frontend",
 7 |     "build": "tsc && vite build",
 8 |     "dev:backend": "convex dev",
 9 |     "dev:frontend": "vite --open --clearScreen false",
10 |     "predev": "convex dev --until-success"
11 |   },
12 |   "dependencies": {
13 |     "@phosphor-icons/react": "^2.0.10",
14 |     "@rewind-ui/core": "^0.12.2",
15 |     "@tailwindcss/forms": "^0.5.3",
16 |     "@tailwindcss/typography": "^0.5.9",
17 |     "convex": "^1.16.0",
18 |     "convex-helpers": "^0.1.58",
19 |     "langchain": "^0.0.92",
20 |     "openai": "^3.2.1",
21 |     "react": "^18.2.0",
22 |     "react-dom": "^18.2.0",
23 |     "react-hook-form": "^7.44.3",
24 |     "tailwind-scrollbar": "^3.0.4"
25 |   },
26 |   "devDependencies": {
27 |     "@types/node": "^20.2.3",
28 |     "@types/react": "^18.0.28",
29 |     "@types/react-dom": "^18.0.11",
30 |     "@typescript-eslint/eslint-plugin": "^5.57.1",
31 |     "@typescript-eslint/parser": "^5.57.1",
32 |     "@vitejs/plugin-react": "^4.0.0",
33 |     "autoprefixer": "^10.4.14",
34 |     "eslint": "^8.38.0",
35 |     "eslint-plugin-react-hooks": "^4.6.0",
36 |     "eslint-plugin-react-refresh": "^0.3.4",
37 |     "npm-run-all": "^4.1.5",
38 |     "postcss": "^8.4.24",
39 |     "tailwindcss": "^3.3.2",
40 |     "typescript": "^5.0.2",
41 |     "vite": "^4.3.2"
42 |   }
43 | }
44 | 


--------------------------------------------------------------------------------
/convex/lib/embeddings.ts:
--------------------------------------------------------------------------------
 1 | import { query } from "../_generated/server";
 2 | 
 3 | export async function fetchEmbeddingBatch(texts: string[]) {
 4 |   const start = Date.now();
 5 |   const result = await fetch("https://api.openai.com/v1/embeddings", {
 6 |     method: "POST",
 7 |     headers: {
 8 |       "Content-Type": "application/json",
 9 |       Authorization: "Bearer " + process.env.OPENAI_API_KEY,
10 |     },
11 | 
12 |     body: JSON.stringify({
13 |       model: "text-embedding-ada-002",
14 |       input: texts.map((text) => text.replace(/\n/g, " ")),
15 |     }),
16 |   });
17 |   const embeddingMs = Date.now() - start;
18 | 
19 |   const jsonresults = await result.json();
20 |   if (jsonresults.data.length !== texts.length) {
21 |     console.error(result);
22 |     throw new Error("Unexpected number of embeddings");
23 |   }
24 |   const allembeddings = jsonresults.data as {
25 |     embedding: number[];
26 |     index: number;
27 |   }[];
28 |   allembeddings.sort((a, b) => a.index - b.index);
29 |   return {
30 |     embeddings: allembeddings.map(({ embedding }) => embedding),
31 |     totalTokens: jsonresults.usage.total_tokens,
32 |     embeddingMs,
33 |   };
34 | }
35 | 
36 | export async function fetchEmbedding(text: string) {
37 |   const { embeddings, ...stats } = await fetchEmbeddingBatch([text]);
38 |   return { embedding: embeddings[0], ...stats };
39 | }
40 | 
41 | export const envCheck = query(async () => {
42 |   return {
43 |     OPENAI_API_KEY: !!process.env.OPENAI_API_KEY,
44 |   };
45 | });
46 | 


--------------------------------------------------------------------------------
/scripts/addURL.py:
--------------------------------------------------------------------------------
 1 | """ Import files into Convex using Langchain document loaders
 2 | 
 3 | Setup:
 4 | !pip install "playwright"
 5 | !pip install "unstructured"
 6 | !pip install "convex"
 7 | !pip install "python-dotenv"
 8 | !pip install tiktoken
 9 | 
10 | !playwright install
11 | """
12 | 
13 | import os
14 | import sys
15 | from dotenv import load_dotenv
16 | from convex import ConvexClient
17 | from langchain.document_loaders import PlaywrightURLLoader
18 | from langchain.text_splitter import CharacterTextSplitter
19 | 
20 | 
21 | urls = sys.argv[1:]
22 | loader = PlaywrightURLLoader(urls=urls, remove_selectors=["header", "footer"])
23 | data = loader.load()
24 | text_splitter = CharacterTextSplitter.from_tiktoken_encoder(
25 |     chunk_size=100, chunk_overlap=0
26 | )
27 | texts = text_splitter.split_text(data[0].page_content)
28 | 
29 | 
30 | load_dotenv(".env.local")
31 | load_dotenv()
32 | 
33 | backend = os.getenv("VITE_CONVEX_URL")
34 | if not backend:
35 |     raise KeyError("Missing VITE_CONVEX_URL")
36 | 
37 | client = ConvexClient(backend)
38 | print(
39 |     client.action(
40 |         "sources:add",
41 |         dict(
42 |             name=data[0].metadata["source"],
43 |             chunks=list(
44 |                 map(
45 |                     lambda chunk: dict(
46 |                         text=chunk,
47 |                         # TODO: add real line numbers
48 |                         lines={"from": 0, "to": 0},
49 |                     ),
50 |                     texts,
51 |                 )
52 |             ),
53 |         ),
54 |     )
55 | )
56 | 


--------------------------------------------------------------------------------
/src/Compare.tsx:
--------------------------------------------------------------------------------
 1 | import { useQuery } from "convex/react";
 2 | import { Button, Text } from "@rewind-ui/core";
 3 | import { api } from "../convex/_generated/api";
 4 | import { Loading } from "./Loading";
 5 | import { AllChunks, Chunks } from "./Chunks";
 6 | import { Target, CompareFn } from "./useComparison";
 7 | 
 8 | export function Compare({
 9 |   target,
10 |   compare,
11 | }: {
12 |   target?: Target;
13 |   compare: CompareFn;
14 | }) {
15 |   return (
16 |     <>
17 |       {target ? (
18 |         <ComparisonResults target={target} compare={compare} />
19 |       ) : (
20 |         <AllChunks compare={compare} />
21 |       )}
22 |     </>
23 |   );
24 | }
25 | 
26 | function ComparisonResults({
27 |   target: { comparisonId, chunkId },
28 |   compare,
29 | }: {
30 |   target: Target;
31 |   compare: CompareFn;
32 | }) {
33 |   const comparison = useQuery(api.comparisons.get, { comparisonId });
34 |   return comparison?.relatedChunks.length ? (
35 |     <>
36 |       <div className="flex">
37 |         <Text className="line-clamp-3">
38 |           Results for {comparison.target.sourceName} (
39 |           {comparison.target.chunkIndex}): "{comparison.target.text}
40 |         </Text>
41 |         <Button onClick={() => compare()} color="red">
42 |           Clear
43 |         </Button>
44 |       </div>
45 |       <Chunks chunks={comparison.relatedChunks} compare={compare} />{" "}
46 |     </>
47 |   ) : (
48 |     <>
49 |       <Text className="line-clamp-3">
50 |         Comparing against "{comparison?.target?.text ?? chunkId}"...
51 |       </Text>
52 |       <Loading />
53 |     </>
54 |   );
55 | }
56 | 


--------------------------------------------------------------------------------
/src/App.tsx:
--------------------------------------------------------------------------------
 1 | import "./App.css";
 2 | 
 3 | import { Tabs, Text } from "@rewind-ui/core";
 4 | import { AddSource } from "./AddSource";
 5 | import { EnvCheck } from "./EnvCheck";
 6 | import { Sources } from "./Sources";
 7 | import { Search } from "./Search";
 8 | import { Compare } from "./Compare";
 9 | import { useComparison } from "./useComparison";
10 | import { Prompt } from "./Prompt";
11 | import { useRef } from "react";
12 | 
13 | function App() {
14 |   const compareRef = useRef<HTMLButtonElement>(null);
15 |   console.log(compareRef);
16 |   const [target, compare] = useComparison(compareRef);
17 |   return (
18 |     <EnvCheck>
19 |       <Tabs defaultTab="sources">
20 |         <Tabs.List>
21 |           <Tabs.Tab anchor="sources">Sources</Tabs.Tab>
22 |           <Tabs.Tab anchor="search">Search</Tabs.Tab>
23 |           <Tabs.Tab ref={compareRef} anchor="compare">
24 |             Compare
25 |           </Tabs.Tab>
26 |           <Tabs.Tab id="prompt" anchor="prompt">
27 |             Prompt
28 |           </Tabs.Tab>
29 |         </Tabs.List>
30 |         <Tabs.Content anchor="search">
31 |           <Search compare={compare} />
32 |         </Tabs.Content>
33 |         <Tabs.Content anchor="prompt">
34 |           <Prompt />
35 |         </Tabs.Content>
36 |         <Tabs.Content anchor="compare">
37 |           <Compare target={target} compare={compare} />
38 |         </Tabs.Content>
39 |         <Tabs.Content anchor="sources">
40 |           <Text size="lg">Add a source</Text>
41 |           <p>
42 |             This is the data you will be able to search over and compare
43 |             semantically.
44 |           </p>
45 |           <AddSource />
46 |           <Sources />
47 |         </Tabs.Content>
48 |       </Tabs>
49 |     </EnvCheck>
50 |   );
51 | }
52 | 
53 | export default App;
54 | 


--------------------------------------------------------------------------------
/src/Sources.tsx:
--------------------------------------------------------------------------------
 1 | import { useMutation, usePaginatedQuery } from "convex/react";
 2 | import { Button, Table } from "@rewind-ui/core";
 3 | import { api } from "../convex/_generated/api";
 4 | 
 5 | export function Sources() {
 6 |   const {
 7 |     status,
 8 |     loadMore,
 9 |     results: sources,
10 |   } = usePaginatedQuery(api.sources.paginate, {}, { initialNumItems: 10 });
11 |   const deleteSource = useMutation(api.sources.deleteSource);
12 |   return (
13 |     <>
14 |       <Table>
15 |         <Table.Thead>
16 |           <Table.Tr>
17 |             <Table.Th align="left">Source</Table.Th>
18 |             <Table.Th align="left">Chunks</Table.Th>
19 |             <Table.Th align="left">Content</Table.Th>
20 |             <Table.Th align="left">Tokens</Table.Th>
21 |             <Table.Th align="left"></Table.Th>
22 |           </Table.Tr>
23 |         </Table.Thead>
24 |         <Table.Tbody>
25 |           {sources.map((source) => (
26 |             <Table.Tr key={source._id}>
27 |               <Table.Td>{source.name}</Table.Td>
28 |               <Table.Td>{source.chunkIds.length}</Table.Td>
29 |               <Table.Td>
30 |                 <p className="line-clamp-3">{source.firstChunkText}</p>
31 |               </Table.Td>
32 |               <Table.Td>
33 |                 {source.saved ? source.totalTokens : "Unsaved"}
34 |               </Table.Td>
35 |               <Table.Td>
36 |                 <Button
37 |                   color="red"
38 |                   onClick={() => deleteSource({ id: source._id })}
39 |                 >
40 |                   Delete
41 |                 </Button>
42 |               </Table.Td>
43 |             </Table.Tr>
44 |           ))}
45 |         </Table.Tbody>
46 |       </Table>
47 |       {status !== "Exhausted" && (
48 |         <Button
49 |           onClick={() => loadMore(10)}
50 |           disabled={status !== "CanLoadMore"}
51 |         >
52 |           Load More
53 |         </Button>
54 |       )}
55 |     </>
56 |   );
57 | }
58 | 


--------------------------------------------------------------------------------
/convex/_generated/dataModel.d.ts:
--------------------------------------------------------------------------------
 1 | /* prettier-ignore-start */
 2 | 
 3 | /* eslint-disable */
 4 | /**
 5 |  * Generated data model types.
 6 |  *
 7 |  * THIS CODE IS AUTOMATICALLY GENERATED.
 8 |  *
 9 |  * To regenerate, run `npx convex dev`.
10 |  * @module
11 |  */
12 | 
13 | import type {
14 |   DataModelFromSchemaDefinition,
15 |   DocumentByName,
16 |   TableNamesInDataModel,
17 |   SystemTableNames,
18 | } from "convex/server";
19 | import type { GenericId } from "convex/values";
20 | import schema from "../schema.js";
21 | 
22 | /**
23 |  * The names of all of your Convex tables.
24 |  */
25 | export type TableNames = TableNamesInDataModel<DataModel>;
26 | 
27 | /**
28 |  * The type of a document stored in Convex.
29 |  *
30 |  * @typeParam TableName - A string literal type of the table name (like "users").
31 |  */
32 | export type Doc<TableName extends TableNames> = DocumentByName<
33 |   DataModel,
34 |   TableName
35 | >;
36 | 
37 | /**
38 |  * An identifier for a document in Convex.
39 |  *
40 |  * Convex documents are uniquely identified by their `Id`, which is accessible
41 |  * on the `_id` field. To learn more, see [Document IDs](https://docs.convex.dev/using/document-ids).
42 |  *
43 |  * Documents can be loaded using `db.get(id)` in query and mutation functions.
44 |  *
45 |  * IDs are just strings at runtime, but this type can be used to distinguish them from other
46 |  * strings when type checking.
47 |  *
48 |  * @typeParam TableName - A string literal type of the table name (like "users").
49 |  */
50 | export type Id<TableName extends TableNames | SystemTableNames> =
51 |   GenericId<TableName>;
52 | 
53 | /**
54 |  * A type describing your Convex data model.
55 |  *
56 |  * This type includes information about what tables you have, the type of
57 |  * documents stored in those tables, and the indexes defined on them.
58 |  *
59 |  * This type is used to parameterize methods like `queryGeneric` and
60 |  * `mutationGeneric` to make them type-safe.
61 |  */
62 | export type DataModel = DataModelFromSchemaDefinition<typeof schema>;
63 | 
64 | /* prettier-ignore-end */
65 | 


--------------------------------------------------------------------------------
/src/AddSource.tsx:
--------------------------------------------------------------------------------
 1 | import { useEffect, useState } from "react";
 2 | import { useMutation } from "convex/react";
 3 | import { RecursiveCharacterTextSplitter } from "langchain/text_splitter";
 4 | import { useForm, Controller } from "react-hook-form";
 5 | import { Alert, Input, Textarea, Button } from "@rewind-ui/core";
 6 | import { api } from "../convex/_generated/api";
 7 | 
 8 | export function AddSource() {
 9 |   const createSource = useMutation(api.sources.add);
10 |   const [added, setAdded] = useState("");
11 | 
12 |   const { formState, handleSubmit, control, reset } = useForm<{
13 |     name: string;
14 |     text: string;
15 |   }>({});
16 |   const onSubmit = handleSubmit(({ name, text }) => {
17 |     const textSplitter = new RecursiveCharacterTextSplitter({
18 |       chunkSize: 1000,
19 |     });
20 |     textSplitter.createDocuments([text]).then((docs) => {
21 |       createSource({
22 |         name,
23 |         chunks: docs.map((doc) => ({
24 |           text: doc.pageContent,
25 |           lines: doc.metadata.loc.lines,
26 |         })),
27 |       }).then(() => {
28 |         setAdded(name);
29 |         setTimeout(
30 |           () => setAdded((state) => (state === name ? "" : state)),
31 |           1000
32 |         );
33 |       });
34 |     });
35 |   });
36 |   useEffect(() => {
37 |     if (formState.isSubmitSuccessful) {
38 |       reset();
39 |     }
40 |   }, [formState, reset]);
41 | 
42 |   return (
43 |     <form onSubmit={onSubmit}>
44 |       {added.length ? (
45 |         <Alert color="green">Successfully added {added}</Alert>
46 |       ) : null}
47 |       <Controller
48 |         name="name"
49 |         control={control}
50 |         rules={{ required: true }}
51 |         render={({ field }) => <Input placeholder="Name" {...field} />}
52 |         defaultValue=""
53 |       />
54 |       <Controller
55 |         name="text"
56 |         control={control}
57 |         rules={{ required: true }}
58 |         render={({ field }) => <Textarea placeholder="Text" {...field} />}
59 |         defaultValue=""
60 |       />
61 |       <Button type="submit" color="green">
62 |         Submit
63 |       </Button>
64 |     </form>
65 |   );
66 | }
67 | 


--------------------------------------------------------------------------------
/src/Chunks.tsx:
--------------------------------------------------------------------------------
 1 | import { Button, Table } from "@rewind-ui/core";
 2 | import { Doc } from "../convex/_generated/dataModel";
 3 | import { usePaginatedQuery } from "convex/react";
 4 | import { api } from "../convex/_generated/api";
 5 | import { CompareFn } from "./useComparison";
 6 | 
 7 | export function Chunks({
 8 |   chunks,
 9 |   compare,
10 | }: {
11 |   chunks: (Doc<"chunks"> & { sourceName: string; score?: number })[];
12 |   compare?: CompareFn;
13 | }) {
14 |   return (
15 |     <Table>
16 |       <Table.Thead>
17 |         <Table.Tr>
18 |           <Table.Th align="left">Source</Table.Th>
19 |           <Table.Th align="left">Index</Table.Th>
20 |           <Table.Th align="left">Content</Table.Th>
21 |           <Table.Th align="left">Score</Table.Th>
22 |           {compare && <Table.Th align="left">Compare</Table.Th>}
23 |         </Table.Tr>
24 |       </Table.Thead>
25 |       <Table.Tbody>
26 |         {chunks.map((chunk) => (
27 |           <Table.Tr key={chunk._id}>
28 |             <Table.Td>{chunk.sourceName}</Table.Td>
29 |             <Table.Td>{chunk.chunkIndex}</Table.Td>
30 |             <Table.Td>
31 |               <p className="line-clamp-3">{chunk.text}</p>
32 |             </Table.Td>
33 |             <Table.Td>
34 |               {chunk.score ? `${(chunk.score * 100).toFixed(2)}%` : "-"}
35 |             </Table.Td>
36 |             {compare && (
37 |               <Table.Td>
38 |                 <Button color="blue" onClick={() => compare(chunk._id)}>
39 |                   Compare
40 |                 </Button>
41 |               </Table.Td>
42 |             )}
43 |           </Table.Tr>
44 |         ))}
45 |       </Table.Tbody>
46 |     </Table>
47 |   );
48 | }
49 | 
50 | export function AllChunks({ compare }: { compare: CompareFn }) {
51 |   const {
52 |     status,
53 |     loadMore,
54 |     results: chunks,
55 |   } = usePaginatedQuery(
56 |     api.sources.paginateChunks,
57 |     {},
58 |     { initialNumItems: 10 }
59 |   );
60 |   return (
61 |     <>
62 |       <Chunks chunks={chunks} compare={compare} />
63 |       {status !== "Exhausted" && (
64 |         <Button
65 |           onClick={() => loadMore(10)}
66 |           disabled={status !== "CanLoadMore"}
67 |         >
68 |           Load More
69 |         </Button>
70 |       )}
71 |     </>
72 |   );
73 | }
74 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # Embeddings Playground with OpenAI and Convex
 2 | 
 3 | An example of working with embeddings and vector databases in Convex.
 4 | 
 5 | [Embeddings](https://stack.convex.dev/the-magic-of-embeddings) enable all sorts
 6 | of use cases, but it's hard to know how they'll perform on comparisons and
 7 | queries without playing around with them.
 8 | 
 9 | This project allows you to add source data, generate embeddings via OpenAI,
10 | compare them to each other, and compare semantic and word searches over them.
11 | 
12 | You can then use the queried source data to include in a ChatGPT prompt (WIP).
13 | 
14 | ![](./screenshot.png)
15 | 
16 | UI:
17 | 
18 | - React
19 | - Tailwindcss
20 | - Rewind-UI
21 | - Vite
22 | 
23 | Backend:
24 | 
25 | - OpenAI API for creating vector embeddings.
26 | - Convex for storing vectors, application data, and running server-side functions.
27 | 
28 | Work planned:
29 | 
30 | - [x] Add a python script that scrapes URLs and imports the data.
31 | - [x] Add a node script that imports local files (.pdf, .md, .txt).
32 | - [ ] Allow picking which sources to use in a ChatGPT prompt, and what template to use, to iterate on templates.
33 | - [ ] Configuration to fetch the most 20, 40, or 80 documents when searching (hard-coded to 10 currently).
34 | 
35 | ## Setup
36 | 
37 | ### Prerequisites:
38 | 
39 | 1. A Convex backend: it will be configured automatically on `npm run dev`.
40 |    By running this first, you can enter environment variables for (2) and (3) on
41 |    the [dashboard](https://dashboard.convex.dev).
42 | 
43 | 2. An [OpenAI](https://platform.openai.com/) API key.
44 |    Environment variable: `OPEN_API_KEY` (should start with `sk-`).
45 |    Run `npx convex env set OPEN_API_KEY sk-XXXX # --prod`
46 | 
47 | ## Run:
48 | 
49 | ```bash
50 | npm install
51 | npm run dev
52 | ```
53 | 
54 | ## Upload sources from a URL
55 | 
56 | You can add a source from a URL using the scripts/addURL.py python script:
57 | 
58 | ```sh
59 | pip install dotenv convex langchain
60 | python scripts/addURL.py https://example.com
61 | ```
62 | 
63 | ## Upload sources from a folder
64 | 
65 | You can add .txt, .md, and .pdf files as sources to your project via:
66 | 
67 | ```sh
68 | export VITE_CONVEX_URL= # your backend url - see .env.local (dev) or .env (prod)
69 | npx ts-node-esm scripts/addFiles.ts ./path/to/folder
70 | ```
71 | 
72 | By default it'll check in a documents folder at the root of the repo.
73 | It will upload in chunks
74 | 


--------------------------------------------------------------------------------
/scripts/addFiles.ts:
--------------------------------------------------------------------------------
 1 | import { api } from "../convex/_generated/api.js";
 2 | import path from "path";
 3 | import { ConvexHttpClient } from "convex/browser";
 4 | import { RecursiveCharacterTextSplitter } from "langchain/text_splitter";
 5 | import { TextLoader } from "langchain/document_loaders/fs/text";
 6 | import { DirectoryLoader } from "langchain/document_loaders/fs/directory";
 7 | import { PDFLoader } from "langchain/document_loaders/fs/pdf";
 8 | 
 9 | const ChunkBatchSize = 100;
10 | const ChunkSize = 1000;
11 | 
12 | // Recursively processes all files with these extensions in the given dir.
13 | const loader = new DirectoryLoader(process.argv[2] || "../documents", {
14 |   ".txt": (path) => new TextLoader(path),
15 |   ".md": (path) => new TextLoader(path),
16 |   ".pdf": (path) => new PDFLoader(path),
17 | });
18 | 
19 | // The Convex backend where we're uploading results
20 | const address = process.env.VITE_CONVEX_URL;
21 | if (!address) throw new Error("Specify VITE_CONVEX_URL env variable");
22 | const client = new ConvexHttpClient(address);
23 | 
24 | export const uploadDocuments = async () => {
25 |   const start = Date.now();
26 |   const docs = await loader.load();
27 |   let batch: {
28 |     name: string;
29 |     chunks: { text: string; lines: { from: number; to: number } }[];
30 |   }[] = [];
31 |   for (const doc of docs) {
32 |     console.log(`Processing document: ${doc.metadata.source}`);
33 |     const textSplitter = new RecursiveCharacterTextSplitter({
34 |       chunkSize: ChunkSize,
35 |     });
36 |     const chunks = await textSplitter.createDocuments([doc.pageContent]);
37 |     // If this doc will put us over the batch limit, let's send it
38 |     if (
39 |       batch.length > 0 &&
40 |       batch.reduce((sum, val) => sum + val.chunks.length, 0) >=
41 |         ChunkBatchSize - chunks.length
42 |     ) {
43 |       console.log("Sending up a batch:", batch.length);
44 |       await client.mutation(api.sources.addBatch, { batch });
45 |       batch = [];
46 |     }
47 |     batch.push({
48 |       name: path.parse(doc.metadata.source).base,
49 |       chunks: chunks.map((chunk) => ({
50 |         text: chunk.pageContent,
51 |         lines: chunk.metadata.loc.lines,
52 |       })),
53 |     });
54 |   }
55 |   if (batch.length) {
56 |     console.log("Sending up a final batch:", batch.length);
57 |     await client.mutation(api.sources.addBatch, { batch });
58 |     batch = [];
59 |   }
60 |   console.log("Finished embedding documents. ms:", Date.now() - start);
61 |   return "success";
62 | };
63 | 
64 | uploadDocuments().then(console.log);
65 | 


--------------------------------------------------------------------------------
/convex/schema.ts:
--------------------------------------------------------------------------------
 1 | import { v } from "convex/values";
 2 | import { defineSchema, defineTable } from "convex/server";
 3 | import { deprecated } from "convex-helpers/validators";
 4 | 
 5 | export default defineSchema({
 6 |   // Chunks are one part of a Source, broken up to generate embeddings.
 7 |   chunks: defineTable({
 8 |     // raw text: ~1k bytes or less
 9 |     text: v.string(),
10 |     sourceId: v.id("sources"),
11 |     // Where in a larger document is this text.
12 |     chunkIndex: v.number(),
13 |     lines: v.object({
14 |       from: v.number(),
15 |       to: v.number(),
16 |     }),
17 |     // Approx: estimated based on the total batch size.
18 |     tokens: v.optional(v.number()),
19 |     embeddingId: v.optional(v.id("chunkEmbeddings")),
20 |   })
21 |     .searchIndex("text", { searchField: "text" })
22 |     .index("embeddingId", ["embeddingId"]),
23 | 
24 |   chunkEmbeddings: defineTable({
25 |     vector: v.array(v.number()),
26 |   }).vectorIndex("vector", { vectorField: "vector", dimensions: 1536 }),
27 | 
28 |   // Sources are materials to search over / compare, made of chunks of text.
29 |   sources: defineTable({
30 |     name: v.string(),
31 |     // Max 1k chunks (otherwise remove this and use an index on sourceId)
32 |     chunkIds: v.array(v.id("chunks")),
33 |     // Whether the embeddings have been saved.
34 |     saved: v.boolean(),
35 |     // stats
36 |     totalTokens: v.optional(v.number()),
37 |     embeddingMs: v.optional(v.number()),
38 |   }),
39 | 
40 |   // Searches track a comparison between an input string and related chunks.
41 |   searches: defineTable({
42 |     input: v.string(),
43 |     float32Buffer: v.optional(v.bytes()),
44 |     relatedChunks: v.optional(
45 |       v.array(
46 |         v.object({
47 |           id: v.id("chunkEmbeddings"),
48 |           score: v.optional(v.number()),
49 |         }),
50 |       ),
51 |     ),
52 |     // stats
53 |     count: v.number(),
54 |     inputTokens: v.optional(v.number()),
55 |     embeddingMs: v.optional(v.number()),
56 |     queryMs: v.optional(v.number()),
57 |     saveSearchMs: deprecated,
58 |     embeddingId: v.optional(v.id("searchEmbeddings")),
59 |   })
60 |     .index("input", ["input"])
61 |     .index("embeddingId", ["embeddingId"]),
62 | 
63 |   searchEmbeddings: defineTable({
64 |     vector: v.array(v.number()),
65 |   }).vectorIndex("vector", { vectorField: "vector", dimensions: 1536 }),
66 | 
67 |   // Comparisons track a comparison between one chunk and other chunks.
68 |   comparisons: defineTable({
69 |     target: v.id("chunks"),
70 |     relatedChunks: v.optional(
71 |       v.array(
72 |         v.object({
73 |           id: v.id("chunkEmbeddings"),
74 |           score: v.optional(v.number()),
75 |         }),
76 |       ),
77 |     ),
78 |     // stats
79 |     count: v.number(),
80 |     queryMs: v.optional(v.number()),
81 |   }).index("target", ["target"]),
82 | });
83 | 


--------------------------------------------------------------------------------
/convex/README.md:
--------------------------------------------------------------------------------
 1 | # Welcome to your Convex functions directory!
 2 | 
 3 | Write your Convex functions here. See
 4 | https://docs.convex.dev/using/writing-convex-functions for more.
 5 | 
 6 | A query function that takes two arguments looks like:
 7 | 
 8 | ```ts
 9 | // functions.js
10 | import { query } from "./_generated/server";
11 | import { v } from "convex/values";
12 | 
13 | export const myQueryFunction = query({
14 |   // Validators for arguments.
15 |   args: {
16 |     first: v.number(),
17 |     second: v.string(),
18 |   },
19 | 
20 |   // Function implementation.
21 |   hander: async (ctx, args) => {
22 |     // Read the database as many times as you need here.
23 |     // See https://docs.convex.dev/database/reading-data.
24 |     const documents = await ctx.db.query("tablename").collect();
25 | 
26 |     // Arguments passed from the client are properties of the args object.
27 |     console.log(args.first, args.second);
28 | 
29 |     // Write arbitrary JavaScript here: filter, aggregate, build derived data,
30 |     // remove non-public properties, or create new objects.
31 |     return documents;
32 |   },
33 | });
34 | ```
35 | 
36 | Using this query function in a React component looks like:
37 | 
38 | ```ts
39 | const data = useQuery(api.functions.myQueryFunction, {
40 |   first: 10,
41 |   second: "hello",
42 | });
43 | ```
44 | 
45 | A mutation function looks like:
46 | 
47 | ```ts
48 | // functions.js
49 | import { mutation } from "./_generated/server";
50 | import { v } from "convex/values";
51 | 
52 | export const myMutationFunction = mutation({
53 |   // Validators for arguments.
54 |   args: {
55 |     first: v.string(),
56 |     second: v.string(),
57 |   },
58 | 
59 |   // Function implementation.
60 |   hander: async (ctx, args) => {
61 |     // Insert or modify documents in the database here.
62 |     // Mutations can also read from the database like queries.
63 |     // See https://docs.convex.dev/database/writing-data.
64 |     const message = { body: args.first, author: args.second };
65 |     const id = await ctx.db.insert("messages", message);
66 | 
67 |     // Optionally, return a value from your mutation.
68 |     return await ctx.db.get(id);
69 |   },
70 | });
71 | ```
72 | 
73 | Using this mutation function in a React component looks like:
74 | 
75 | ```ts
76 | const mutation = useMutation(api.functions.myMutationFunction);
77 | function handleButtonPress() {
78 |   // fire and forget, the most common way to use mutations
79 |   mutation({ first: "Hello!", second: "me" });
80 |   // OR
81 |   // use the result once the mutation has completed
82 |   mutation({ first: "Hello!", second: "me" }).then((result) =>
83 |     console.log(result)
84 |   );
85 | }
86 | ```
87 | 
88 | Use the Convex CLI to push your functions to a deployment. See everything
89 | the Convex CLI can do by running `npx convex -h` in your project root
90 | directory. To learn more, launch the docs with `npx convex docs`.
91 | 


--------------------------------------------------------------------------------
/convex/_generated/server.js:
--------------------------------------------------------------------------------
 1 | /* prettier-ignore-start */
 2 | 
 3 | /* eslint-disable */
 4 | /**
 5 |  * Generated utilities for implementing server-side Convex query and mutation functions.
 6 |  *
 7 |  * THIS CODE IS AUTOMATICALLY GENERATED.
 8 |  *
 9 |  * To regenerate, run `npx convex dev`.
10 |  * @module
11 |  */
12 | 
13 | import {
14 |   actionGeneric,
15 |   httpActionGeneric,
16 |   queryGeneric,
17 |   mutationGeneric,
18 |   internalActionGeneric,
19 |   internalMutationGeneric,
20 |   internalQueryGeneric,
21 | } from "convex/server";
22 | 
23 | /**
24 |  * Define a query in this Convex app's public API.
25 |  *
26 |  * This function will be allowed to read your Convex database and will be accessible from the client.
27 |  *
28 |  * @param func - The query function. It receives a {@link QueryCtx} as its first argument.
29 |  * @returns The wrapped query. Include this as an `export` to name it and make it accessible.
30 |  */
31 | export const query = queryGeneric;
32 | 
33 | /**
34 |  * Define a query that is only accessible from other Convex functions (but not from the client).
35 |  *
36 |  * This function will be allowed to read from your Convex database. It will not be accessible from the client.
37 |  *
38 |  * @param func - The query function. It receives a {@link QueryCtx} as its first argument.
39 |  * @returns The wrapped query. Include this as an `export` to name it and make it accessible.
40 |  */
41 | export const internalQuery = internalQueryGeneric;
42 | 
43 | /**
44 |  * Define a mutation in this Convex app's public API.
45 |  *
46 |  * This function will be allowed to modify your Convex database and will be accessible from the client.
47 |  *
48 |  * @param func - The mutation function. It receives a {@link MutationCtx} as its first argument.
49 |  * @returns The wrapped mutation. Include this as an `export` to name it and make it accessible.
50 |  */
51 | export const mutation = mutationGeneric;
52 | 
53 | /**
54 |  * Define a mutation that is only accessible from other Convex functions (but not from the client).
55 |  *
56 |  * This function will be allowed to modify your Convex database. It will not be accessible from the client.
57 |  *
58 |  * @param func - The mutation function. It receives a {@link MutationCtx} as its first argument.
59 |  * @returns The wrapped mutation. Include this as an `export` to name it and make it accessible.
60 |  */
61 | export const internalMutation = internalMutationGeneric;
62 | 
63 | /**
64 |  * Define an action in this Convex app's public API.
65 |  *
66 |  * An action is a function which can execute any JavaScript code, including non-deterministic
67 |  * code and code with side-effects, like calling third-party services.
68 |  * They can be run in Convex's JavaScript environment or in Node.js using the "use node" directive.
69 |  * They can interact with the database indirectly by calling queries and mutations using the {@link ActionCtx}.
70 |  *
71 |  * @param func - The action. It receives an {@link ActionCtx} as its first argument.
72 |  * @returns The wrapped action. Include this as an `export` to name it and make it accessible.
73 |  */
74 | export const action = actionGeneric;
75 | 
76 | /**
77 |  * Define an action that is only accessible from other Convex functions (but not from the client).
78 |  *
79 |  * @param func - The function. It receives an {@link ActionCtx} as its first argument.
80 |  * @returns The wrapped function. Include this as an `export` to name it and make it accessible.
81 |  */
82 | export const internalAction = internalActionGeneric;
83 | 
84 | /**
85 |  * Define a Convex HTTP action.
86 |  *
87 |  * @param func - The function. It receives an {@link ActionCtx} as its first argument, and a `Request` object
88 |  * as its second.
89 |  * @returns The wrapped endpoint function. Route a URL path to this function in `convex/http.js`.
90 |  */
91 | export const httpAction = httpActionGeneric;
92 | 
93 | /* prettier-ignore-end */
94 | 


--------------------------------------------------------------------------------
/src/assets/react.svg:
--------------------------------------------------------------------------------
1 | <svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" class="iconify iconify--logos" width="35.93" height="32" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 228"><path fill="#00D8FF" d="M210.483 73.824a171.49 171.49 0 0 0-8.24-2.597c.465-1.9.893-3.777 1.273-5.621c6.238-30.281 2.16-54.676-11.769-62.708c-13.355-7.7-35.196.329-57.254 19.526a171.23 171.23 0 0 0-6.375 5.848a155.866 155.866 0 0 0-4.241-3.917C100.759 3.829 77.587-4.822 63.673 3.233C50.33 10.957 46.379 33.89 51.995 62.588a170.974 170.974 0 0 0 1.892 8.48c-3.28.932-6.445 1.924-9.474 2.98C17.309 83.498 0 98.307 0 113.668c0 15.865 18.582 31.778 46.812 41.427a145.52 145.52 0 0 0 6.921 2.165a167.467 167.467 0 0 0-2.01 9.138c-5.354 28.2-1.173 50.591 12.134 58.266c13.744 7.926 36.812-.22 59.273-19.855a145.567 145.567 0 0 0 5.342-4.923a168.064 168.064 0 0 0 6.92 6.314c21.758 18.722 43.246 26.282 56.54 18.586c13.731-7.949 18.194-32.003 12.4-61.268a145.016 145.016 0 0 0-1.535-6.842c1.62-.48 3.21-.974 4.76-1.488c29.348-9.723 48.443-25.443 48.443-41.52c0-15.417-17.868-30.326-45.517-39.844Zm-6.365 70.984c-1.4.463-2.836.91-4.3 1.345c-3.24-10.257-7.612-21.163-12.963-32.432c5.106-11 9.31-21.767 12.459-31.957c2.619.758 5.16 1.557 7.61 2.4c23.69 8.156 38.14 20.213 38.14 29.504c0 9.896-15.606 22.743-40.946 31.14Zm-10.514 20.834c2.562 12.94 2.927 24.64 1.23 33.787c-1.524 8.219-4.59 13.698-8.382 15.893c-8.067 4.67-25.32-1.4-43.927-17.412a156.726 156.726 0 0 1-6.437-5.87c7.214-7.889 14.423-17.06 21.459-27.246c12.376-1.098 24.068-2.894 34.671-5.345a134.17 134.17 0 0 1 1.386 6.193ZM87.276 214.515c-7.882 2.783-14.16 2.863-17.955.675c-8.075-4.657-11.432-22.636-6.853-46.752a156.923 156.923 0 0 1 1.869-8.499c10.486 2.32 22.093 3.988 34.498 4.994c7.084 9.967 14.501 19.128 21.976 27.15a134.668 134.668 0 0 1-4.877 4.492c-9.933 8.682-19.886 14.842-28.658 17.94ZM50.35 144.747c-12.483-4.267-22.792-9.812-29.858-15.863c-6.35-5.437-9.555-10.836-9.555-15.216c0-9.322 13.897-21.212 37.076-29.293c2.813-.98 5.757-1.905 8.812-2.773c3.204 10.42 7.406 21.315 12.477 32.332c-5.137 11.18-9.399 22.249-12.634 32.792a134.718 134.718 0 0 1-6.318-1.979Zm12.378-84.26c-4.811-24.587-1.616-43.134 6.425-47.789c8.564-4.958 27.502 2.111 47.463 19.835a144.318 144.318 0 0 1 3.841 3.545c-7.438 7.987-14.787 17.08-21.808 26.988c-12.04 1.116-23.565 2.908-34.161 5.309a160.342 160.342 0 0 1-1.76-7.887Zm110.427 27.268a347.8 347.8 0 0 0-7.785-12.803c8.168 1.033 15.994 2.404 23.343 4.08c-2.206 7.072-4.956 14.465-8.193 22.045a381.151 381.151 0 0 0-7.365-13.322Zm-45.032-43.861c5.044 5.465 10.096 11.566 15.065 18.186a322.04 322.04 0 0 0-30.257-.006c4.974-6.559 10.069-12.652 15.192-18.18ZM82.802 87.83a323.167 323.167 0 0 0-7.227 13.238c-3.184-7.553-5.909-14.98-8.134-22.152c7.304-1.634 15.093-2.97 23.209-3.984a321.524 321.524 0 0 0-7.848 12.897Zm8.081 65.352c-8.385-.936-16.291-2.203-23.593-3.793c2.26-7.3 5.045-14.885 8.298-22.6a321.187 321.187 0 0 0 7.257 13.246c2.594 4.48 5.28 8.868 8.038 13.147Zm37.542 31.03c-5.184-5.592-10.354-11.779-15.403-18.433c4.902.192 9.899.29 14.978.29c5.218 0 10.376-.117 15.453-.343c-4.985 6.774-10.018 12.97-15.028 18.486Zm52.198-57.817c3.422 7.8 6.306 15.345 8.596 22.52c-7.422 1.694-15.436 3.058-23.88 4.071a382.417 382.417 0 0 0 7.859-13.026a347.403 347.403 0 0 0 7.425-13.565Zm-16.898 8.101a358.557 358.557 0 0 1-12.281 19.815a329.4 329.4 0 0 1-23.444.823c-7.967 0-15.716-.248-23.178-.732a310.202 310.202 0 0 1-12.513-19.846h.001a307.41 307.41 0 0 1-10.923-20.627a310.278 310.278 0 0 1 10.89-20.637l-.001.001a307.318 307.318 0 0 1 12.413-19.761c7.613-.576 15.42-.876 23.31-.876H128c7.926 0 15.743.303 23.354.883a329.357 329.357 0 0 1 12.335 19.695a358.489 358.489 0 0 1 11.036 20.54a329.472 329.472 0 0 1-11 20.722Zm22.56-122.124c8.572 4.944 11.906 24.881 6.52 51.026c-.344 1.668-.73 3.367-1.15 5.09c-10.622-2.452-22.155-4.275-34.23-5.408c-7.034-10.017-14.323-19.124-21.64-27.008a160.789 160.789 0 0 1 5.888-5.4c18.9-16.447 36.564-22.941 44.612-18.3ZM128 90.808c12.625 0 22.86 10.235 22.86 22.86s-10.235 22.86-22.86 22.86s-22.86-10.235-22.86-22.86s10.235-22.86 22.86-22.86Z"></path></svg>


--------------------------------------------------------------------------------
/convex/comparisons.ts:
--------------------------------------------------------------------------------
  1 | import { internal } from "./_generated/api";
  2 | import {
  3 |   internalAction,
  4 |   internalMutation,
  5 |   mutation,
  6 |   query,
  7 | } from "./_generated/server";
  8 | import { v } from "convex/values";
  9 | import { pruneNull } from "./lib/utils";
 10 | 
 11 | export const upsert = mutation({
 12 |   args: { target: v.id("chunks"), count: v.optional(v.number()) },
 13 |   handler: async (ctx, { target, count }) => {
 14 |     const topK = count || 10;
 15 |     const existing = await ctx.db
 16 |       .query("comparisons")
 17 |       .withIndex("target", (q) => q.eq("target", target))
 18 |       .filter((q) => q.gte(q.field("count"), topK))
 19 |       .unique();
 20 |     if (existing) return existing._id;
 21 |     const comparisonId = await ctx.db.insert("comparisons", {
 22 |       target,
 23 |       count: topK,
 24 |     });
 25 |     const chunk = await ctx.db.get(target);
 26 |     if (!chunk) throw new Error("Unknown chunk");
 27 |     if (!chunk.embeddingId) throw new Error("Chunk has no embedding yet");
 28 |     const embedding = await ctx.db.get(chunk.embeddingId);
 29 |     if (!embedding) throw new Error("Unknown embedding");
 30 |     await ctx.scheduler.runAfter(0, internal.comparisons.compare, {
 31 |       vector: embedding.vector,
 32 |       comparisonId,
 33 |       topK,
 34 |     });
 35 |     return comparisonId;
 36 |   },
 37 | });
 38 | 
 39 | export const compare = internalAction({
 40 |   args: {
 41 |     vector: v.array(v.number()),
 42 |     comparisonId: v.optional(v.id("comparisons")),
 43 |     topK: v.number(),
 44 |   },
 45 |   handler: async (ctx, { vector, comparisonId, topK }) => {
 46 |     const start = Date.now();
 47 |     const matches = await ctx.vectorSearch("chunkEmbeddings", "vector", {
 48 |       vector,
 49 |       limit: topK,
 50 |     });
 51 |     if (!matches) throw new Error("Vector search matches are empty");
 52 |     const relatedChunks = matches.map(({ _id, _score }) => ({
 53 |       id: _id,
 54 |       score: _score,
 55 |     }));
 56 |     const queryMs = Date.now() - start;
 57 |     console.log({
 58 |       queryMs,
 59 |     });
 60 |     if (comparisonId) {
 61 |       await ctx.runMutation(internal.comparisons.patch, {
 62 |         id: comparisonId,
 63 |         patch: {
 64 |           relatedChunks,
 65 |           // stats
 66 |           queryMs,
 67 |         },
 68 |       });
 69 |     }
 70 |     return relatedChunks;
 71 |   },
 72 | });
 73 | 
 74 | export const get = query({
 75 |   args: { comparisonId: v.id("comparisons") },
 76 |   handler: async (ctx, { comparisonId }) => {
 77 |     const comparison = await ctx.db.get(comparisonId);
 78 |     if (!comparison) throw new Error("Unknown comparison");
 79 |     if (!comparison.relatedChunks) return null;
 80 |     const target = await ctx.db.get(comparison.target);
 81 |     return (
 82 |       target && {
 83 |         ...comparison,
 84 |         relatedChunks: pruneNull(
 85 |           await Promise.all(
 86 |             comparison.relatedChunks.map(async ({ id, score }) => {
 87 |               const chunk = await ctx.db
 88 |                 .query("chunks")
 89 |                 .withIndex("embeddingId", (q) => q.eq("embeddingId", id))
 90 |                 .unique();
 91 |               if (chunk?._id === comparison.target) return null;
 92 |               if (!chunk) throw new Error("Unknown chunk for embedding" + id);
 93 |               const source = await ctx.db.get(chunk.sourceId);
 94 |               if (!source) throw new Error("Unknown source" + chunk.sourceId);
 95 |               return { ...chunk, score, sourceName: source!.name };
 96 |             }),
 97 |           ),
 98 |         ),
 99 |         target: {
100 |           ...target,
101 |           sourceName: (await ctx.db.get(target.sourceId))!.name,
102 |         },
103 |       }
104 |     );
105 |   },
106 | });
107 | 
108 | export const patch = internalMutation({
109 |   args: { id: v.id("comparisons"), patch: v.any() },
110 |   handler: async (ctx, { id, patch }) => {
111 |     return await ctx.db.patch(id, patch);
112 |   },
113 | });
114 | 


--------------------------------------------------------------------------------
/convex/searches.ts:
--------------------------------------------------------------------------------
  1 | import { api, internal } from "./_generated/api";
  2 | import { action, internalMutation, mutation, query } from "./_generated/server";
  3 | import { fetchEmbedding } from "./lib/embeddings";
  4 | import { v } from "convex/values";
  5 | import { pruneNull } from "./lib/utils";
  6 | import { pick } from "convex-helpers";
  7 | import schema from "./schema";
  8 | import { crud } from "convex-helpers/server/crud";
  9 | import { getOrThrow } from "convex-helpers/server/relationships";
 10 | 
 11 | export const upsert = mutation({
 12 |   args: { input: v.string(), count: v.optional(v.number()) },
 13 |   handler: async (ctx, { input, count: countOpt }) => {
 14 |     const count = countOpt || 10;
 15 |     const existing = await ctx.db
 16 |       .query("searches")
 17 |       .withIndex("input", (q) => q.eq("input", input))
 18 |       .filter((q) => q.gte(q.field("count"), count))
 19 |       .unique();
 20 |     if (existing) {
 21 |       console.log("Re-using search for", input);
 22 |       return existing._id;
 23 |     }
 24 |     const searchId = await ctx.db.insert("searches", { input, count });
 25 |     console.log("Starting search for", input);
 26 |     await ctx.scheduler.runAfter(0, api.searches.search, {
 27 |       input,
 28 |       searchId,
 29 |       topK: count,
 30 |     });
 31 |     return searchId;
 32 |   },
 33 | });
 34 | 
 35 | export const search = action({
 36 |   args: {
 37 |     input: v.string(),
 38 |     topK: v.number(),
 39 |     searchId: v.optional(v.id("searches")),
 40 |   },
 41 |   handler: async (ctx, { input, topK, searchId }) => {
 42 |     const {
 43 |       embedding,
 44 |       totalTokens: inputTokens,
 45 |       embeddingMs,
 46 |     } = await fetchEmbedding(input);
 47 |     const start = Date.now();
 48 |     const relatedChunks = (
 49 |       await ctx.vectorSearch("chunkEmbeddings", "vector", {
 50 |         vector: embedding,
 51 |         limit: topK,
 52 |       })
 53 |     ).map(({ _id: id, _score: score }) => ({ id, score }));
 54 |     const queryMs = Date.now() - start;
 55 |     if (searchId) {
 56 |       await ctx.runMutation(internal.searches.storeEmbedding, {
 57 |         embedding,
 58 |         searchId,
 59 |         relatedChunks,
 60 |         // stats
 61 |         inputTokens,
 62 |         embeddingMs,
 63 |         queryMs,
 64 |       });
 65 |       const saveSearchMs = Date.now() - start - queryMs;
 66 |       console.log({
 67 |         inputTokens,
 68 |         embeddingMs,
 69 |         queryMs,
 70 |         saveSearchMs,
 71 |       });
 72 |     }
 73 |   },
 74 | });
 75 | 
 76 | export const storeEmbedding = internalMutation({
 77 |   args: {
 78 |     searchId: v.id("searches"),
 79 |     embedding: v.array(v.number()),
 80 |     ...pick(schema.tables.searches.validator.fields, [
 81 |       "inputTokens",
 82 |       "embeddingMs",
 83 |       "queryMs",
 84 |       "relatedChunks",
 85 |     ]),
 86 |   },
 87 |   handler: async (ctx, { searchId, embedding, ...patch }) => {
 88 |     const search = await getOrThrow(ctx, searchId);
 89 |     if (search.embeddingId) {
 90 |       await ctx.db.patch(search.embeddingId, { vector: embedding });
 91 |     } else {
 92 |       const embeddingId = await ctx.db.insert("searchEmbeddings", {
 93 |         vector: embedding,
 94 |       });
 95 |       await ctx.db.patch(searchId, { embeddingId, ...patch });
 96 |     }
 97 |   },
 98 | });
 99 | 
100 | export const wordSearch = query({
101 |   args: { input: v.string(), count: v.number() },
102 |   handler: async (ctx, { input, count }) => {
103 |     const results = await ctx.db
104 |       .query("chunks")
105 |       .withSearchIndex("text", (q) => q.search("text", input))
106 |       .take(count);
107 |     return Promise.all(
108 |       results.map(async (chunk) => {
109 |         const source = await ctx.db.get(chunk.sourceId);
110 |         if (!source) throw new Error("Missing source for chunk " + chunk._id);
111 |         return { ...chunk, sourceName: source.name };
112 |       })
113 |     );
114 |   },
115 | });
116 | 
117 | export const semanticSearch = query({
118 |   args: { searchId: v.id("searches") },
119 |   handler: async (ctx, { searchId }) => {
120 |     const search = await ctx.db.get(searchId);
121 |     if (!search) throw new Error("Unknown search " + searchId);
122 |     if (!search.relatedChunks) return null;
123 |     return pruneNull(
124 |       await Promise.all(
125 |         search.relatedChunks.map(async ({ id, score }) => {
126 |           const chunk = await ctx.db
127 |             .query("chunks")
128 |             .withIndex("embeddingId", (q) => q.eq("embeddingId", id))
129 |             .unique();
130 |           if (!chunk) return null;
131 |           const source = await ctx.db.get(chunk.sourceId);
132 |           return { ...chunk, score, sourceName: source!.name };
133 |         })
134 |       )
135 |     );
136 |   },
137 | });
138 | 
139 | export const { paginate } = crud(schema, "searches", query);
140 | 


--------------------------------------------------------------------------------
/convex/_generated/server.d.ts:
--------------------------------------------------------------------------------
  1 | /* prettier-ignore-start */
  2 | 
  3 | /* eslint-disable */
  4 | /**
  5 |  * Generated utilities for implementing server-side Convex query and mutation functions.
  6 |  *
  7 |  * THIS CODE IS AUTOMATICALLY GENERATED.
  8 |  *
  9 |  * To regenerate, run `npx convex dev`.
 10 |  * @module
 11 |  */
 12 | 
 13 | import {
 14 |   ActionBuilder,
 15 |   HttpActionBuilder,
 16 |   MutationBuilder,
 17 |   QueryBuilder,
 18 |   GenericActionCtx,
 19 |   GenericMutationCtx,
 20 |   GenericQueryCtx,
 21 |   GenericDatabaseReader,
 22 |   GenericDatabaseWriter,
 23 | } from "convex/server";
 24 | import type { DataModel } from "./dataModel.js";
 25 | 
 26 | /**
 27 |  * Define a query in this Convex app's public API.
 28 |  *
 29 |  * This function will be allowed to read your Convex database and will be accessible from the client.
 30 |  *
 31 |  * @param func - The query function. It receives a {@link QueryCtx} as its first argument.
 32 |  * @returns The wrapped query. Include this as an `export` to name it and make it accessible.
 33 |  */
 34 | export declare const query: QueryBuilder<DataModel, "public">;
 35 | 
 36 | /**
 37 |  * Define a query that is only accessible from other Convex functions (but not from the client).
 38 |  *
 39 |  * This function will be allowed to read from your Convex database. It will not be accessible from the client.
 40 |  *
 41 |  * @param func - The query function. It receives a {@link QueryCtx} as its first argument.
 42 |  * @returns The wrapped query. Include this as an `export` to name it and make it accessible.
 43 |  */
 44 | export declare const internalQuery: QueryBuilder<DataModel, "internal">;
 45 | 
 46 | /**
 47 |  * Define a mutation in this Convex app's public API.
 48 |  *
 49 |  * This function will be allowed to modify your Convex database and will be accessible from the client.
 50 |  *
 51 |  * @param func - The mutation function. It receives a {@link MutationCtx} as its first argument.
 52 |  * @returns The wrapped mutation. Include this as an `export` to name it and make it accessible.
 53 |  */
 54 | export declare const mutation: MutationBuilder<DataModel, "public">;
 55 | 
 56 | /**
 57 |  * Define a mutation that is only accessible from other Convex functions (but not from the client).
 58 |  *
 59 |  * This function will be allowed to modify your Convex database. It will not be accessible from the client.
 60 |  *
 61 |  * @param func - The mutation function. It receives a {@link MutationCtx} as its first argument.
 62 |  * @returns The wrapped mutation. Include this as an `export` to name it and make it accessible.
 63 |  */
 64 | export declare const internalMutation: MutationBuilder<DataModel, "internal">;
 65 | 
 66 | /**
 67 |  * Define an action in this Convex app's public API.
 68 |  *
 69 |  * An action is a function which can execute any JavaScript code, including non-deterministic
 70 |  * code and code with side-effects, like calling third-party services.
 71 |  * They can be run in Convex's JavaScript environment or in Node.js using the "use node" directive.
 72 |  * They can interact with the database indirectly by calling queries and mutations using the {@link ActionCtx}.
 73 |  *
 74 |  * @param func - The action. It receives an {@link ActionCtx} as its first argument.
 75 |  * @returns The wrapped action. Include this as an `export` to name it and make it accessible.
 76 |  */
 77 | export declare const action: ActionBuilder<DataModel, "public">;
 78 | 
 79 | /**
 80 |  * Define an action that is only accessible from other Convex functions (but not from the client).
 81 |  *
 82 |  * @param func - The function. It receives an {@link ActionCtx} as its first argument.
 83 |  * @returns The wrapped function. Include this as an `export` to name it and make it accessible.
 84 |  */
 85 | export declare const internalAction: ActionBuilder<DataModel, "internal">;
 86 | 
 87 | /**
 88 |  * Define an HTTP action.
 89 |  *
 90 |  * This function will be used to respond to HTTP requests received by a Convex
 91 |  * deployment if the requests matches the path and method where this action
 92 |  * is routed. Be sure to route your action in `convex/http.js`.
 93 |  *
 94 |  * @param func - The function. It receives an {@link ActionCtx} as its first argument.
 95 |  * @returns The wrapped function. Import this function from `convex/http.js` and route it to hook it up.
 96 |  */
 97 | export declare const httpAction: HttpActionBuilder;
 98 | 
 99 | /**
100 |  * A set of services for use within Convex query functions.
101 |  *
102 |  * The query context is passed as the first argument to any Convex query
103 |  * function run on the server.
104 |  *
105 |  * This differs from the {@link MutationCtx} because all of the services are
106 |  * read-only.
107 |  */
108 | export type QueryCtx = GenericQueryCtx<DataModel>;
109 | 
110 | /**
111 |  * A set of services for use within Convex mutation functions.
112 |  *
113 |  * The mutation context is passed as the first argument to any Convex mutation
114 |  * function run on the server.
115 |  */
116 | export type MutationCtx = GenericMutationCtx<DataModel>;
117 | 
118 | /**
119 |  * A set of services for use within Convex action functions.
120 |  *
121 |  * The action context is passed as the first argument to any Convex action
122 |  * function run on the server.
123 |  */
124 | export type ActionCtx = GenericActionCtx<DataModel>;
125 | 
126 | /**
127 |  * An interface to read from the database within Convex query functions.
128 |  *
129 |  * The two entry points are {@link DatabaseReader.get}, which fetches a single
130 |  * document by its {@link Id}, or {@link DatabaseReader.query}, which starts
131 |  * building a query.
132 |  */
133 | export type DatabaseReader = GenericDatabaseReader<DataModel>;
134 | 
135 | /**
136 |  * An interface to read from and write to the database within Convex mutation
137 |  * functions.
138 |  *
139 |  * Convex guarantees that all writes within a single mutation are
140 |  * executed atomically, so you never have to worry about partial writes leaving
141 |  * your data in an inconsistent state. See [the Convex Guide](https://docs.convex.dev/understanding/convex-fundamentals/functions#atomicity-and-optimistic-concurrency-control)
142 |  * for the guarantees Convex provides your functions.
143 |  */
144 | export type DatabaseWriter = GenericDatabaseWriter<DataModel>;
145 | 
146 | /* prettier-ignore-end */
147 | 


--------------------------------------------------------------------------------
/src/Search.tsx:
--------------------------------------------------------------------------------
  1 | import {
  2 |   useAction,
  3 |   useMutation,
  4 |   usePaginatedQuery,
  5 |   useQuery,
  6 | } from "convex/react";
  7 | import { useRef, useState } from "react";
  8 | import { Dispatch } from "react";
  9 | import { MagnifyingGlass } from "@phosphor-icons/react";
 10 | import { Accordion, Button, Table, Text, InputGroup } from "@rewind-ui/core";
 11 | import { Id } from "../convex/_generated/dataModel";
 12 | import { api } from "../convex/_generated/api";
 13 | import { Loading } from "./Loading";
 14 | import { Chunks } from "./Chunks";
 15 | import { CompareFn } from "./useComparison";
 16 | 
 17 | type Target = { text: string; searchId: Id<"searches"> };
 18 | 
 19 | export function Search({ compare }: { compare?: CompareFn }) {
 20 |   const [target, setTarget] = useState<Target>();
 21 |   const [input, setInput] = useState("");
 22 | 
 23 |   return (
 24 |     <>
 25 |       <SearchBar setInput={setInput} setTarget={setTarget} />
 26 |       {target && (
 27 |         <Accordion defaultItem="semantic-search">
 28 |           <Accordion.Item anchor="semantic-search">
 29 |             <Accordion.Header>Semantic</Accordion.Header>
 30 |             <Accordion.Body>
 31 |               <SemanticSearch target={target} compare={compare} />
 32 |             </Accordion.Body>
 33 |           </Accordion.Item>
 34 |           <Accordion.Item anchor="word-search">
 35 |             <Accordion.Header>Word-Based</Accordion.Header>
 36 |             <Accordion.Body>
 37 |               <WordSearch input={input} compare={compare} />
 38 |             </Accordion.Body>
 39 |           </Accordion.Item>
 40 |         </Accordion>
 41 |       )}
 42 |       <Text size="lg">Search History</Text>
 43 |       <PreviousSearches
 44 |         reuseSearch={(target) => {
 45 |           setTarget(target);
 46 |           setInput(target.text);
 47 |         }}
 48 |       />
 49 |     </>
 50 |   );
 51 | }
 52 | 
 53 | function SearchBar({
 54 |   setInput,
 55 |   setTarget,
 56 | }: {
 57 |   setInput: Dispatch<string>;
 58 |   setTarget: Dispatch<Target>;
 59 | }) {
 60 |   const addSearch = useMutation(api.searches.upsert);
 61 |   const searchRef = useRef<HTMLInputElement>(null);
 62 |   return (
 63 |     <form
 64 |       className="flex"
 65 |       onSubmit={(e) => {
 66 |         e.preventDefault();
 67 |         const input = searchRef.current?.value;
 68 |         if (!input) return;
 69 |         setInput(input);
 70 |         addSearch({
 71 |           input,
 72 |         }).then((searchId) => setTarget({ text: input, searchId }));
 73 |       }}
 74 |     >
 75 |       <InputGroup className="w-full">
 76 |         <InputGroup.Text>
 77 |           <MagnifyingGlass weight="duotone" />
 78 |         </InputGroup.Text>
 79 |         <InputGroup.Input
 80 |           ref={searchRef}
 81 |           placeholder={"Search..."}
 82 |           type="search"
 83 |           withRing={false}
 84 |         />
 85 |         <InputGroup.Button type="submit" withRing={false}>
 86 |           Submit
 87 |         </InputGroup.Button>
 88 |       </InputGroup>
 89 |     </form>
 90 |   );
 91 | }
 92 | 
 93 | function WordSearch({
 94 |   input,
 95 |   compare,
 96 | }: {
 97 |   input: string;
 98 |   compare?: CompareFn;
 99 | }) {
100 |   const wordBased = useQuery(api.searches.wordSearch, { input, count: 10 });
101 |   return wordBased ? (
102 |     <>
103 |       <Text>Results for "{input}":</Text>
104 |       <Chunks chunks={wordBased} compare={compare} />
105 |     </>
106 |   ) : (
107 |     <>
108 |       <Text>Searching for "{input}"...</Text>
109 |       <Loading />
110 |     </>
111 |   );
112 | }
113 | 
114 | function SemanticSearch({
115 |   target: { searchId, text },
116 |   compare,
117 | }: {
118 |   target: Target;
119 |   compare?: CompareFn;
120 | }) {
121 |   const semantic = useQuery(api.searches.semanticSearch, { searchId });
122 |   return (
123 |     <>
124 |       {semantic ? (
125 |         <>
126 |           <Text>Results for "{text}":</Text>
127 |           <Chunks chunks={semantic} compare={compare} />
128 |         </>
129 |       ) : (
130 |         <>
131 |           <Text>Searching for "{text}"...</Text>
132 |           <Loading />
133 |         </>
134 |       )}
135 |     </>
136 |   );
137 | }
138 | 
139 | export type UseSearchFn = (target: Target) => void;
140 | 
141 | export function PreviousSearches({
142 |   reuseSearch,
143 | }: {
144 |   reuseSearch: UseSearchFn;
145 | }) {
146 |   const { status, loadMore, results } = usePaginatedQuery(
147 |     api.searches.paginate,
148 |     {},
149 |     { initialNumItems: 10 },
150 |   );
151 |   const search = useAction(api.searches.search);
152 |   return (
153 |     <>
154 |       <Table>
155 |         <Table.Thead>
156 |           <Table.Tr>
157 |             <Table.Th align="left">Input</Table.Th>
158 |             <Table.Th align="left">Tokens</Table.Th>
159 |             <Table.Th align="left">Results</Table.Th>
160 |             <Table.Th align="left">Embedding</Table.Th>
161 |             <Table.Th align="left">Search</Table.Th>
162 |             <Table.Th align="left"></Table.Th>
163 |           </Table.Tr>
164 |         </Table.Thead>
165 |         <Table.Tbody>
166 |           {results.map((result) => (
167 |             <Table.Tr key={result._id}>
168 |               <Table.Td>{result.input}</Table.Td>
169 |               <Table.Td>{result.inputTokens}</Table.Td>
170 |               <Table.Td>
171 |                 {result.relatedChunks?.length || result.count}
172 |               </Table.Td>
173 |               <Table.Td>{result.embeddingMs?.toFixed(0) + " ms"}</Table.Td>
174 |               <Table.Td>{result.queryMs?.toFixed(0) + " ms"}</Table.Td>
175 |               <Table.Td>
176 |                 <Button
177 |                   onClick={() =>
178 |                     reuseSearch({ text: result.input, searchId: result._id })
179 |                   }
180 |                   color="green"
181 |                 >
182 |                   Re-use
183 |                 </Button>
184 |                 <span className="mx-1"></span>
185 |                 <Button
186 |                   onClick={() =>
187 |                     search({
188 |                       input: result.input,
189 |                       topK: result.count,
190 |                       searchId: result._id,
191 |                     }).then(() =>
192 |                       reuseSearch({ text: result.input, searchId: result._id }),
193 |                     )
194 |                   }
195 |                   color="red"
196 |                 >
197 |                   Re-run
198 |                 </Button>
199 |               </Table.Td>
200 |             </Table.Tr>
201 |           ))}
202 |         </Table.Tbody>
203 |       </Table>
204 | 
205 |       {status !== "Exhausted" && (
206 |         <button
207 |           onClick={() => loadMore(10)}
208 |           disabled={status !== "CanLoadMore"}
209 |         >
210 |           Load More
211 |         </button>
212 |       )}
213 |     </>
214 |   );
215 | }
216 | 


--------------------------------------------------------------------------------
/convex/sources.ts:
--------------------------------------------------------------------------------
  1 | import { paginationOptsValidator } from "convex/server";
  2 | import { internal } from "./_generated/api";
  3 | import {
  4 |   DatabaseWriter,
  5 |   internalAction,
  6 |   internalMutation,
  7 |   mutation,
  8 |   query,
  9 | } from "./_generated/server";
 10 | import { fetchEmbeddingBatch } from "./lib/embeddings";
 11 | import { v, Infer } from "convex/values";
 12 | import { Doc } from "./_generated/dataModel";
 13 | import { crud } from "convex-helpers/server/crud";
 14 | import schema from "./schema";
 15 | import { getOrThrow } from "convex-helpers/server/relationships";
 16 | 
 17 | const InputChunk = v.object({
 18 |   text: v.string(),
 19 |   lines: v.object({ from: v.number(), to: v.number() }),
 20 | });
 21 | type InputChunk = Infer<typeof InputChunk>;
 22 | 
 23 | async function addSource(
 24 |   db: DatabaseWriter,
 25 |   name: string,
 26 |   chunks: InputChunk[]
 27 | ) {
 28 |   const sourceId = await db.insert("sources", {
 29 |     name,
 30 |     chunkIds: [],
 31 |     saved: false,
 32 |   });
 33 |   const chunkIds = await Promise.all(
 34 |     chunks.map(({ text, lines }, chunkIndex) =>
 35 |       db.insert("chunks", {
 36 |         text,
 37 |         sourceId,
 38 |         chunkIndex,
 39 |         lines,
 40 |       })
 41 |     )
 42 |   );
 43 |   await db.patch(sourceId, { chunkIds });
 44 |   return (await db.get(sourceId))!;
 45 | }
 46 | 
 47 | // Insert the source into the DB, along with the associated chunks.
 48 | export const add = mutation({
 49 |   args: { name: v.string(), chunks: v.array(InputChunk) },
 50 |   handler: async (ctx, { name, chunks }) => {
 51 |     const source = await addSource(ctx.db, name, chunks);
 52 |     await ctx.scheduler.runAfter(0, internal.sources.addEmbedding, {
 53 |       source,
 54 |       texts: chunks.map(({ text }) => text),
 55 |     });
 56 |   },
 57 | });
 58 | 
 59 | // Make embeddings for a source's chunks and store them.
 60 | export const addEmbedding = internalAction({
 61 |   handler: async (
 62 |     ctx,
 63 |     { source, texts }: { source: Doc<"sources">; texts: string[] }
 64 |   ) => {
 65 |     const { embeddings, embeddingMs, totalTokens } =
 66 |       await fetchEmbeddingBatch(texts);
 67 |     console.log({
 68 |       batchSize: texts.length,
 69 |       totalTokens,
 70 |       embeddingMs,
 71 |     });
 72 |     await ctx.runMutation(internal.sources.storeEmbeddings, {
 73 |       embeddings: embeddings.map((embedding, idx) => ({
 74 |         chunkId: source.chunkIds[idx],
 75 |         embedding,
 76 |       })),
 77 |     });
 78 |     await ctx.runMutation(internal.sources.patch, {
 79 |       id: source._id,
 80 |       patch: { saved: true, totalTokens, embeddingMs },
 81 |     });
 82 |   },
 83 | });
 84 | 
 85 | /**
 86 |  * Add a batch of sources, where each one is a named source with all chunks.
 87 |  */
 88 | export const addBatch = mutation({
 89 |   args: {
 90 |     batch: v.array(v.object({ name: v.string(), chunks: v.array(InputChunk) })),
 91 |   },
 92 |   handler: async (ctx, { batch }) => {
 93 |     await ctx.scheduler.runAfter(0, internal.sources.addEmbeddingBatch, {
 94 |       batch: await Promise.all(
 95 |         batch.map(async ({ name, chunks }) => ({
 96 |           source: await addSource(ctx.db, name, chunks),
 97 |           texts: chunks.map(({ text }) => text),
 98 |         }))
 99 |       ),
100 |     });
101 |   },
102 | });
103 | 
104 | export function chunk<T>(items: T[], chunkSize?: number): T[][] {
105 |   const chunks = [];
106 |   const size = chunkSize || 100;
107 |   for (let i = 0; i < items.length; i += size) {
108 |     chunks.push(items.slice(i, i + size));
109 |   }
110 |   return chunks;
111 | }
112 | 
113 | export const addEmbeddingBatch = internalAction(
114 |   async (
115 |     ctx,
116 |     { batch }: { batch: { source: Doc<"sources">; texts: string[] }[] }
117 |   ) => {
118 |     const chunks = chunk(
119 |       batch.flatMap(({ texts, source }) =>
120 |         texts.map((text, i) => ({ text, chunkId: source.chunkIds[i] }))
121 |       )
122 |     );
123 |     let totalTokens_ = 0;
124 |     let embeddingMs_ = 0;
125 |     for (const chunkBatch of chunks) {
126 |       // Calculate all the embeddings for all sources at once.
127 |       const { embeddings, totalTokens, embeddingMs } =
128 |         await fetchEmbeddingBatch(chunkBatch.map(({ text }) => text));
129 |       totalTokens_ += totalTokens;
130 |       embeddingMs_ += embeddingMs;
131 |       console.log({ batchSize: embeddings.length, totalTokens, embeddingMs });
132 |       await ctx.runMutation(internal.sources.storeEmbeddings, {
133 |         embeddings: embeddings.map((embedding, i) => ({
134 |           chunkId: chunkBatch[i].chunkId,
135 |           embedding,
136 |         })),
137 |       });
138 |     }
139 |     // The length of all strings put together.
140 |     const totalLength = batch.reduce(
141 |       (sum, { texts }) => sum + textLength(texts),
142 |       0
143 |     );
144 |     await Promise.all(
145 |       batch.map(async ({ source, texts }) => {
146 |         const sourceLength = textLength(texts);
147 |         const portion = sourceLength / totalLength;
148 |         await ctx.runMutation(internal.sources.patch, {
149 |           id: source._id,
150 |           patch: {
151 |             saved: true,
152 |             embeddingMs: Math.ceil(embeddingMs_ * portion),
153 |             totalTokens: Math.ceil(totalTokens_ * portion),
154 |           },
155 |         });
156 |       })
157 |     );
158 |   }
159 | );
160 | 
161 | export const storeEmbeddings = internalMutation({
162 |   args: {
163 |     embeddings: v.array(
164 |       v.object({ chunkId: v.id("chunks"), embedding: v.array(v.number()) })
165 |     ),
166 |   },
167 |   handler: async (ctx, args) => {
168 |     await Promise.all(
169 |       args.embeddings.map(async ({ chunkId, embedding }) => {
170 |         const chunk = await getOrThrow(ctx, chunkId);
171 |         if (chunk.embeddingId) {
172 |           await ctx.db.patch(chunk.embeddingId, { vector: embedding });
173 |         } else {
174 |           const embeddingId = await ctx.db.insert("chunkEmbeddings", {
175 |             vector: embedding,
176 |           });
177 |           await ctx.db.patch(chunkId, { embeddingId });
178 |         }
179 |       })
180 |     );
181 |   },
182 | });
183 | 
184 | function textLength(texts: string[]) {
185 |   return texts.reduce((sum, cur) => sum + cur.length, 0);
186 | }
187 | 
188 | export const { update: patch } = crud(schema, "sources");
189 | 
190 | export const paginate = query({
191 |   args: { paginationOpts: paginationOptsValidator },
192 |   handler: async (ctx, { paginationOpts }) => {
193 |     const results = await ctx.db
194 |       .query("sources")
195 |       .order("desc")
196 |       .paginate(paginationOpts);
197 | 
198 |     return {
199 |       ...results,
200 |       page: await Promise.all(
201 |         results.page.map(async (source) => {
202 |           let firstChunkText = "";
203 |           if (source.chunkIds.length) {
204 |             const firstChunk = await ctx.db.get(source.chunkIds[0]);
205 |             firstChunkText = firstChunk?.text ?? "";
206 |           }
207 |           return { ...source, firstChunkText };
208 |         })
209 |       ),
210 |     };
211 |   },
212 | });
213 | 
214 | export const deleteSource = mutation({
215 |   args: { id: v.id("sources") },
216 |   handler: async (ctx, { id }) => {
217 |     const source = await ctx.db.get(id);
218 |     if (!source) return;
219 |     await ctx.db.delete(id);
220 |     await Promise.all(
221 |       source.chunkIds.map(async (id) => {
222 |         const chunk = await ctx.db.get(id);
223 |         if (!chunk) return;
224 |         ctx.db.delete(id);
225 |         if (chunk.embeddingId) {
226 |           ctx.db.delete(chunk.embeddingId);
227 |         }
228 |       })
229 |     );
230 |   },
231 | });
232 | 
233 | export const paginateChunks = query({
234 |   args: { paginationOpts: paginationOptsValidator },
235 |   handler: async (ctx, { paginationOpts }) => {
236 |     const results = await ctx.db.query("chunks").paginate(paginationOpts);
237 | 
238 |     return {
239 |       ...results,
240 |       page: await Promise.all(
241 |         results.page.map(async (chunk) => {
242 |           const source = await ctx.db.get(chunk.sourceId);
243 |           return { ...chunk, sourceName: source!.name };
244 |         })
245 |       ),
246 |     };
247 |   },
248 | });
249 | 
250 | export const getChunk = query({
251 |   args: { id: v.id("chunks") },
252 |   handler: async (ctx, { id }) => {
253 |     const doc = await ctx.db.get(id);
254 |     if (!doc) {
255 |       throw new Error("Document not found: " + id);
256 |     }
257 |     return doc;
258 |   },
259 | });
260 | 


--------------------------------------------------------------------------------