Customizable personas (#772)

Also includes a small fix to LLM filtering when combined with reranking
2025-09-19 20:24:32 +02:00 · 2023-11-28 00:57:48 -08:00
parent 87beb1f4d1
commit 78d1ae0379
49 changed files with 1846 additions and 408 deletions
--- a/web/src/app/admin/personas/PersonaEditor.tsx
+++ b/web/src/app/admin/personas/PersonaEditor.tsx
@@ -0,0 +1,411 @@
+"use client";
+
+import {
+  BooleanFormField,
+  TextArrayField,
+  TextFormField,
+} from "@/components/admin/connectors/Field";
+import { DocumentSet } from "@/lib/types";
+import { Button, Divider, Text, Title } from "@tremor/react";
+import {
+  ArrayHelpers,
+  ErrorMessage,
+  Field,
+  FieldArray,
+  Form,
+  Formik,
+} from "formik";
+
+import * as Yup from "yup";
+import { buildFinalPrompt, createPersona, updatePersona } from "./lib";
+import { useRouter } from "next/navigation";
+import { usePopup } from "@/components/admin/connectors/Popup";
+import { Persona } from "./interfaces";
+import Link from "next/link";
+import { useEffect, useState } from "react";
+
+function SectionHeader({ children }: { children: string | JSX.Element }) {
+  return <div className="mb-4 font-bold text-lg">{children}</div>;
+}
+
+function Label({ children }: { children: string | JSX.Element }) {
+  return (
+    <div className="block font-medium text-base text-gray-200">{children}</div>
+  );
+}
+
+function SubLabel({ children }: { children: string | JSX.Element }) {
+  return <div className="text-sm text-gray-300 mb-2">{children}</div>;
+}
+
+// TODO: make this the default text input across all forms
+function PersonaTextInput({
+  name,
+  label,
+  subtext,
+  placeholder,
+  onChange,
+  type = "text",
+  isTextArea = false,
+  disabled = false,
+  autoCompleteDisabled = true,
+}: {
+  name: string;
+  label: string;
+  subtext?: string | JSX.Element;
+  placeholder?: string;
+  onChange?: (e: React.ChangeEvent<HTMLInputElement>) => void;
+  type?: string;
+  isTextArea?: boolean;
+  disabled?: boolean;
+  autoCompleteDisabled?: boolean;
+}) {
+  return (
+    <div className="mb-4">
+      <Label>{label}</Label>
+      {subtext && <SubLabel>{subtext}</SubLabel>}
+      <Field
+        as={isTextArea ? "textarea" : "input"}
+        type={type}
+        name={name}
+        id={name}
+        className={
+          `
+        border 
+        text-gray-200 
+        border-gray-600 
+        rounded 
+        w-full 
+        py-2 
+        px-3 
+        mt-1
+        ${isTextArea ? " h-28" : ""}
+      ` + (disabled ? " bg-gray-900" : " bg-gray-800")
+        }
+        disabled={disabled}
+        placeholder={placeholder}
+        autoComplete={autoCompleteDisabled ? "off" : undefined}
+        {...(onChange ? { onChange } : {})}
+      />
+      <ErrorMessage
+        name={name}
+        component="div"
+        className="text-red-500 text-sm mt-1"
+      />
+    </div>
+  );
+}
+
+function PersonaBooleanInput({
+  name,
+  label,
+  subtext,
+}: {
+  name: string;
+  label: string;
+  subtext?: string | JSX.Element;
+}) {
+  return (
+    <div className="mb-4">
+      <Label>{label}</Label>
+      {subtext && <SubLabel>{subtext}</SubLabel>}
+      <Field
+        type="checkbox"
+        name={name}
+        id={name}
+        className={`
+        ml-2
+        border 
+        text-gray-200 
+        border-gray-600 
+        rounded 
+        py-2 
+        px-3 
+        mt-1
+      `}
+      />
+      <ErrorMessage
+        name={name}
+        component="div"
+        className="text-red-500 text-sm mt-1"
+      />
+    </div>
+  );
+}
+
+export function PersonaEditor({
+  existingPersona,
+  documentSets,
+}: {
+  existingPersona?: Persona | null;
+  documentSets: DocumentSet[];
+}) {
+  const router = useRouter();
+  const { popup, setPopup } = usePopup();
+
+  const [finalPrompt, setFinalPrompt] = useState<string | null>("");
+
+  const triggerFinalPromptUpdate = async (
+    systemPrompt: string,
+    taskPrompt: string
+  ) => {
+    const response = await buildFinalPrompt(systemPrompt, taskPrompt);
+    if (response.ok) {
+      setFinalPrompt((await response.json()).final_prompt_template);
+    }
+  };
+
+  const isUpdate = existingPersona !== undefined && existingPersona !== null;
+
+  useEffect(() => {
+    if (isUpdate) {
+      triggerFinalPromptUpdate(
+        existingPersona.system_prompt,
+        existingPersona.task_prompt
+      );
+    }
+  }, []);
+
+  return (
+    <div className="dark">
+      {popup}
+      <Formik
+        initialValues={{
+          name: existingPersona?.name ?? "",
+          description: existingPersona?.description ?? "",
+          system_prompt: existingPersona?.system_prompt ?? "",
+          task_prompt: existingPersona?.task_prompt ?? "",
+          document_set_ids:
+            existingPersona?.document_sets?.map(
+              (documentSet) => documentSet.id
+            ) ?? ([] as number[]),
+          num_chunks: existingPersona?.num_chunks ?? null,
+          apply_llm_relevance_filter:
+            existingPersona?.apply_llm_relevance_filter ?? false,
+        }}
+        validationSchema={Yup.object().shape({
+          name: Yup.string().required("Must give the Persona a name!"),
+          description: Yup.string().required(
+            "Must give the Persona a description!"
+          ),
+          system_prompt: Yup.string().required(
+            "Must give the Persona a system prompt!"
+          ),
+          task_prompt: Yup.string().required(
+            "Must give the Persona a task prompt!"
+          ),
+          document_set_ids: Yup.array().of(Yup.number()),
+          num_chunks: Yup.number().max(20).nullable(),
+          apply_llm_relevance_filter: Yup.boolean().required(),
+        })}
+        onSubmit={async (values, formikHelpers) => {
+          formikHelpers.setSubmitting(true);
+
+          let response;
+          if (isUpdate) {
+            response = await updatePersona({
+              id: existingPersona.id,
+              ...values,
+              num_chunks: values.num_chunks || null,
+            });
+          } else {
+            response = await createPersona({
+              ...values,
+              num_chunks: values.num_chunks || null,
+            });
+          }
+          if (response.ok) {
+            router.push("/admin/personas");
+            return;
+          }
+
+          setPopup({
+            type: "error",
+            message: `Failed to create Persona - ${await response.text()}`,
+          });
+          formikHelpers.setSubmitting(false);
+        }}
+      >
+        {({ isSubmitting, values, setFieldValue }) => (
+          <Form>
+            <div className="pb-6">
+              <SectionHeader>Who am I?</SectionHeader>
+
+              <PersonaTextInput
+                name="name"
+                label="Name"
+                disabled={isUpdate}
+                subtext="Users will be able to select this Persona based on this name."
+              />
+
+              <PersonaTextInput
+                name="description"
+                label="Description"
+                subtext="Provide a short descriptions which gives users a hint as to what they should use this Persona for."
+              />
+
+              <Divider />
+
+              <SectionHeader>Customize my response style</SectionHeader>
+
+              <PersonaTextInput
+                name="system_prompt"
+                label="System Prompt"
+                isTextArea={true}
+                subtext={
+                  'Give general info about what the Persona is about. For example, "You are an assistant for On-Call engineers. Your goal is to read the provided context documents and give recommendations as to how to resolve the issue."'
+                }
+                onChange={(e) => {
+                  setFieldValue("system_prompt", e.target.value);
+                  triggerFinalPromptUpdate(e.target.value, values.task_prompt);
+                }}
+              />
+
+              <PersonaTextInput
+                name="task_prompt"
+                label="Task Prompt"
+                isTextArea={true}
+                subtext={
+                  'Give specific instructions as to what to do with the user query. For example, "Find any relevant sections from the provided documents that can help the user resolve their issue and explain how they are relevant."'
+                }
+                onChange={(e) => {
+                  setFieldValue("task_prompt", e.target.value);
+                  triggerFinalPromptUpdate(
+                    values.system_prompt,
+                    e.target.value
+                  );
+                }}
+              />
+
+              <Label>Final Prompt</Label>
+
+              {finalPrompt ? (
+                <pre className="text-sm mt-2 whitespace-pre-wrap">
+                  {finalPrompt.replaceAll("\\n", "\n")}
+                </pre>
+              ) : (
+                "-"
+              )}
+
+              <Divider />
+
+              <SectionHeader>What data should I have access to?</SectionHeader>
+
+              <FieldArray
+                name="document_set_ids"
+                render={(arrayHelpers: ArrayHelpers) => (
+                  <div>
+                    <div>
+                      <SubLabel>
+                        <>
+                          Select which{" "}
+                          <Link
+                            href="/admin/documents/sets"
+                            className="text-blue-500"
+                            target="_blank"
+                          >
+                            Document Sets
+                          </Link>{" "}
+                          that this Persona should search through. If none are
+                          specified, the Persona will search through all
+                          available documents in order to try and response to
+                          queries.
+                        </>
+                      </SubLabel>
+                    </div>
+                    <div className="mb-3 mt-2 flex gap-2 flex-wrap text-sm">
+                      {documentSets.map((documentSet) => {
+                        const ind = values.document_set_ids.indexOf(
+                          documentSet.id
+                        );
+                        let isSelected = ind !== -1;
+                        return (
+                          <div
+                            key={documentSet.id}
+                            className={
+                              `
+                              px-3 
+                              py-1
+                              rounded-lg 
+                              border
+                              border-gray-700 
+                              w-fit 
+                              flex 
+                              cursor-pointer ` +
+                              (isSelected
+                                ? " bg-gray-600"
+                                : " bg-gray-900 hover:bg-gray-700")
+                            }
+                            onClick={() => {
+                              if (isSelected) {
+                                arrayHelpers.remove(ind);
+                              } else {
+                                arrayHelpers.push(documentSet.id);
+                              }
+                            }}
+                          >
+                            <div className="my-auto">{documentSet.name}</div>
+                          </div>
+                        );
+                      })}
+                    </div>
+                  </div>
+                )}
+              />
+
+              <Divider />
+
+              <SectionHeader>[Advanced] Retrieval Customization</SectionHeader>
+
+              <PersonaTextInput
+                name="num_chunks"
+                label="Number of Chunks"
+                subtext={
+                  <div>
+                    How many chunks should we feed into the LLM when generating
+                    the final response? Each chunk is ~400 words long. If you
+                    are using gpt-3.5-turbo or other similar models, setting
+                    this to a value greater than 5 will result in errors at
+                    query time due to the model&apos;s input length limit.
+                    <br />
+                    <br />
+                    If unspecified, will use 5 chunks.
+                  </div>
+                }
+                onChange={(e) => {
+                  const value = e.target.value;
+                  // Allow only integer values
+                  if (value === "" || /^[0-9]+$/.test(value)) {
+                    setFieldValue("num_chunks", value);
+                  }
+                }}
+              />
+
+              <PersonaBooleanInput
+                name="apply_llm_relevance_filter"
+                label="Apply LLM Relevance Filter"
+                subtext={
+                  "If enabled, the LLM will filter out chunks that are not relevant to the user query."
+                }
+              />
+
+              <Divider />
+
+              <div className="flex">
+                <Button
+                  className="mx-auto"
+                  variant="secondary"
+                  size="md"
+                  type="submit"
+                  disabled={isSubmitting}
+                >
+                  {isUpdate ? "Update!" : "Create!"}
+                </Button>
+              </div>
+            </div>
+          </Form>
+        )}
+      </Formik>
+    </div>
+  );
+}
--- a/web/src/app/admin/personas/PersonaTable.tsx
+++ b/web/src/app/admin/personas/PersonaTable.tsx
@@ -0,0 +1,52 @@
+"use client";
+
+import {
+  Table,
+  TableHead,
+  TableRow,
+  TableHeaderCell,
+  TableBody,
+  TableCell,
+} from "@tremor/react";
+import { Persona } from "./interfaces";
+import Link from "next/link";
+import { EditButton } from "@/components/EditButton";
+import { useRouter } from "next/navigation";
+
+export function PersonasTable({ personas }: { personas: Persona[] }) {
+  const router = useRouter();
+
+  const sortedPersonas = [...personas];
+  sortedPersonas.sort((a, b) => a.name.localeCompare(b.name));
+
+  return (
+    <div className="dark">
+      <Table className="overflow-visible">
+        <TableHead>
+          <TableRow>
+            <TableHeaderCell>Name</TableHeaderCell>
+            <TableHeaderCell>Description</TableHeaderCell>
+            <TableHeaderCell></TableHeaderCell>
+          </TableRow>
+        </TableHead>
+        <TableBody>
+          {sortedPersonas.map((persona) => {
+            return (
+              <TableRow key={persona.id}>
+                <TableCell className="whitespace-normal break-all">
+                  <p className="text font-medium">{persona.name}</p>
+                </TableCell>
+                <TableCell>{persona.description}</TableCell>
+                <TableCell>
+                  <EditButton
+                    onClick={() => router.push(`/admin/personas/${persona.id}`)}
+                  />
+                </TableCell>
+              </TableRow>
+            );
+          })}
+        </TableBody>
+      </Table>
+    </div>
+  );
+}
--- a/web/src/app/admin/personas/[personaId]/DeletePersonaButton.tsx
+++ b/web/src/app/admin/personas/[personaId]/DeletePersonaButton.tsx
@@ -0,0 +1,29 @@
+"use client";
+
+import { Button } from "@tremor/react";
+import { FiTrash } from "react-icons/fi";
+import { deletePersona } from "../lib";
+import { useRouter } from "next/navigation";
+
+export function DeletePersonaButton({ personaId }: { personaId: number }) {
+  const router = useRouter();
+
+  return (
+    <Button
+      variant="secondary"
+      size="xs"
+      color="red"
+      onClick={async () => {
+        const response = await deletePersona(personaId);
+        if (response.ok) {
+          router.push("/admin/personas");
+        } else {
+          alert(`Failed to delete persona - ${await response.text()}`);
+        }
+      }}
+      icon={FiTrash}
+    >
+      Delete
+    </Button>
+  );
+}
--- a/web/src/app/admin/personas/[personaId]/page.tsx
+++ b/web/src/app/admin/personas/[personaId]/page.tsx
@@ -0,0 +1,62 @@
+import { ErrorCallout } from "@/components/ErrorCallout";
+import { fetchSS } from "@/lib/utilsSS";
+import { FaRobot } from "react-icons/fa";
+import { Persona } from "../interfaces";
+import { PersonaEditor } from "../PersonaEditor";
+import { DocumentSet } from "@/lib/types";
+import { RobotIcon } from "@/components/icons/icons";
+import { BackButton } from "@/components/BackButton";
+import { Card, Title, Text, Divider, Button } from "@tremor/react";
+import { FiTrash } from "react-icons/fi";
+import { DeletePersonaButton } from "./DeletePersonaButton";
+
+export default async function Page({
+  params,
+}: {
+  params: { personaId: string };
+}) {
+  const personaResponse = await fetchSS(`/persona/${params.personaId}`);
+
+  if (!personaResponse.ok) {
+    return (
+      <ErrorCallout
+        errorTitle="Something went wrong :("
+        errorMsg={`Failed to fetch Persona - ${await personaResponse.text()}`}
+      />
+    );
+  }
+
+  const documentSetsResponse = await fetchSS("/manage/document-set");
+
+  if (!documentSetsResponse.ok) {
+    return (
+      <ErrorCallout
+        errorTitle="Something went wrong :("
+        errorMsg={`Failed to fetch document sets - ${await documentSetsResponse.text()}`}
+      />
+    );
+  }
+
+  const documentSets = (await documentSetsResponse.json()) as DocumentSet[];
+  const persona = (await personaResponse.json()) as Persona;
+
+  return (
+    <div className="dark">
+      <BackButton />
+      <div className="pb-2 mb-4 flex">
+        <h1 className="text-3xl font-bold pl-2">Edit Persona</h1>
+      </div>
+
+      <Card>
+        <PersonaEditor existingPersona={persona} documentSets={documentSets} />
+      </Card>
+
+      <div className="mt-12">
+        <Title>Delete Persona</Title>
+        <div className="flex mt-6">
+          <DeletePersonaButton personaId={persona.id} />
+        </div>
+      </div>
+    </div>
+  );
+}
--- a/web/src/app/admin/personas/interfaces.ts
+++ b/web/src/app/admin/personas/interfaces.ts
@@ -0,0 +1,12 @@
+import { DocumentSet } from "@/lib/types";
+
+export interface Persona {
+  id: number;
+  name: string;
+  description: string;
+  system_prompt: string;
+  task_prompt: string;
+  document_sets: DocumentSet[];
+  num_chunks?: number;
+  apply_llm_relevance_filter?: boolean;
+}
--- a/web/src/app/admin/personas/lib.ts
+++ b/web/src/app/admin/personas/lib.ts
@@ -0,0 +1,61 @@
+interface PersonaCreationRequest {
+  name: string;
+  description: string;
+  system_prompt: string;
+  task_prompt: string;
+  document_set_ids: number[];
+  num_chunks: number | null;
+  apply_llm_relevance_filter: boolean | null;
+}
+
+export function createPersona(personaCreationRequest: PersonaCreationRequest) {
+  return fetch("/api/admin/persona", {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+    },
+    body: JSON.stringify(personaCreationRequest),
+  });
+}
+
+interface PersonaUpdateRequest {
+  id: number;
+  description: string;
+  system_prompt: string;
+  task_prompt: string;
+  document_set_ids: number[];
+  num_chunks: number | null;
+  apply_llm_relevance_filter: boolean | null;
+}
+
+export function updatePersona(personaUpdateRequest: PersonaUpdateRequest) {
+  const { id, ...requestBody } = personaUpdateRequest;
+
+  return fetch(`/api/admin/persona/${id}`, {
+    method: "PATCH",
+    headers: {
+      "Content-Type": "application/json",
+    },
+    body: JSON.stringify(requestBody),
+  });
+}
+
+export function deletePersona(personaId: number) {
+  return fetch(`/api/admin/persona/${personaId}`, {
+    method: "DELETE",
+  });
+}
+
+export function buildFinalPrompt(systemPrompt: string, taskPrompt: string) {
+  let queryString = Object.entries({
+    system_prompt: systemPrompt,
+    task_prompt: taskPrompt,
+  })
+    .map(
+      ([key, value]) =>
+        `${encodeURIComponent(key)}=${encodeURIComponent(value)}`
+    )
+    .join("&");
+
+  return fetch(`/api/persona-utils/prompt-explorer?${queryString}`);
+}
--- a/web/src/app/admin/personas/new/page.tsx
+++ b/web/src/app/admin/personas/new/page.tsx
@@ -0,0 +1,37 @@
+import { FaRobot } from "react-icons/fa";
+import { PersonaEditor } from "../PersonaEditor";
+import { fetchSS } from "@/lib/utilsSS";
+import { ErrorCallout } from "@/components/ErrorCallout";
+import { DocumentSet } from "@/lib/types";
+import { RobotIcon } from "@/components/icons/icons";
+import { BackButton } from "@/components/BackButton";
+import { Card } from "@tremor/react";
+
+export default async function Page() {
+  const documentSetsResponse = await fetchSS("/manage/document-set");
+
+  if (!documentSetsResponse.ok) {
+    return (
+      <ErrorCallout
+        errorTitle="Something went wrong :("
+        errorMsg={`Failed to fetch document sets - ${await documentSetsResponse.text()}`}
+      />
+    );
+  }
+
+  const documentSets = (await documentSetsResponse.json()) as DocumentSet[];
+
+  return (
+    <div className="dark">
+      <BackButton />
+      <div className="border-solid border-gray-600 border-b pb-2 mb-4 flex">
+        <RobotIcon size={32} />
+        <h1 className="text-3xl font-bold pl-2">Create a New Persona</h1>
+      </div>
+
+      <Card>
+        <PersonaEditor documentSets={documentSets} />
+      </Card>
+    </div>
+  );
+}
--- a/web/src/app/admin/personas/page.tsx
+++ b/web/src/app/admin/personas/page.tsx
@@ -0,0 +1,64 @@
+import { PersonasTable } from "./PersonaTable";
+import { FiPlusSquare } from "react-icons/fi";
+import Link from "next/link";
+import { Divider, Text, Title } from "@tremor/react";
+import { fetchSS } from "@/lib/utilsSS";
+import { ErrorCallout } from "@/components/ErrorCallout";
+import { Persona } from "./interfaces";
+import { RobotIcon } from "@/components/icons/icons";
+
+export default async function Page() {
+  const personaResponse = await fetchSS("/persona");
+
+  if (!personaResponse.ok) {
+    return (
+      <ErrorCallout
+        errorTitle="Something went wrong :("
+        errorMsg={`Failed to fetch personas - ${await personaResponse.text()}`}
+      />
+    );
+  }
+
+  const personas = (await personaResponse.json()) as Persona[];
+
+  return (
+    <div>
+      <div className="border-solid border-gray-600 border-b pb-2 mb-4 flex">
+        <RobotIcon size={32} />
+        <h1 className="text-3xl font-bold pl-2">Personas</h1>
+      </div>
+
+      <div className="text-gray-300 text-sm mb-2">
+        Personas are a way to build custom search/question-answering experiences
+        for different use cases.
+        <p className="mt-2">They allow you to customize:</p>
+        <ul className="list-disc mt-2 ml-4">
+          <li>
+            The prompt used by your LLM of choice to respond to the user query
+          </li>
+          <li>The documents that are used as context</li>
+        </ul>
+      </div>
+
+      <div className="dark">
+        <Divider />
+
+        <Title>Create a Persona</Title>
+        <Link
+          href="/admin/personas/new"
+          className="text-gray-100 flex py-2 px-4 mt-2 border border-gray-800 h-fit cursor-pointer hover:bg-gray-800 text-sm w-36"
+        >
+          <div className="mx-auto flex">
+            <FiPlusSquare className="my-auto mr-2" />
+            New Persona
+          </div>
+        </Link>
+
+        <Divider />
+
+        <Title>Existing Personas</Title>
+        <PersonasTable personas={personas} />
+      </div>
+    </div>
+  );
+}
--- a/web/src/app/page.tsx
+++ b/web/src/app/page.tsx
@@ -8,10 +8,11 @@ import {
 import { redirect } from "next/navigation";
 import { HealthCheckBanner } from "@/components/health/healthcheck";
 import { ApiKeyModal } from "@/components/openai/ApiKeyModal";
-import { buildUrl } from "@/lib/utilsSS";
+import { buildUrl, fetchSS } from "@/lib/utilsSS";
 import { Connector, DocumentSet, User } from "@/lib/types";
 import { cookies } from "next/headers";
 import { SearchType } from "@/lib/search/interfaces";
+import { Persona } from "./admin/personas/interfaces";

 export default async function Home() {
  const tasks = [
@@ -29,6 +30,7 @@ export default async function Home() {
        cookie: processCookies(cookies()),
      },
    }),
+    fetchSS("/persona"),
  ];

  // catch cases where the backend is completely unreachable here
