Unified UI (#3308)

* fix typing * add filters display
onyx-dot-app · Dec 2, 2024 · 7c618c9 · 7c618c9
1 parent 03e2789
commit 7c618c9
Show file tree

Hide file tree

Showing 13 changed files with 246 additions and 49 deletions.
diff --git a/backend/alembic/env.py b/backend/alembic/env.py
@@ -1,5 +1,5 @@
 from sqlalchemy.engine.base import Connection
-from typing import Any
+from typing import Literal
 import asyncio
 from logging.config import fileConfig
 import logging
@@ -8,6 +8,7 @@
 from sqlalchemy import pool
 from sqlalchemy.ext.asyncio import create_async_engine
 from sqlalchemy.sql import text
+from sqlalchemy.sql.schema import SchemaItem
 
 from shared_configs.configs import MULTI_TENANT
 from danswer.db.engine import build_connection_string
@@ -35,7 +36,18 @@
 
 
 def include_object(
-    object: Any, name: str, type_: str, reflected: bool, compare_to: Any
+    object: SchemaItem,
+    name: str | None,
+    type_: Literal[
+        "schema",
+        "table",
+        "column",
+        "index",
+        "unique_constraint",
+        "foreign_key_constraint",
+    ],
+    reflected: bool,
+    compare_to: SchemaItem | None,
 ) -> bool:
     """
     Determines whether a database object should be included in migrations.

diff --git a/backend/alembic_tenants/env.py b/backend/alembic_tenants/env.py
@@ -1,5 +1,6 @@
 import asyncio
 from logging.config import fileConfig
+from typing import Literal
 
 from sqlalchemy import pool
 from sqlalchemy.engine import Connection
@@ -37,8 +38,15 @@
 
 def include_object(
     object: SchemaItem,
-    name: str,
-    type_: str,
+    name: str | None,
+    type_: Literal[
+        "schema",
+        "table",
+        "column",
+        "index",
+        "unique_constraint",
+        "foreign_key_constraint",
+    ],
     reflected: bool,
     compare_to: SchemaItem | None,
 ) -> bool:

diff --git a/backend/danswer/chat/chat_utils.py b/backend/danswer/chat/chat_utils.py
@@ -31,6 +31,7 @@ def llm_doc_from_inference_section(inference_section: InferenceSection) -> LlmDo
         if inference_section.center_chunk.source_links
         else None,
         source_links=inference_section.center_chunk.source_links,
+        match_highlights=inference_section.center_chunk.match_highlights,
     )
 
 

diff --git a/backend/danswer/chat/models.py b/backend/danswer/chat/models.py
@@ -25,6 +25,7 @@ class LlmDoc(BaseModel):
     updated_at: datetime | None
     link: str | None
     source_links: dict[int, str] | None
+    match_highlights: list[str] | None
 
 
 # First chunk of info for streaming QA

diff --git a/backend/danswer/tools/tool_implementations/internet_search/internet_search_tool.py b/backend/danswer/tools/tool_implementations/internet_search/internet_search_tool.py
@@ -77,6 +77,7 @@ def llm_doc_from_internet_search_result(result: InternetSearchResult) -> LlmDoc:
         updated_at=datetime.now(),
         link=result.link,
         source_links={0: result.link},
+        match_highlights=[],
     )
 
 

diff --git a/backend/tests/integration/tests/personas/test_persona_categories.py b/backend/tests/integration/tests/personas/test_persona_categories.py
@@ -44,6 +44,7 @@ def test_persona_category_management(reset: None) -> None:
             category=updated_persona_category,
             user_performing_action=regular_user,
         )
+    assert exc_info.value.response is not None
     assert exc_info.value.response.status_code == 403
 
     assert PersonaCategoryManager.verify(

diff --git a/backend/tests/unit/danswer/llm/answering/conftest.py b/backend/tests/unit/danswer/llm/answering/conftest.py
@@ -64,6 +64,7 @@ def mock_search_results() -> list[LlmDoc]:
             updated_at=datetime(2023, 1, 1),
             link="https://example.com/doc1",
             source_links={0: "https://example.com/doc1"},
+            match_highlights=[],
         ),
         LlmDoc(
             content="Search result 2",
@@ -75,6 +76,7 @@ def mock_search_results() -> list[LlmDoc]:
             updated_at=datetime(2023, 1, 2),
             link="https://example.com/doc2",
             source_links={0: "https://example.com/doc2"},
+            match_highlights=[],
         ),
     ]
 

diff --git a/backend/tests/unit/danswer/llm/answering/stream_processing/test_citation_processing.py b/backend/tests/unit/danswer/llm/answering/stream_processing/test_citation_processing.py
@@ -46,6 +46,7 @@
         updated_at=datetime.now(),
         link=f"https://{int(id/2)}.com" if int(id / 2) % 2 == 0 else None,
         source_links={0: "https://mintlify.com/docs/settings/broken-links"},
+        match_highlights=[],
     )
     for id in range(10)
 ]

diff --git a/backend/tests/unit/danswer/llm/answering/stream_processing/test_quote_processing.py b/backend/tests/unit/danswer/llm/answering/stream_processing/test_quote_processing.py
@@ -20,6 +20,7 @@
         updated_at=datetime.now(),
         link=f"https://{int(id/2)}.com" if int(id / 2) % 2 == 0 else None,
         source_links={0: "https://mintlify.com/docs/settings/broken-links"},
+        match_highlights=[],
     )
     for id in range(10)
 ]

diff --git a/web/src/app/chat/input/ChatInputBar.tsx b/web/src/app/chat/input/ChatInputBar.tsx
@@ -18,7 +18,7 @@ import {
   SendIcon,
   StopGeneratingIcon,
 } from "@/components/icons/icons";
-import { DanswerDocument } from "@/lib/search/interfaces";
+import { DanswerDocument, SourceMetadata } from "@/lib/search/interfaces";
 import { AssistantIcon } from "@/components/assistants/AssistantIcon";
 import {
   Tooltip,
@@ -37,9 +37,41 @@ import { AssistantsTab } from "../modal/configuration/AssistantsTab";
 import { IconType } from "react-icons";
 import { LlmTab } from "../modal/configuration/LlmTab";
 import { XIcon } from "lucide-react";
+import { FilterPills } from "./FilterPills";
+import { Tag } from "@/lib/types";
+import FiltersDisplay from "./FilterDisplay";
 
 const MAX_INPUT_HEIGHT = 200;
 
+interface ChatInputBarProps {
+  removeFilters: () => void;
+  removeDocs: () => void;
+  openModelSettings: () => void;
+  showDocs: () => void;
+  showConfigureAPIKey: () => void;
+  selectedDocuments: DanswerDocument[];
+  message: string;
+  setMessage: (message: string) => void;
+  stopGenerating: () => void;
+  onSubmit: () => void;
+  filterManager: FilterManager;
+  llmOverrideManager: LlmOverrideManager;
+  chatState: ChatState;
+  alternativeAssistant: Persona | null;
+  inputPrompts: InputPrompt[];
+  // assistants
+  selectedAssistant: Persona;
+  setSelectedAssistant: (assistant: Persona) => void;
+  setAlternativeAssistant: (alternativeAssistant: Persona | null) => void;
+
+  files: FileDescriptor[];
+  setFiles: (files: FileDescriptor[]) => void;
+  handleFileUpload: (files: File[]) => void;
+  textAreaRef: React.RefObject<HTMLTextAreaElement>;
+  chatSessionId?: string;
+  toggleFilters?: () => void;
+}
+
 export function ChatInputBar({
   removeFilters,
   removeDocs,
@@ -68,32 +100,7 @@ export function ChatInputBar({
   chatSessionId,
   inputPrompts,
   toggleFilters,
-}: {
-  removeFilters: () => void;
-  removeDocs: () => void;
-  showConfigureAPIKey: () => void;
-  openModelSettings: () => void;
-  chatState: ChatState;
-  stopGenerating: () => void;
-  showDocs: () => void;
-  selectedDocuments: DanswerDocument[];
-  setAlternativeAssistant: (alternativeAssistant: Persona | null) => void;
-  setSelectedAssistant: (assistant: Persona) => void;
-  inputPrompts: InputPrompt[];
-  message: string;
-  setMessage: (message: string) => void;
-  onSubmit: () => void;
-  filterManager: FilterManager;
-  llmOverrideManager: LlmOverrideManager;
-  selectedAssistant: Persona;
-  alternativeAssistant: Persona | null;
-  files: FileDescriptor[];
-  setFiles: (files: FileDescriptor[]) => void;
-  handleFileUpload: (files: File[]) => void;
-  textAreaRef: React.RefObject<HTMLTextAreaElement>;
-  chatSessionId?: string;
-  toggleFilters?: () => void;
-}) {
+}: ChatInputBarProps) {
   useEffect(() => {
     const textarea = textAreaRef.current;
     if (textarea) {
@@ -340,23 +347,26 @@ export function ChatInputBar({
               className="text-sm absolute inset-x-0 top-0 w-full transform -translate-y-full"
             >
               <div className="rounded-lg py-1.5 bg-white border border-border-medium overflow-hidden shadow-lg mx-2 px-1.5 mt-2 rounded z-10">
-                {filteredPrompts.map((currentPrompt, index) => (
-                  <button
-                    key={index}
-                    className={`px-2 ${
-                      tabbingIconIndex == index && "bg-hover"
-                    } rounded content-start flex gap-x-1 py-1.5 w-full  hover:bg-hover cursor-pointer`}
-                    onClick={() => {
-                      updateInputPrompt(currentPrompt);
-                    }}
-                  >
-                    <p className="font-bold">{currentPrompt.prompt}:</p>
-                    <p className="text-left flex-grow mr-auto line-clamp-1">
-                      {currentPrompt.id == selectedAssistant.id && "(default) "}
-                      {currentPrompt.content?.trim()}
-                    </p>
-                  </button>
-                ))}
+                {filteredPrompts.map(
+                  (currentPrompt: InputPrompt, index: number) => (
+                    <button
+                      key={index}
+                      className={`px-2 ${
+                        tabbingIconIndex == index && "bg-hover"
+                      } rounded content-start flex gap-x-1 py-1.5 w-full  hover:bg-hover cursor-pointer`}
+                      onClick={() => {
+                        updateInputPrompt(currentPrompt);
+                      }}
+                    >
+                      <p className="font-bold">{currentPrompt.prompt}:</p>
+                      <p className="text-left flex-grow mr-auto line-clamp-1">
+                        {currentPrompt.id == selectedAssistant.id &&
+                          "(default) "}
+                        {currentPrompt.content?.trim()}
+                      </p>
+                    </button>
+                  )
+                )}
 
                 <a
                   key={filteredPrompts.length}
@@ -430,6 +440,7 @@ export function ChatInputBar({
                 </div>
               </div>
             )}
+
             {(selectedDocuments.length > 0 || files.length > 0) && (
               <div className="flex gap-x-2 px-2 pt-2">
                 <div className="flex gap-x-1 px-2 overflow-visible overflow-x-scroll items-end miniscroll">
@@ -564,6 +575,16 @@ export function ChatInputBar({
                   onClick={toggleFilters}
                 />
               )}
+              {(filterManager.selectedSources.length > 0 ||
+                filterManager.selectedDocumentSets.length > 0 ||
+                filterManager.selectedTags.length > 0 ||
+                filterManager.timeRange) &&
+                toggleFilters && (
+                  <FiltersDisplay
+                    filterManager={filterManager}
+                    toggleFilters={toggleFilters}
+                  />
+                )}
             </div>
 
             <div className="absolute bottom-2.5 mobile:right-4 desktop:right-10">

diff --git a/web/src/app/chat/input/FilterDisplay.tsx b/web/src/app/chat/input/FilterDisplay.tsx
@@ -0,0 +1,109 @@
+import React from "react";
+import { XIcon } from "lucide-react";
+
+import { FilterPills } from "./FilterPills";
+import { SourceMetadata } from "@/lib/search/interfaces";
+import { FilterManager } from "@/lib/hooks";
+import { Tag } from "@/lib/types";
+
+interface FiltersDisplayProps {
+  filterManager: FilterManager;
+  toggleFilters: () => void;
+}
+export default function FiltersDisplay({
+  filterManager,
+  toggleFilters,
+}: FiltersDisplayProps) {
+  return (
+    <div className="flex my-auto flex-wrap gap-2 px-2">
+      {(() => {
+        const allFilters = [
+          ...filterManager.selectedSources,
+          ...filterManager.selectedDocumentSets,
+          ...filterManager.selectedTags,
+          ...(filterManager.timeRange ? [filterManager.timeRange] : []),
+        ];
+        const filtersToShow = allFilters.slice(0, 2);
+        const remainingFilters = allFilters.length - 2;
+
+        return (
+          <>
+            {filtersToShow.map((filter, index) => {
+              if (typeof filter === "object" && "displayName" in filter) {
+                return (
+                  <FilterPills<SourceMetadata>
+                    key={index}
+                    item={filter}
+                    itemToString={(source) => source.displayName}
+                    onRemove={(source) =>
+                      filterManager.setSelectedSources((prev) =>
+                        prev.filter(
+                          (s) => s.internalName !== source.internalName
+                        )
+                      )
+                    }
+                    toggleFilters={toggleFilters}
+                  />
+                );
+              } else if (typeof filter === "string") {
+                return (
+                  <FilterPills<string>
+                    key={index}
+                    item={filter}
+                    itemToString={(set) => set}
+                    onRemove={(set) =>
+                      filterManager.setSelectedDocumentSets((prev) =>
+                        prev.filter((s) => s !== set)
+                      )
+                    }
+                    toggleFilters={toggleFilters}
+                  />
+                );
+              } else if ("tag_key" in filter) {
+                return (
+                  <FilterPills<Tag>
+                    key={index}
+                    item={filter}
+                    itemToString={(tag) => `${tag.tag_key}:${tag.tag_value}`}
+                    onRemove={(tag) =>
+                      filterManager.setSelectedTags((prev) =>
+                        prev.filter(
+                          (t) =>
+                            t.tag_key !== tag.tag_key ||
+                            t.tag_value !== tag.tag_value
+                        )
+                      )
+                    }
+                    toggleFilters={toggleFilters}
+                  />
+                );
+              } else if ("from" in filter && "to" in filter) {
+                return (
+                  <div
+                    key={index}
+                    className="flex items-center bg-background-150 rounded-full px-3 py-1 text-sm"
+                  >
+                    <span>
+                      {filter.from.toLocaleDateString()} -{" "}
+                      {filter.to.toLocaleDateString()}
+                    </span>
+                    <XIcon
+                      onClick={() => filterManager.setTimeRange(null)}
+                      size={16}
+                      className="ml-2 text-text-400 hover:text-text-600 cursor-pointer"
+                    />
+                  </div>
+                );
+              }
+            })}
+            {remainingFilters > 0 && (
+              <div className="flex items-center bg-background-150 rounded-full px-3 py-1 text-sm">
+                <span>+{remainingFilters} more</span>
+              </div>
+            )}
+          </>
+        );
+      })()}
+    </div>
+  );
+}