From af5c4a43d8898fb760a2119c6272bbd29dd7aeb4 Mon Sep 17 00:00:00 2001
From: Julian Benegas <julianbenegas99@gmail.com>
Date: Thu, 14 Nov 2024 13:38:38 -0300
Subject: [PATCH] ask ai

---
 packages/basehub/src/search/index.ts      |  15 ++-
 packages/basehub/src/search/primitive.tsx | 134 ++++++++++++++++++++++
 2 files changed, 147 insertions(+), 2 deletions(-)

diff --git a/packages/basehub/src/search/index.ts b/packages/basehub/src/search/index.ts
index 288d49c5..127e23ac 100644
--- a/packages/basehub/src/search/index.ts
+++ b/packages/basehub/src/search/index.ts
@@ -1,2 +1,13 @@
-export { getSearchClient as getSearchClientRaw, search } from "./primitive";
-export type { Highlight, Hit, SearchOptions, SearchResult } from "./primitive";
+export {
+  getSearchClient as getSearchClientRaw,
+  search,
+  askAi,
+} from "./primitive";
+export type {
+  Highlight,
+  Hit,
+  SearchOptions,
+  SearchResult,
+  AskAiResult,
+  AskAiOptions,
+} from "./primitive";
diff --git a/packages/basehub/src/search/primitive.tsx b/packages/basehub/src/search/primitive.tsx
index 9f189d97..f553183e 100644
--- a/packages/basehub/src/search/primitive.tsx
+++ b/packages/basehub/src/search/primitive.tsx
@@ -191,6 +191,140 @@ export const search = async <Document extends Record<string, unknown>>(
   return newResult;
 };
 
+/* -------------------------------------------------------------------------------------------------
+ * Ask Ai
+ * -----------------------------------------------------------------------------------------------*/
+
+export type AskAiOptions = {
+  conversationId?: string;
+};
+
+export type AskAiResult<Doc = Record<string, unknown>> = SearchResult<Doc> & {
+  conversation: {
+    id: string;
+    answer: string;
+    history: {
+      question: string;
+      answer: string;
+    }[];
+  };
+};
+
+export const askAi = async <Document extends Record<string, unknown>>(
+  _searchKey: string | null,
+  q: string,
+  opts: SearchOptions
+): Promise<SearchResult<Document>> => {
+  if (!_searchKey) throw new Error("Not enabled");
+  const { collectionName, valid } = decodeKey(_searchKey);
+  if (!valid) throw new Error("Invalid _searchKey");
+  const client = getSearchClient(_searchKey);
+  if (!client) throw new Error("Couldn't get search client");
+
+  const options: Record<string, unknown> = {
+    q,
+    conversation: true,
+    conversation_model_id: "conv-model-1",
+    queryBy: ["__embedding"],
+    excludeFields: ["__embedding"],
+  };
+  Object.entries(opts).forEach(([key, value]) => {
+    options[camelToSnake(key)] = value;
+  });
+
+  const rawResult = await client
+    .collections(collectionName)
+    .documents()
+    .search(options);
+
+  if (!rawResult.conversation) {
+    throw new Error("Ask AI failed: conversation wasn't found in response.");
+  }
+
+  const rawConversationHistory = (rawResult.conversation
+    .conversation_history as unknown as undefined) ?? [
+    { user: q },
+    { assistant: rawResult.conversation.answer },
+  ];
+
+  const newResult: AskAiResult<Document> = {
+    empty: !rawResult.found,
+    found: rawResult.found,
+    outOf: rawResult.out_of,
+    page: rawResult.page,
+    searchTimeMs: rawResult.search_time_ms,
+    hits:
+      rawResult.hits?.map((hit) => {
+        const document = deFlatten(hit.document) as Document & BaseDoc;
+        const highlightRecord = {} as Record<string, Highlight>;
+        const highlights =
+          hit.highlights?.map((highlight) => {
+            const fieldPath = highlight.field as string;
+
+            const cast: Highlight = {
+              fieldPath,
+              fieldValue: get(document, fieldPath) as unknown,
+              indices: highlight.indices ?? [],
+              matchedTokens: highlight.matched_tokens,
+              snippet: highlight.snippet,
+              snippets: highlight.snippets ?? [],
+              value: highlight.value,
+            };
+
+            highlightRecord[highlight.field as string] = cast;
+
+            return cast;
+          }) ?? [];
+
+        const _key = getHitKey(hit);
+
+        const _getField = (fieldPath: string) => {
+          return get(document, fieldPath) as unknown;
+        };
+
+        const fullHit: Hit<Document> = {
+          _key,
+          curated: hit.curated ?? false,
+          document,
+          highlight: highlightRecord,
+          highlights,
+          textMatch: hit.text_match,
+          _getField,
+          _getFieldHighlight: () => null,
+        };
+
+        fullHit._getFieldHighlight = (
+          fieldPath: string,
+          fallbackFieldPaths?: string[]
+        ) => {
+          return _getFieldHighlightImpl({
+            fieldPath,
+            fallbackFieldPaths,
+            includeFallback: true,
+            hit: fullHit,
+          });
+        };
+
+        return fullHit;
+      }) ?? [],
+    conversation: {
+      answer: rawResult.conversation.answer,
+      id: rawResult.conversation.conversation_id,
+      history: rawConversationHistory
+        .map((h, i) => {
+          const question = h.user;
+          if (!question) return null;
+          const answer = rawConversationHistory[i + 1]?.assistant;
+          if (!answer) return null;
+          return { question, answer };
+        })
+        .filter((h): h is { question: string; answer: string } => h !== null),
+    },
+  };
+
+  return newResult;
+};
+
 /* -------------------------------------------------------------------------------------------------
  * Utils
  * -----------------------------------------------------------------------------------------------*/