@@ -44,6 +46,7 @@ export default async function Home() {
  const user = results[1] as User | null;
  const connectorsResponse = results[2] as Response | null;
  const documentSetsResponse = results[3] as Response | null;
+  const personaResponse = results[4] as Response | null;

  if (!authDisabled && !user) {
    return redirect("/auth/login");
@@ -65,6 +68,13 @@ export default async function Home() {
    );
  }

+  let personas: Persona[] = [];
+  if (personaResponse?.ok) {
+    personas = await personaResponse.json();
+  } else {
+    console.log(`Failed to fetch personas - ${personaResponse?.status}`);
+  }
+
  // needs to be done in a non-client side component due to nextjs
  const storedSearchType = cookies().get("searchType")?.value as
    | string
@@ -87,6 +97,7 @@ export default async function Home() {
          <SearchSection
            connectors={connectors}
            documentSets={documentSets}
+            personas={personas}
            defaultSearchType={searchTypeDefault}
          />
        </div>
--- a/web/src/components/Dropdown.tsx
+++ b/web/src/components/Dropdown.tsx
@@ -14,11 +14,7 @@ interface DropdownProps {
  onSelect: (selected: Option) => void;
 }

-export const Dropdown: FC<DropdownProps> = ({
-  options,
-  selected,
-  onSelect,
-}) => {
+export const Dropdown = ({ options, selected, onSelect }: DropdownProps) => {
  const [isOpen, setIsOpen] = useState(false);
  const dropdownRef = useRef<HTMLDivElement>(null);

--- a/web/src/components/EditButton.tsx
+++ b/web/src/components/EditButton.tsx
@@ -0,0 +1,27 @@
+"use client";
+
+import { useRouter } from "next/navigation";
+
+import { FiChevronLeft, FiEdit } from "react-icons/fi";
+
+export function EditButton({ onClick }: { onClick: () => void }) {
+  return (
+    <div
+      className={`
+        my-auto 
+        flex 
+        mb-1 
+        hover:bg-gray-800 
+        w-fit 
+        p-2 
+        cursor-pointer 
+        rounded-lg
+        border-gray-800 
+        text-sm`}
+      onClick={onClick}
+    >
+      <FiEdit className="mr-1 my-auto" />
+      Edit
+    </div>
+  );
+}
--- a/web/src/components/admin/Layout.tsx
+++ b/web/src/components/admin/Layout.tsx
@@ -28,6 +28,7 @@ import {
  GongIcon,
  ZoomInIcon,
  ZendeskIcon,
+  RobotIcon,
 } from "@/components/icons/icons";
 import { getAuthDisabledSS, getCurrentUserSS } from "@/lib/userSS";
 import { redirect } from "next/navigation";
@@ -314,13 +315,22 @@ export async function Layout({ children }: { children: React.ReactNode }) {
              ],
            },
            {
-              name: "Bots",
+              name: "Custom Assistants",
              items: [
+                {
+                  name: (
+                    <div className="flex">
+                      <RobotIcon size={18} />
+                      <div className="ml-1">Personas</div>
+                    </div>
+                  ),
+                  link: "/admin/personas",
+                },
                {
                  name: (
                    <div className="flex">
                      <CPUIcon size={18} />
-                      <div className="ml-1">Slack Bot</div>
+                      <div className="ml-1">Slack Bots</div>
                    </div>
                  ),
                  link: "/admin/bot",
--- a/web/src/components/icons/icons.tsx
+++ b/web/src/components/icons/icons.tsx
@@ -49,6 +49,7 @@ import hubSpotIcon from "../../../public/HubSpot.png";
 import document360Icon from "../../../public/Document360.png";
 import googleSitesIcon from "../../../public/GoogleSites.png";
 import zendeskIcon from "../../../public/Zendesk.svg";
+import { FaRobot } from "react-icons/fa";

 interface IconProps {
  size?: number;
@@ -281,6 +282,13 @@ export const CPUIcon = ({
  return <FiCpu size={size} className={className} />;
 };

+export const RobotIcon = ({
+  size = 16,
+  className = defaultTailwindCSS,
+}: IconProps) => {
+  return <FaRobot size={size} className={className} />;
+};
+
 //
 //  COMPANY LOGOS
 //
--- a/web/src/components/search/DocumentDisplay.tsx
+++ b/web/src/components/search/DocumentDisplay.tsx
@@ -173,7 +173,7 @@ export const DocumentDisplay = ({
                ml-auto
                mr-2`}
            >
-              {document.score.toFixed(2)}
+              {Math.abs(document.score).toFixed(2)}
            </div>
          </div>
        )}
--- a/web/src/components/search/PersonaSelector.tsx
+++ b/web/src/components/search/PersonaSelector.tsx
@@ -0,0 +1,120 @@
+import { Persona } from "@/app/admin/personas/interfaces";
+import { CustomDropdown } from "../Dropdown";
+import { FiCheck, FiChevronDown } from "react-icons/fi";
+import { FaRobot } from "react-icons/fa";
+
+function PersonaItem({
+  id,
+  name,
+  onSelect,
+  isSelected,
+  isFinal,
+}: {
+  id: number;
+  name: string;
+  onSelect: (personaId: number) => void;
+  isSelected: boolean;
+  isFinal: boolean;
+}) {
+  return (
+    <div
+      key={id}
+      className={`
+    flex
+    px-3 
+    text-sm 
+    text-gray-200 
+    py-2.5 
+    select-none 
+    cursor-pointer 
+    ${isFinal ? "" : "border-b border-gray-800"} 
+    ${
+      isSelected
+        ? "bg-dark-tremor-background-muted"
+        : "hover:bg-dark-tremor-background-muted "
+    }
+  `}
+      onClick={() => {
+        onSelect(id);
+      }}
+    >
+      {name}
+      {isSelected && (
+        <div className="ml-auto mr-1">
+          <FiCheck />
+        </div>
+      )}
+    </div>
+  );
+}
+
+export function PersonaSelector({
+  personas,
+  selectedPersonaId,
+  onPersonaChange,
+}: {
+  personas: Persona[];
+  selectedPersonaId: number | null;
+  onPersonaChange: (persona: Persona | null) => void;
+}) {
+  const currentlySelectedPersona = personas.find(
+    (persona) => persona.id === selectedPersonaId
+  );
+
+  return (
+    <CustomDropdown
+      dropdown={
+        <div
+          className={`
+            border 
+            border-gray-800 
+            rounded-lg 
+            flex 
+            flex-col 
+            w-64 
+            max-h-96 
+            overflow-y-auto 
+            flex
+            overscroll-contain`}
+        >
+          <PersonaItem
+            key={-1}
+            id={-1}
+            name="Default"
+            onSelect={() => {
+              onPersonaChange(null);
+            }}
+            isSelected={selectedPersonaId === null}
+            isFinal={false}
+          />
+          {personas.map((persona, ind) => {
+            const isSelected = persona.id === selectedPersonaId;
+            return (
+              <PersonaItem
+                key={persona.id}
+                id={persona.id}
+                name={persona.name}
+                onSelect={(clickedPersonaId) => {
+                  const clickedPersona = personas.find(
+                    (persona) => persona.id === clickedPersonaId
+                  );
+                  if (clickedPersona) {
+                    onPersonaChange(clickedPersona);
+                  }
+                }}
+                isSelected={isSelected}
+                isFinal={ind === personas.length - 1}
+              />
+            );
+          })}
+        </div>
+      }
+    >
+      <div className="select-none text-sm flex text-gray-300 px-1 py-1.5 cursor-pointer w-64">
+        <FaRobot className="my-auto mr-2" />
+        {currentlySelectedPersona?.name || "Default"}{" "}
+        <FiChevronDown className="my-auto ml-2" />
+      </div>
+    </CustomDropdown>
+  );
+}
--- a/web/src/components/search/SearchBar.tsx
+++ b/web/src/components/search/SearchBar.tsx
@@ -7,11 +7,7 @@ interface SearchBarProps {
  onSearch: () => void;
 }

-export const SearchBar: React.FC<SearchBarProps> = ({
-  query,
-  setQuery,
-  onSearch,
-}) => {
+export const SearchBar = ({ query, setQuery, onSearch }: SearchBarProps) => {
  const handleChange = (event: ChangeEvent<HTMLTextAreaElement>) => {
    const target = event.target;
    setQuery(target.value);
@@ -30,7 +26,7 @@ export const SearchBar: React.FC<SearchBarProps> = ({
  };

  return (
-    <div className="flex justify-center py-3">
+    <div className="flex justify-center">
      <div className="flex items-center w-full border-2 border-gray-600 rounded px-4 py-2 focus-within:border-blue-500">
        <MagnifyingGlass className="text-gray-400" />
        <textarea
--- a/web/src/components/search/SearchResultsDisplay.tsx
+++ b/web/src/components/search/SearchResultsDisplay.tsx
@@ -20,6 +20,7 @@ import {
 } from "@/lib/search/aiThoughtUtils";
 import { ThreeDots } from "react-loader-spinner";
 import { usePopup } from "../admin/connectors/Popup";
+import { AlertIcon } from "../icons/icons";

 const removeDuplicateDocs = (documents: DanswerDocument[]) => {
  const seen = new Set<string>();
@@ -49,14 +50,16 @@ interface SearchResultsDisplayProps {
  validQuestionResponse: ValidQuestionResponse;
  isFetching: boolean;
  defaultOverrides: SearchDefaultOverrides;
+  personaName?: string | null;
 }

-export const SearchResultsDisplay: React.FC<SearchResultsDisplayProps> = ({
+export const SearchResultsDisplay = ({
  searchResponse,
  validQuestionResponse,
  isFetching,
  defaultOverrides,
-}) => {
+  personaName = null,
+}: SearchResultsDisplayProps) => {
  const { popup, setPopup } = usePopup();
  const [isAIThoughtsOpen, setIsAIThoughtsOpen] = React.useState<boolean>(
    getAIThoughtsIsOpenSavedValue()
@@ -70,6 +73,7 @@ export const SearchResultsDisplay: React.FC<SearchResultsDisplayProps> = ({
    return null;
  }

+  const isPersona = personaName !== null;
  const { answer, quotes, documents, error, queryEventId } = searchResponse;

  if (isFetching && !answer && !documents) {
@@ -92,6 +96,17 @@ export const SearchResultsDisplay: React.FC<SearchResultsDisplayProps> = ({
  }

  if (answer === null && documents === null && quotes === null) {
+    if (error) {
+      return (
+        <div className="text-red-500 text-sm">
+          <div className="flex">
+            <AlertIcon size={16} className="text-red-500 my-auto mr-1" />
+            <p className="italic">{error}</p>
+          </div>
+        </div>
+      );
+    }
+
    return <div className="text-gray-300">No matching documents found.</div>;
  }

@@ -132,34 +147,38 @@ export const SearchResultsDisplay: React.FC<SearchResultsDisplayProps> = ({
              <h2 className="text font-bold my-auto mb-1 w-full">AI Answer</h2>
            </div>

-            <div className="mb-2 w-full">
-              <ResponseSection
-                status={questionValidityCheckStatus}
-                header={
-                  validQuestionResponse.answerable === null ? (
-                    <div className="flex ml-2">Evaluating question...</div>
-                  ) : (
-                    <div className="flex ml-2">AI thoughts</div>
-                  )
-                }
-                body={<div>{validQuestionResponse.reasoning}</div>}
-                desiredOpenStatus={isAIThoughtsOpen}
-                setDesiredOpenStatus={handleAIThoughtToggle}
-              />
-            </div>
+            {!isPersona && (
+              <div className="mb-2 w-full">
+                <ResponseSection
+                  status={questionValidityCheckStatus}
+                  header={
+                    validQuestionResponse.answerable === null ? (
+                      <div className="flex ml-2">Evaluating question...</div>
+                    ) : (
+                      <div className="flex ml-2">AI thoughts</div>
+                    )
+                  }
+                  body={<div>{validQuestionResponse.reasoning}</div>}
+                  desiredOpenStatus={isAIThoughtsOpen}
+                  setDesiredOpenStatus={handleAIThoughtToggle}
+                />
+              </div>
+            )}

            <div className="mb-2 pt-1 border-t border-gray-700 w-full">
              <AnswerSection
                answer={answer}
                quotes={quotes}
                error={error}
-                isAnswerable={validQuestionResponse.answerable}
+                isAnswerable={
+                  validQuestionResponse.answerable || (isPersona ? true : null)
+                }
                isFetching={isFetching}
                aiThoughtsIsOpen={isAIThoughtsOpen}
              />
            </div>

-            {quotes !== null && answer && (
+            {quotes !== null && answer && !isPersona && (
              <div className="pt-1 border-t border-gray-700 w-full">
                <QuotesSection
                  quotes={dedupedQuotes}
--- a/web/src/components/search/SearchSection.tsx
+++ b/web/src/components/search/SearchSection.tsx
@@ -20,8 +20,11 @@ import { SearchHelper } from "./SearchHelper";
 import { CancellationToken, cancellable } from "@/lib/search/cancellable";
 import { NEXT_PUBLIC_DISABLE_STREAMING } from "@/lib/constants";
 import { searchRequest } from "@/lib/search/qa";
-import { useFilters, useObjectState, useTimeRange } from "@/lib/hooks";
+import { useFilters, useObjectState } from "@/lib/hooks";
 import { questionValidationStreamed } from "@/lib/search/streamingQuestionValidation";
+import { createChatSession } from "@/lib/search/chatSessions";
+import { Persona } from "@/app/admin/personas/interfaces";
+import { PersonaSelector } from "./PersonaSelector";

 const SEARCH_DEFAULT_OVERRIDES_START: SearchDefaultOverrides = {
  forceDisplayQA: false,
@@ -36,14 +39,16 @@ const VALID_QUESTION_RESPONSE_DEFAULT: ValidQuestionResponse = {
 interface SearchSectionProps {
  connectors: Connector<any>[];
  documentSets: DocumentSet[];
+  personas: Persona[];
  defaultSearchType: SearchType;
 }

-export const SearchSection: React.FC<SearchSectionProps> = ({
+export const SearchSection = ({
  connectors,
  documentSets,
+  personas,
  defaultSearchType,
-}) => {
+}: SearchSectionProps) => {
  // Search Bar
  const [query, setQuery] = useState<string>("");

@@ -63,6 +68,8 @@ export const SearchSection: React.FC<SearchSectionProps> = ({
  const [selectedSearchType, setSelectedSearchType] =
    useState<SearchType>(defaultSearchType);

+  const [selectedPersona, setSelectedPersona] = useState<number | null>(null);
+
  // Overrides for default behavior that only last a single query
  const [defaultOverrides, setDefaultOverrides] =
    useState<SearchDefaultOverrides>(SEARCH_DEFAULT_OVERRIDES_START);
@@ -134,11 +141,23 @@ export const SearchSection: React.FC<SearchSectionProps> = ({
    setSearchResponse(initialSearchResponse);
    setValidQuestionResponse(VALID_QUESTION_RESPONSE_DEFAULT);

+    const chatSessionResponse = await createChatSession(selectedPersona);
+    if (!chatSessionResponse.ok) {
+      updateError(
+        `Unable to create chat session - ${await chatSessionResponse.text()}`
+      );
+      setIsFetching(false);
+      return;
+    }
+    const chatSessionId = (await chatSessionResponse.json())
+      .chat_session_id as number;
+
    const searchFn = NEXT_PUBLIC_DISABLE_STREAMING
      ? searchRequest
      : searchRequestStreamed;
    const searchFnArgs = {
      query,
+      chatSessionId,
      sources: filterManager.selectedSources,
      documentSets: filterManager.selectedDocumentSets,
      timeRange: filterManager.timeRange,
@@ -180,6 +199,7 @@ export const SearchSection: React.FC<SearchSectionProps> = ({

    const questionValidationArgs = {
      query,
+      chatSessionId,
      update: setValidQuestionResponse,
    };

@@ -226,6 +246,20 @@ export const SearchSection: React.FC<SearchSectionProps> = ({
        </div>
      </div>
      <div className="w-[800px] mx-auto">
+        {personas.length > 0 ? (
+          <div className="flex mb-2 w-64">
+            <PersonaSelector
+              personas={personas}
+              selectedPersonaId={selectedPersona}
+              onPersonaChange={(persona) =>
+                setSelectedPersona(persona ? persona.id : null)
+              }
+            />
+          </div>
+        ) : (
+          <div className="pt-3" />
+        )}
+
        <SearchBar
          query={query}
          setQuery={setQuery}
@@ -241,6 +275,11 @@ export const SearchSection: React.FC<SearchSectionProps> = ({
            validQuestionResponse={validQuestionResponse}
            isFetching={isFetching}
            defaultOverrides={defaultOverrides}
+            personaName={
+              selectedPersona
+                ? personas.find((p) => p.id === selectedPersona)?.name
+                : null
+            }
          />
        </div>
      </div>
--- a/web/src/lib/search/chatSessions.ts
+++ b/web/src/lib/search/chatSessions.ts
@@ -0,0 +1,12 @@
+export async function createChatSession(personaId?: number | null) {
+  const chatSessionResponse = await fetch("/api/chat/create-chat-session", {
+    method: "POST",
+    body: JSON.stringify({
+      persona_id: personaId,
+    }),
+    headers: {
+      "Content-Type": "application/json",
+    },
+  });
+  return chatSessionResponse;
+}
--- a/web/src/lib/search/interfaces.ts
+++ b/web/src/lib/search/interfaces.ts
@@ -92,6 +92,7 @@ export interface Filters {

 export interface SearchRequestArgs {
  query: string;
+  chatSessionId: number;
  sources: Source[];
  documentSets: string[];
  timeRange: DateRangePickerValue | null;
--- a/web/src/lib/search/streamingQa.ts
+++ b/web/src/lib/search/streamingQa.ts
@@ -14,6 +14,7 @@ import { buildFilters } from "./utils";

 export const searchRequestStreamed = async ({
  query,
+  chatSessionId,
  sources,
  documentSets,
  timeRange,
@@ -35,6 +36,7 @@ export const searchRequestStreamed = async ({
    const response = await fetch("/api/stream-direct-qa", {
      method: "POST",
      body: JSON.stringify({
+        chat_session_id: chatSessionId,
        query,
        collection: "danswer_index",
        filters,
--- a/web/src/lib/search/streamingQuestionValidation.ts
+++ b/web/src/lib/search/streamingQuestionValidation.ts
@@ -3,11 +3,13 @@ import { processRawChunkString } from "./streamingUtils";

 export interface QuestionValidationArgs {
  query: string;
+  chatSessionId: number;
  update: (update: Partial<ValidQuestionResponse>) => void;
 }

 export const questionValidationStreamed = async <T>({
  query,
+  chatSessionId,
  update,
 }: QuestionValidationArgs) => {
  const emptyFilters = {
@@ -20,6 +22,7 @@ export const questionValidationStreamed = async <T>({
    method: "POST",
    body: JSON.stringify({
      query,
+      chat_session_id: chatSessionId,
      collection: "danswer_index",
      filters: emptyFilters,
      enable_auto_detect_filters: false,