From d6f2c11a79261e750b3173a1fe8e7a8247d679cf Mon Sep 17 00:00:00 2001
From: Abigail MacKenzie-Armes <70247056+abimacarmes@users.noreply.github.com>
Date: Thu, 22 Aug 2024 10:09:19 -0400
Subject: [PATCH 1/6] fix(toolkit): tool list and switch design fixes (#718)

* fix tool list styling

* update switch styling with on/off

* use changes in create assistant flow
---
 .../AgentSettingsForm/ToolsStep.tsx           |  3 +-
 .../components/Composer/DataSourceMenu.tsx    |  1 +
 .../src/components/UI/Switch.tsx              | 68 +++++++++++++------
 .../src/components/UI/ToggleCard.tsx          | 11 ++-
 4 files changed, 57 insertions(+), 26 deletions(-)

diff --git a/src/interfaces/assistants_web/src/components/AgentSettingsForm/ToolsStep.tsx b/src/interfaces/assistants_web/src/components/AgentSettingsForm/ToolsStep.tsx
index 7033aa1df3..1f65fd1c0a 100644
--- a/src/interfaces/assistants_web/src/components/AgentSettingsForm/ToolsStep.tsx
+++ b/src/interfaces/assistants_web/src/components/AgentSettingsForm/ToolsStep.tsx
@@ -57,7 +57,7 @@ const ToolRow: React.FC<{
   handleSwitch: (checked: boolean) => void;
 }> = ({ name, description, icon, checked, handleSwitch }) => {
   return (
-    <div className="flex flex-col space-y-4 rounded-md border p-4 dark:border-volcanic-300 dark:bg-volcanic-100">
+    <div className="flex flex-col space-y-1 rounded-md border p-4 dark:border-volcanic-300 dark:bg-volcanic-100">
       <div className="flex justify-between">
         <div className="flex items-center space-x-2">
           <div className="flex h-5 w-5 items-center justify-center rounded-sm bg-marble-800 dark:bg-volcanic-200">
@@ -70,6 +70,7 @@ const ToolRow: React.FC<{
         <Switch
           checked={checked}
           onChange={(checked: boolean) => !!name && handleSwitch(checked)}
+          showCheckedState
         />
       </div>
       <Text className="dark:text-marble-800">{description}</Text>
diff --git a/src/interfaces/assistants_web/src/components/Composer/DataSourceMenu.tsx b/src/interfaces/assistants_web/src/components/Composer/DataSourceMenu.tsx
index 765976a3da..2513ed8cec 100644
--- a/src/interfaces/assistants_web/src/components/Composer/DataSourceMenu.tsx
+++ b/src/interfaces/assistants_web/src/components/Composer/DataSourceMenu.tsx
@@ -113,6 +113,7 @@ export const DataSourceMenu: React.FC<Props> = ({ agent, tools }) => {
                     theme="evolved-blue"
                     checked={!!paramsTools?.find((t) => t.name === tool.name)}
                     onChange={(checked) => handleToggle(tool.name!, checked)}
+                    showCheckedState
                   />
                 )}
               </div>
diff --git a/src/interfaces/assistants_web/src/components/UI/Switch.tsx b/src/interfaces/assistants_web/src/components/UI/Switch.tsx
index 5ccf10d275..fd14128221 100644
--- a/src/interfaces/assistants_web/src/components/UI/Switch.tsx
+++ b/src/interfaces/assistants_web/src/components/UI/Switch.tsx
@@ -19,6 +19,7 @@ type Props = {
   label?: string;
   name?: string;
   theme?: COHERE_BRANDED_COLORS;
+  showCheckedState?: boolean;
   className?: string;
 };
 
@@ -30,8 +31,30 @@ export const Switch: React.FC<Props> = ({
   tooltip,
   theme = 'evolved-green',
   name,
+  showCheckedState = false,
   className = '',
 }) => {
+  const themeColors: Partial<Record<COHERE_BRANDED_COLORS, string>> = {
+    blue: 'bg-blue-500 group-hover:bg-blue-400',
+    'evolved-green': 'bg-evolved-green-700 group-hover:bg-evolved-green-500',
+    quartz: 'bg-quartz-500 group-hover:bg-quartz-400',
+    green: 'bg-green-250 group-hover:bg-green-200',
+    mushroom: 'bg-mushroom-600 group-hover:bg-mushroom-500',
+    coral: 'bg-coral-600 group-hover:bg-coral-500',
+    'evolved-blue': 'bg-evolved-blue-500 group-hover:bg-blue-400',
+    'evolved-mushroom': 'bg-evolved-mushroom-500 group-hover:bg-evolved-mushroom-600',
+    'evolved-quartz': 'bg-evolved-quartz-500 group-hover:bg-evolved-quartz-700',
+  };
+
+  const checkedColor = checked
+    ? themeColors[theme]
+    : 'bg-mushroom-900 hover:bg-mushroom-800 dark:bg-volcanic-500 dark:group-hover:bg-volcanic-400';
+
+  const toggleTextClassNames = cn(
+    'pointer-events-none absolute',
+    'pt-0.5 font-variable text-[6px] font-medium uppercase',
+    'transform transition-all duration-300 ease-in-out'
+  );
   return (
     <div className="group flex w-10 items-center">
       <Field>
@@ -47,34 +70,41 @@ export const Switch: React.FC<Props> = ({
             checked={checked}
             onChange={onChange}
             className={cn(
-              'relative inline-flex h-5 w-10 shrink-0 cursor-pointer rounded border-2 border-transparent',
+              'relative inline-flex h-4 w-[30px] shrink-0 cursor-pointer rounded border-2 border-transparent',
               'transition-colors duration-300 ease-in-out',
               'focus-visible:outline focus-visible:outline-offset-4 focus-visible:outline-volcanic-500',
-              {
-                'bg-mushroom-900 hover:bg-mushroom-800 dark:bg-volcanic-500 dark:group-hover:bg-volcanic-400':
-                  !checked,
-                'bg-blue-500 group-hover:bg-blue-400': checked && theme === 'blue',
-                'bg-evolved-green-700 group-hover:bg-evolved-green-500':
-                  checked && theme === 'evolved-green',
-                'bg-quartz-500 group-hover:bg-quartz-400': checked && theme === 'quartz',
-                'bg-green-250 group-hover:bg-green-200': checked && theme === 'green',
-                'bg-mushroom-600 group-hover:bg-mushroom-500': checked && theme === 'mushroom',
-                'bg-coral-600 group-hover:bg-coral-500': checked && theme === 'coral',
-                'bg-evolved-blue-500 group-hover:bg-blue-400': checked && theme === 'evolved-blue',
-                'bg-evolved-mushroom-500 group-hover:bg-evolved-mushroom-600':
-                  checked && theme === 'evolved-mushroom',
-                'bg-evolved-quartz-500 group-hover:bg-evolved-quartz-700':
-                  checked && theme === 'evolved-quartz',
-              }
+              checkedColor
             )}
           >
+            {showCheckedState && (
+              <>
+                <Text
+                  aria-hidden="true"
+                  className={cn(toggleTextClassNames, {
+                    'translate-x-[3px] opacity-100': checked,
+                    'translate-x-[14px] opacity-0': !checked,
+                  })}
+                >
+                  On
+                </Text>
+                <Text
+                  aria-hidden="true"
+                  className={cn(toggleTextClassNames, 'text-mushroom-700 dark:text-volcanic-800', {
+                    'translate-x-[14px] opacity-100': !checked,
+                    'translate-x-0 opacity-0': checked,
+                  })}
+                >
+                  Off
+                </Text>
+              </>
+            )}
             <span
               aria-hidden="true"
               className={cn(
-                'pointer-events-none inline-block h-4 w-5 rounded shadow-lg',
+                'pointer-events-none absolute h-3 w-3 rounded shadow-lg',
                 'transform transition-all duration-300 ease-in-out',
                 {
-                  'translate-x-4': checked,
+                  'translate-x-[14px]': checked,
                   'translate-x-0': !checked,
                   'bg-mushroom-700 dark:bg-volcanic-800': !checked,
                   'bg-marble-950': checked,
diff --git a/src/interfaces/assistants_web/src/components/UI/ToggleCard.tsx b/src/interfaces/assistants_web/src/components/UI/ToggleCard.tsx
index f760a6be5c..69564965a4 100644
--- a/src/interfaces/assistants_web/src/components/UI/ToggleCard.tsx
+++ b/src/interfaces/assistants_web/src/components/UI/ToggleCard.tsx
@@ -43,13 +43,13 @@ export const ToggleCard: React.FC<Props> = ({
   return (
     <div className="flex flex-col rounded-md border border-volcanic-800 bg-volcanic-950 p-4 dark:border-volcanic-300 dark:bg-volcanic-150">
       <div className="flex items-start gap-x-6">
-        <div className="flex flex-grow flex-col gap-y-2">
+        <div className="flex flex-grow flex-col gap-y-1">
           <div className="flex items-center gap-x-2">
-            <div className="flex size-6 items-center justify-center rounded bg-volcanic-900 dark:bg-volcanic-200">
+            <div className="flex size-5 items-center justify-center rounded bg-volcanic-900 dark:bg-volcanic-200">
               <Icon
                 name={icon}
                 kind="outline"
-                size="sm"
+                size="xs"
                 className="fill-volcanic-300 dark:fill-marble-950"
               />
             </div>
@@ -57,9 +57,7 @@ export const ToggleCard: React.FC<Props> = ({
               {label}
             </Text>
           </div>
-          <Text styleAs="p-sm" className="text-mushroom-300 dark:text-marble-800">
-            {description}
-          </Text>
+          <Text className="text-mushroom-300 dark:text-marble-800">{description}</Text>
           {errorMessage && (
             <Text styleAs="p-sm" className="text-danger-350">
               Error: {errorMessage}
@@ -72,6 +70,7 @@ export const ToggleCard: React.FC<Props> = ({
             onChange={onToggle}
             className="flex-shrink-0 gap-0"
             theme={theme}
+            showCheckedState
           />
         )}
       </div>

From 0a7cba5cbdd89a9202b0a0db2dc58010b23aba75 Mon Sep 17 00:00:00 2001
From: Tianjing Li <tianjinglimail@gmail.com>
Date: Thu, 22 Aug 2024 10:16:09 -0400
Subject: [PATCH 2/6] bump: update postcss 8.4.31 -> 8.4.39 in coral-web (#727)

update postcss
---
 src/interfaces/coral_web/package-lock.json | 16 ++++++++--------
 src/interfaces/coral_web/package.json      |  2 +-
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/interfaces/coral_web/package-lock.json b/src/interfaces/coral_web/package-lock.json
index 6ec9a2eacf..745b67977c 100644
--- a/src/interfaces/coral_web/package-lock.json
+++ b/src/interfaces/coral_web/package-lock.json
@@ -38,7 +38,7 @@
         "next-runtime-env": "^3.2.2",
         "node-html-parser": "^6.1.4",
         "parse-domain": "^8.0.1",
-        "postcss": "^8.4.31",
+        "postcss": "^8.4.39",
         "postcss-viewport-height-correction": "^1.1.1",
         "react": "^18.2.0",
         "react-aria": "^3.24.0",
@@ -8883,9 +8883,9 @@
       "dev": true
     },
     "node_modules/picocolors": {
-      "version": "1.0.0",
-      "resolved": "https://registry.npmjs.org/picocolors/-/picocolors-1.0.0.tgz",
-      "integrity": "sha512-1fygroTLlHu66zi26VoTDv8yRgm0Fccecssto+MhsZ0D/DGW2sm8E8AjW7NU5VVTRt5GxbeZ5qBuJr+HyLYkjQ=="
+      "version": "1.0.1",
+      "resolved": "https://registry.npmjs.org/picocolors/-/picocolors-1.0.1.tgz",
+      "integrity": "sha512-anP1Z8qwhkbmu7MFP5iTt+wQKXgwzf7zTyGlcdzabySa9vd0Xt392U0rVmz9poOaBj0uHJKyyo9/upk0HrEQew=="
     },
     "node_modules/picomatch": {
       "version": "2.3.1",
@@ -8935,9 +8935,9 @@
       }
     },
     "node_modules/postcss": {
-      "version": "8.4.38",
-      "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.4.38.tgz",
-      "integrity": "sha512-Wglpdk03BSfXkHoQa3b/oulrotAkwrlLDRSOb9D0bN86FdRyE9lppSp33aHNPgBa0JKCoB+drFLZkQoRRYae5A==",
+      "version": "8.4.41",
+      "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.4.41.tgz",
+      "integrity": "sha512-TesUflQ0WKZqAvg52PWL6kHgLKP6xB6heTOdoYM0Wt2UHyxNa4K25EZZMgKns3BH1RLVbZCREPpLY0rhnNoHVQ==",
       "funding": [
         {
           "type": "opencollective",
@@ -8954,7 +8954,7 @@
       ],
       "dependencies": {
         "nanoid": "^3.3.7",
-        "picocolors": "^1.0.0",
+        "picocolors": "^1.0.1",
         "source-map-js": "^1.2.0"
       },
       "engines": {
diff --git a/src/interfaces/coral_web/package.json b/src/interfaces/coral_web/package.json
index d4ef3de511..360b0aaab1 100644
--- a/src/interfaces/coral_web/package.json
+++ b/src/interfaces/coral_web/package.json
@@ -52,7 +52,7 @@
     "next-runtime-env": "^3.2.2",
     "node-html-parser": "^6.1.4",
     "parse-domain": "^8.0.1",
-    "postcss": "^8.4.31",
+    "postcss": "^8.4.39",
     "postcss-viewport-height-correction": "^1.1.1",
     "react": "^18.2.0",
     "react-aria": "^3.24.0",

From d8c73f90c16ad90b477cc0089cc19a99f4b68804 Mon Sep 17 00:00:00 2001
From: Khalil Najjar <knajjars@gmail.com>
Date: Thu, 22 Aug 2024 17:14:07 +0200
Subject: [PATCH 3/6] fix(toolkit): address UI bugs from 22/08/2024 (#732)

fix(toolkit): bugs 22082024
---
 .../src/app/(main)/(chat)/Chat.tsx            | 20 +++++++++++++++----
 .../src/app/(main)/edit/[agentId]/page.tsx    |  3 ++-
 .../components/Composer/DataSourceMenu.tsx    |  1 -
 .../src/components/MessageRow/Citation.tsx    |  2 +-
 .../src/components/UI/Button.tsx              |  3 ++-
 .../src/hooks/use-brandedColors.ts            |  6 +++---
 6 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/src/interfaces/assistants_web/src/app/(main)/(chat)/Chat.tsx b/src/interfaces/assistants_web/src/app/(main)/(chat)/Chat.tsx
index 32a350336e..68fc43bbee 100644
--- a/src/interfaces/assistants_web/src/app/(main)/(chat)/Chat.tsx
+++ b/src/interfaces/assistants_web/src/app/(main)/(chat)/Chat.tsx
@@ -5,7 +5,7 @@ import { useEffect } from 'react';
 import { Document, ManagedTool } from '@/cohere-client';
 import { Conversation, ConversationError } from '@/components/Conversation';
 import { TOOL_PYTHON_INTERPRETER_ID } from '@/constants';
-import { useAgent, useConversation, useListTools } from '@/hooks';
+import { useAgent, useAvailableTools, useConversation, useListTools } from '@/hooks';
 import { useCitationsStore, useConversationStore, useParamsStore } from '@/stores';
 import { OutputFiles } from '@/stores/slices/citationsSlice';
 import {
@@ -24,6 +24,7 @@ const Chat: React.FC<{ agentId?: string; conversationId?: string }> = ({
   const { setConversation } = useConversationStore();
   const { addCitation, saveOutputFiles } = useCitationsStore();
   const { setParams, resetFileParams } = useParamsStore();
+  const { availableTools } = useAvailableTools({ agent, managedTools: tools });
 
   const {
     data: conversation,
@@ -39,9 +40,11 @@ const Chat: React.FC<{ agentId?: string; conversationId?: string }> = ({
 
     const agentTools =
       agent?.tools &&
-      ((agent.tools
+      (agent.tools
         .map((name) => (tools ?? [])?.find((t) => t.name === name))
-        .filter((t) => t !== undefined) ?? []) as ManagedTool[]);
+        .filter(
+          (t) => t !== undefined && availableTools.some((at) => at.name === t?.name)
+        ) as ManagedTool[]);
 
     const fileIds = conversation?.files.map((file) => file.id);
 
@@ -53,7 +56,16 @@ const Chat: React.FC<{ agentId?: string; conversationId?: string }> = ({
     if (conversationId) {
       setConversation({ id: conversationId });
     }
-  }, [agent, tools, conversation]);
+  }, [
+    agent,
+    tools,
+    conversation,
+    availableTools,
+    setParams,
+    resetFileParams,
+    setConversation,
+    conversationId,
+  ]);
 
   useEffect(() => {
     if (!conversation) return;
diff --git a/src/interfaces/assistants_web/src/app/(main)/edit/[agentId]/page.tsx b/src/interfaces/assistants_web/src/app/(main)/edit/[agentId]/page.tsx
index a41be96c39..d8445addfe 100644
--- a/src/interfaces/assistants_web/src/app/(main)/edit/[agentId]/page.tsx
+++ b/src/interfaces/assistants_web/src/app/(main)/edit/[agentId]/page.tsx
@@ -1,8 +1,9 @@
 import { NextPage } from 'next';
 
-import { UpdateAgent } from './UpdateAgent';
 import { getCohereServerClient } from '@/server/cohereServerClient';
 
+import { UpdateAgent } from './UpdateAgent';
+
 type Props = {
   params: {
     agentId: string;
diff --git a/src/interfaces/assistants_web/src/components/Composer/DataSourceMenu.tsx b/src/interfaces/assistants_web/src/components/Composer/DataSourceMenu.tsx
index 2513ed8cec..1bbdd5099d 100644
--- a/src/interfaces/assistants_web/src/components/Composer/DataSourceMenu.tsx
+++ b/src/interfaces/assistants_web/src/components/Composer/DataSourceMenu.tsx
@@ -27,7 +27,6 @@ export const DataSourceMenu: React.FC<Props> = ({ agent, tools }) => {
   });
 
   const { text, contrastText, border, bg } = useBrandedColors(agent?.id);
-
   const isBaseAgent = checkIsBaseAgent(agent);
   return (
     <Popover className="relative">
diff --git a/src/interfaces/assistants_web/src/components/MessageRow/Citation.tsx b/src/interfaces/assistants_web/src/components/MessageRow/Citation.tsx
index 1d866fbaa3..0c835e58a3 100644
--- a/src/interfaces/assistants_web/src/components/MessageRow/Citation.tsx
+++ b/src/interfaces/assistants_web/src/components/MessageRow/Citation.tsx
@@ -32,7 +32,7 @@ export const Citation: React.FC<Props> = ({ generationId, citationKey, agentId }
   const safeUrl = getSafeUrl(document.url);
   const { text, lightText, fill, lightFill, dark, light } = useBrandedColors(agentId);
 
-  const brandedClassName = cn(text, lightText, dark(lightFill), light(fill));
+  const brandedClassName = cn(dark(lightText), light(text), dark(lightFill), light(fill));
 
   return (
     <div className="space-y-4">
diff --git a/src/interfaces/assistants_web/src/components/UI/Button.tsx b/src/interfaces/assistants_web/src/components/UI/Button.tsx
index b6e52930ef..00697f1a12 100644
--- a/src/interfaces/assistants_web/src/components/UI/Button.tsx
+++ b/src/interfaces/assistants_web/src/components/UI/Button.tsx
@@ -33,7 +33,8 @@ const getLabelStyles = (kind: ButtonKind, theme: ButtonTheme, disabled: boolean)
           theme === 'blue' ||
           theme === 'quartz' ||
           theme === 'evolved-blue' ||
-          theme === 'evolved-quartz',
+          theme === 'evolved-quartz' ||
+          theme === 'mushroom',
         // dark mode
         'dark:text-volcanic-150 dark:fill-volcanic-150':
           theme === 'evolved-green' ||
diff --git a/src/interfaces/assistants_web/src/hooks/use-brandedColors.ts b/src/interfaces/assistants_web/src/hooks/use-brandedColors.ts
index 8294e51717..7b4129d9f4 100644
--- a/src/interfaces/assistants_web/src/hooks/use-brandedColors.ts
+++ b/src/interfaces/assistants_web/src/hooks/use-brandedColors.ts
@@ -15,21 +15,21 @@ const ASSISTANT_COLORS = [
   'quartz-500',
   'evolved-quartz-700',
   'evolved-mushroom-500',
-  'mushroom-300',
+  'mushroom-500',
 ];
 const ASSISTANT_LIGHT_COLORS = [
   'green-800',
   'quartz-800',
   'evolved-quartz-900',
   'evolved-mushroom-800',
-  'mushroom-600',
+  'mushroom-700',
 ];
 const ASSISTANT_CONTRAST_COLORS = [
   'marble-950',
   'marble-950',
   'marble-950',
   'volcanic-100',
-  'mushroom-800',
+  'mushroom-900',
 ];
 const DEFAULT_COLOR = 'evolved-blue-500';
 const DEFAULT_LIGHT_COLOR = 'blue-800';

From de58a58451ea685ba8783410e995ca6dfb6b13c3 Mon Sep 17 00:00:00 2001
From: Tomeu <tomeu@cohere.com>
Date: Thu, 22 Aug 2024 17:57:42 +0200
Subject: [PATCH 4/6] fix: try to update citations starting points when a code
 block is present (#731)

* fix: try to update citations start and end points when there is a code block present

* fix: decoding data before to generate file

* fix: multiple blocks with multi citations
---
 .../src/stores/slices/citationsSlice.ts       |   4 +-
 .../src/utils/citations.test.ts               | 203 +++++++++++++++++-
 .../assistants_web/src/utils/citations.ts     |  85 +++++++-
 3 files changed, 271 insertions(+), 21 deletions(-)

diff --git a/src/interfaces/assistants_web/src/stores/slices/citationsSlice.ts b/src/interfaces/assistants_web/src/stores/slices/citationsSlice.ts
index e1f5ffafc6..4257301f4a 100644
--- a/src/interfaces/assistants_web/src/stores/slices/citationsSlice.ts
+++ b/src/interfaces/assistants_web/src/stores/slices/citationsSlice.ts
@@ -1,7 +1,7 @@
 import { StateCreator } from 'zustand';
 
 import { Document } from '@/cohere-client';
-import { mapExtensionToMimeType } from '@/utils';
+import { decodeBase64, mapExtensionToMimeType } from '@/utils';
 
 import { StoreState } from '..';
 
@@ -87,7 +87,7 @@ export const createCitationsSlice: StateCreator<StoreState, [], [], CitationsSto
       if (file.downloadUrl) {
         continue;
       }
-      const data = file.data;
+      const data = decodeBase64(file.data);
       const fileExtension = file.name.split('.').pop() || '.txt';
       const mimeType = mapExtensionToMimeType(fileExtension);
       const blob = new Blob([data], { type: mimeType });
diff --git a/src/interfaces/assistants_web/src/utils/citations.test.ts b/src/interfaces/assistants_web/src/utils/citations.test.ts
index 9a7ab4eced..b3759336aa 100644
--- a/src/interfaces/assistants_web/src/utils/citations.test.ts
+++ b/src/interfaces/assistants_web/src/utils/citations.test.ts
@@ -2,6 +2,8 @@ import { describe, expect, test } from 'vitest';
 
 import { Citation } from '@/cohere-client';
 import {
+  CODE_BLOCK_REGEX_EXP,
+  IFRAME_REGEX_EXP,
   fixInlineCitationsForMarkdown,
   isReferenceBetweenSpecialTags,
   replaceTextWithCitations,
@@ -214,11 +216,188 @@ describe('fixInlineCitationsForMarkdown', () => {
       },
     ]);
   });
+
+  test('should fix citations position for to get the correct markdown', () => {
+    const citations = [
+      {
+        text: '(fn',
+        start: 71,
+        end: 74,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'delay',
+        start: 76,
+        end: 81,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'let debounceTimer',
+        start: 86,
+        end: 103,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'return function',
+        start: 107,
+        end: 122,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'args',
+        start: 126,
+        end: 130,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'clearTimeout',
+        start: 135,
+        end: 147,
+        document_ids: ['12345'],
+      },
+      {
+        text: '(debounceTimer',
+        start: 147,
+        end: 161,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'debounceTimer = setTimeout',
+        start: 165,
+        end: 191,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'fn.apply(this, args',
+        start: 201,
+        end: 220,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'delay',
+        start: 227,
+        end: 232,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'code is only executed once per user input',
+        start: 285,
+        end: 326,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'takes two parameters',
+        start: 331,
+        end: 351,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'the function to be debounced',
+        start: 353,
+        end: 381,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'the delay in milliseconds',
+        start: 386,
+        end: 411,
+        document_ids: ['12345'],
+      },
+    ];
+    const text =
+      "Here's a JavaScript debounce function: ```javascript function debounce(fn, delay) {   let debounceTimer;    return function(...args) {     clearTimeout(debounceTimer);     debounceTimer = setTimeout(() => {       fn.apply(this, args);     }, delay);   }; } ```  The debounce function ensures that the code is only executed once per user input. It takes two parameters: the function to be debounced and the delay in milliseconds.";
+    const result = fixInlineCitationsForMarkdown(citations, text);
+    expect(result).toStrictEqual([
+      {
+        text: '(fn',
+        start: 71,
+        end: 74,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'delay',
+        start: 76,
+        end: 81,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'let debounceTimer',
+        start: 86,
+        end: 103,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'return function',
+        start: 107,
+        end: 122,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'args',
+        start: 126,
+        end: 130,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'clearTimeout',
+        start: 135,
+        end: 147,
+        document_ids: ['12345'],
+      },
+      {
+        text: '(debounceTimer',
+        start: 147,
+        end: 161,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'debounceTimer = setTimeout',
+        start: 165,
+        end: 191,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'fn.apply(this, args',
+        start: 201,
+        end: 220,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'delay',
+        start: 227,
+        end: 232,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'code is only executed once per user input',
+        start: 301,
+        end: 342,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'takes two parameters',
+        start: 347,
+        end: 367,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'the function to be debounced',
+        start: 369,
+        end: 397,
+        document_ids: ['12345'],
+      },
+      {
+        text: 'the delay in milliseconds',
+        start: 402,
+        end: 427,
+        document_ids: ['12345'],
+      },
+    ]);
+  });
 });
 
 describe('isReferenceBetweenSpecialTags', () => {
   test('should return true if the citation is between <iframe> tags', () => {
-    const matchRegex = /<iframe>.*<\/iframe>/;
     const text = '<iframe> This is a citation </iframe>';
     const citation: Citation = {
       start: 19,
@@ -226,11 +405,10 @@ describe('isReferenceBetweenSpecialTags', () => {
       text: 'citation',
       document_ids: ['12345'],
     };
-    const result = isReferenceBetweenSpecialTags(matchRegex, text, citation.start);
+    const result = isReferenceBetweenSpecialTags(IFRAME_REGEX_EXP, text, citation.start);
     expect(result).toBe(true);
   });
   test('should return false if the citation is not between <iframe> tags', () => {
-    const matchRegex = /<iframe>.*<\/iframe>/;
     const text = 'This is a citation <iframe> another test citaiton </iframe>';
     const citation: Citation = {
       start: 10,
@@ -238,11 +416,10 @@ describe('isReferenceBetweenSpecialTags', () => {
       text: 'citation',
       document_ids: ['12345'],
     };
-    const result = isReferenceBetweenSpecialTags(matchRegex, text, citation.start);
+    const result = isReferenceBetweenSpecialTags(IFRAME_REGEX_EXP, text, citation.start);
     expect(result).toBe(false);
   });
   test('should return true if the citation is between ``` tags', () => {
-    const matchRegex = /```[\s\S]*?```/g;
     const text = '``` This is a citation ```';
     const citation: Citation = {
       start: 14,
@@ -250,11 +427,10 @@ describe('isReferenceBetweenSpecialTags', () => {
       text: 'citation',
       document_ids: ['12345'],
     };
-    const result = isReferenceBetweenSpecialTags(matchRegex, text, citation.start);
+    const result = isReferenceBetweenSpecialTags(CODE_BLOCK_REGEX_EXP, text, citation.start);
     expect(result).toBe(true);
   });
   test('should return false if the citation is not between ``` tags', () => {
-    const matchRegex = /```[\s\S]*?```/g;
     const text = '``` This is a citation ``` another test citaiton';
     const citation: Citation = {
       start: 40,
@@ -262,7 +438,18 @@ describe('isReferenceBetweenSpecialTags', () => {
       text: 'citation',
       document_ids: ['12345'],
     };
-    const result = isReferenceBetweenSpecialTags(matchRegex, text, citation.start);
+    const result = isReferenceBetweenSpecialTags(CODE_BLOCK_REGEX_EXP, text, citation.start);
     expect(result).toBe(false);
   });
+  test('should be able to check if there are more than one match', () => {
+    const text = '``` This is a citation ``` another test citaiton ``` final citation ``` ';
+    const citation: Citation = {
+      start: 59,
+      end: 67,
+      text: 'citation',
+      document_ids: ['12345'],
+    };
+    const result = isReferenceBetweenSpecialTags(CODE_BLOCK_REGEX_EXP, text, citation.start);
+    expect(result).toBe(true);
+  });
 });
diff --git a/src/interfaces/assistants_web/src/utils/citations.ts b/src/interfaces/assistants_web/src/utils/citations.ts
index 5709b4ec0d..49e39a09d2 100644
--- a/src/interfaces/assistants_web/src/utils/citations.ts
+++ b/src/interfaces/assistants_web/src/utils/citations.ts
@@ -10,13 +10,61 @@ export const fixMarkdownImagesInText = (text: string) => {
 };
 
 const formatter = new Intl.ListFormat('en', { style: 'long', type: 'conjunction' });
+export const IFRAME_REGEX_EXP = /<iframe.*<\/iframe>/;
+export const CODE_BLOCK_REGEX_EXP = /```[\s\S]*?```/;
+const escapeRegex = (str: string) => str.replace(/[-[\]{}()*+?.,\\^$|#\s]/g, '\\$&');
+const hasOddOccurrences = (str: string, sequence: string) => {
+  // Escape special regex characters in the sequence
+  const escapedSequence = sequence.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+
+  // Create the regex pattern
+  const regexPattern = `^(?:[^${escapedSequence}]*${escapedSequence}){1}(?:[^${escapedSequence}]*${escapedSequence}[^${escapedSequence}]*${escapedSequence})*[^${escapedSequence}]*$`;
+
+  const regex = new RegExp(regexPattern);
+
+  return regex.test(str);
+};
+
+const tryToFixCitationsInCodeBlock = (citations: Citation[], originalText: string) => {
+  const citationsCopy = structuredClone(citations);
+
+  for (const citation of citationsCopy) {
+    // since we are pushing out the citations out of the special blocks, we need to check if the citation is inside one of them
+    if (
+      isReferenceBetweenSpecialTags(IFRAME_REGEX_EXP, originalText, citation.start) ||
+      isReferenceBetweenSpecialTags(CODE_BLOCK_REGEX_EXP, originalText, citation.start)
+    ) {
+      continue;
+    }
+
+    try {
+      const match = originalText.match(new RegExp(escapeRegex(citation.text)));
+      if (match) {
+        const [firstMatch] = match;
+        const start = match.index || citation.start;
+        const end = start + firstMatch.length;
+        citation.start = start;
+        citation.end = end;
+      }
+    } catch (e) {
+      console.error('error updating citation', citation, e);
+    }
+  }
+
+  return citationsCopy;
+};
 
 export const fixInlineCitationsForMarkdown = (citations: Citation[], originalText: string) => {
-  const citationsCopy = [...citations];
+  let citationsCopy = structuredClone(citations);
   const markdownFixList = ['`', '*', '**'];
   const matchDownloableMarkdownLink = /\[.*\]\(.*/;
   let carryOver = 0;
 
+  // try to fix citations start and end indexes in case they are a code block on the text
+  if (originalText.match(CODE_BLOCK_REGEX_EXP)) {
+    citationsCopy = tryToFixCitationsInCodeBlock(citationsCopy, originalText);
+  }
+
   for (let citation of citationsCopy) {
     if (carryOver) {
       citation.start += carryOver;
@@ -32,21 +80,31 @@ export const fixInlineCitationsForMarkdown = (citations: Citation[], originalTex
        * Explanation: it will render as `Abra envolves into cite[**Kadabra]{generationId="12345" start="44" end="53"}**`
        * breaking the markdown layout. This fix will ensure the citation is rendered as `Abra envolves into cite[**Kadabra**]{generationId="12345" start="44" end="53"}`
        */
-      if (citation.text.startsWith(markdown)) {
+      if (
+        hasOddOccurrences(citation.text, markdown) && // it has an odd number of markdown characters
+        !isReferenceBetweenSpecialTags(CODE_BLOCK_REGEX_EXP, originalText, citation.start) // is not in a code block
+      ) {
         const canWeIncludeNextCharacterInTheCitation =
-          originalText.charAt(citation.end) === markdown;
+          originalText.charAt(citation.end) === markdown; // the next character the markdown character
         if (canWeIncludeNextCharacterInTheCitation) {
           citation.end += markdown.length;
           citation.text = citation.text + markdown;
         }
       }
 
-      if (citation.text.startsWith(' [')) {
+      const isCitationInsideOfCodeBlock = isReferenceBetweenSpecialTags(
+        CODE_BLOCK_REGEX_EXP,
+        originalText,
+        citation.start
+      );
+      const isOutOfBonds = citation.start >= originalText.length + 1; // if the citation is outside of the text
+
+      if (citation.text.startsWith(' [') && !isOutOfBonds && !isCitationInsideOfCodeBlock) {
         citation.text = citation.text.slice(1);
         carryOver -= 1;
       }
 
-      if (citation.text.startsWith('! [')) {
+      if (citation.text.startsWith('! [') && !isOutOfBonds && !isCitationInsideOfCodeBlock) {
         citation.text = '![' + citation.text.slice(3);
         carryOver -= 1;
       }
@@ -75,8 +133,6 @@ export const replaceTextWithCitations = (
   generationId: string
 ) => {
   if (!citations.length || !generationId) return text;
-  const iFrameRegExp = /<iframe.*<\/iframe>/;
-  const codeBlockRegExp = /```[\s\S]*?```/g;
 
   let notFoundReferences: string[] = [];
   let iFrameReferences: string[] = [];
@@ -100,7 +156,7 @@ export const replaceTextWithCitations = (
         return;
       }
 
-      if (isReferenceBetweenSpecialTags(iFrameRegExp, replacedText, citeStart)) {
+      if (isReferenceBetweenSpecialTags(IFRAME_REGEX_EXP, replacedText, citeStart)) {
         const ref = `Reference #${iFrameReferences.length + 1}`;
         iFrameReferences.push(
           `:cite[${ref}]{generationId="${generationId}" start="${start}" end="${end}"}`
@@ -109,7 +165,7 @@ export const replaceTextWithCitations = (
         return;
       }
 
-      if (isReferenceBetweenSpecialTags(codeBlockRegExp, replacedText, citeStart)) {
+      if (isReferenceBetweenSpecialTags(CODE_BLOCK_REGEX_EXP, replacedText, citeStart)) {
         const ref = `Reference #${codeBlockReferences.length + 1}`;
         codeBlockReferences.push(
           `:cite[${ref}]{generationId="${generationId}" start="${start}" end="${end}"}`
@@ -152,6 +208,13 @@ export function isReferenceBetweenSpecialTags(
   citeStart: number
 ): boolean {
   const regex = new RegExp(regExp, 'g');
-  const match = regex.exec(replacedText);
-  return Boolean(match && match.index < citeStart && citeStart < match.index + match[0].length);
+
+  let match;
+  while ((match = regex.exec(replacedText))) {
+    if (match && match.index < citeStart && citeStart < match.index + match[0].length) {
+      return true;
+    }
+  }
+
+  return false;
 }

From 2eeac0a09a1a657b5e8e1c65ff58266400c446ce Mon Sep 17 00:00:00 2001
From: Eugene P <144219719+EugeneLightsOn@users.noreply.github.com>
Date: Thu, 22 Aug 2024 19:24:55 +0300
Subject: [PATCH 5/6] BE: Organization filters (#706)

* TLK_572 - global organization filter for all requests

* TLK_572 - global organization filter for all requests

* TLK_972 - fixes after discussion

* TLK_572 - sync

* TLK_572 - sync + validation

* TLK_572 - sync + validation

* TLK_572 - sync + validation + tests
---
 .../versions/2024_08_19_c301506b3676_.py      | 61 ++++++++--------
 src/backend/config/routers.py                 | 23 +++++++
 src/backend/crud/organization.py              |  2 +-
 src/backend/database_models/base.py           | 40 ++++++++++-
 src/backend/database_models/database.py       |  3 +-
 src/backend/routers/agent.py                  |  3 +
 src/backend/routers/organization.py           | 28 +++++++-
 src/backend/schemas/context.py                | 41 +++++++++++
 src/backend/services/context.py               | 13 ++++
 src/backend/services/request_validators.py    | 20 ++++++
 src/backend/tests/unit/conftest.py            |  5 +-
 src/backend/tests/unit/factories/agent.py     |  1 +
 .../tests/unit/factories/conversation.py      |  1 +
 src/backend/tests/unit/routers/test_agent.py  | 69 +++++++++++++++++++
 src/backend/tests/unit/routers/test_chat.py   | 10 +--
 .../tests/unit/routers/test_conversation.py   | 29 ++++++++
 .../tests/unit/routers/test_organization.py   | 15 ++++
 src/backend/tools/google_drive/tool.py        |  4 +-
 18 files changed, 325 insertions(+), 43 deletions(-)

diff --git a/src/backend/alembic/versions/2024_08_19_c301506b3676_.py b/src/backend/alembic/versions/2024_08_19_c301506b3676_.py
index afb1a1b1e7..aaa7fc9513 100644
--- a/src/backend/alembic/versions/2024_08_19_c301506b3676_.py
+++ b/src/backend/alembic/versions/2024_08_19_c301506b3676_.py
@@ -5,52 +5,55 @@
 Create Date: 2024-08-19 12:36:34.118536
 
 """
+
 from typing import Sequence, Union
 
 import sqlalchemy as sa
 from alembic import op
 
 # revision identifiers, used by Alembic.
-revision: str = 'c301506b3676'
-down_revision: Union[str, None] = 'a76ebb869eb8'
+revision: str = "c301506b3676"
+down_revision: Union[str, None] = "a76ebb869eb8"
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
 
 
 def upgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    op.create_table('groups',
-    sa.Column('display_name', sa.String(), nullable=False),
-    sa.Column('id', sa.String(), nullable=False),
-    sa.Column('created_at', sa.DateTime(), nullable=True),
-    sa.Column('updated_at', sa.DateTime(), nullable=True),
-    sa.PrimaryKeyConstraint('id'),
-    sa.UniqueConstraint('display_name', name='unique_display_name')
+    op.create_table(
+        "groups",
+        sa.Column("display_name", sa.String(), nullable=False),
+        sa.Column("id", sa.String(), nullable=False),
+        sa.Column("created_at", sa.DateTime(), nullable=True),
+        sa.Column("updated_at", sa.DateTime(), nullable=True),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("display_name", name="unique_display_name"),
     )
-    op.create_table('user_group',
-    sa.Column('user_id', sa.String(), nullable=False),
-    sa.Column('group_id', sa.String(), nullable=False),
-    sa.Column('display', sa.String(), nullable=False),
-    sa.Column('id', sa.String(), nullable=False),
-    sa.Column('created_at', sa.DateTime(), nullable=True),
-    sa.Column('updated_at', sa.DateTime(), nullable=True),
-    sa.ForeignKeyConstraint(['group_id'], ['groups.id'], ondelete='CASCADE'),
-    sa.ForeignKeyConstraint(['user_id'], ['users.id'], ondelete='CASCADE'),
-    sa.PrimaryKeyConstraint('user_id', 'group_id', 'id')
+    op.create_table(
+        "user_group",
+        sa.Column("user_id", sa.String(), nullable=False),
+        sa.Column("group_id", sa.String(), nullable=False),
+        sa.Column("display", sa.String(), nullable=False),
+        sa.Column("id", sa.String(), nullable=False),
+        sa.Column("created_at", sa.DateTime(), nullable=True),
+        sa.Column("updated_at", sa.DateTime(), nullable=True),
+        sa.ForeignKeyConstraint(["group_id"], ["groups.id"], ondelete="CASCADE"),
+        sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"),
+        sa.PrimaryKeyConstraint("user_id", "group_id", "id"),
     )
-    op.add_column('users', sa.Column('user_name', sa.String(), nullable=True))
-    op.add_column('users', sa.Column('external_id', sa.String(), nullable=True))
-    op.add_column('users', sa.Column('active', sa.Boolean(), nullable=True))
-    op.create_unique_constraint('unique_user_name', 'users', ['user_name'])
+    op.add_column("users", sa.Column("user_name", sa.String(), nullable=True))
+    op.add_column("users", sa.Column("external_id", sa.String(), nullable=True))
+    op.add_column("users", sa.Column("active", sa.Boolean(), nullable=True))
+    op.create_unique_constraint("unique_user_name", "users", ["user_name"])
     # ### end Alembic commands ###
 
 
 def downgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    op.drop_constraint('unique_user_name', 'users', type_='unique')
-    op.drop_column('users', 'active')
-    op.drop_column('users', 'external_id')
-    op.drop_column('users', 'user_name')
-    op.drop_table('user_group')
-    op.drop_table('groups')
+    op.drop_constraint("unique_user_name", "users", type_="unique")
+    op.drop_column("users", "active")
+    op.drop_column("users", "external_id")
+    op.drop_column("users", "user_name")
+    op.drop_table("user_group")
+    op.drop_table("groups")
     # ### end Alembic commands ###
diff --git a/src/backend/config/routers.py b/src/backend/config/routers.py
index 03c7c06fa5..79f988373a 100644
--- a/src/backend/config/routers.py
+++ b/src/backend/config/routers.py
@@ -9,6 +9,7 @@
 )
 from backend.services.request_validators import (
     validate_chat_request,
+    validate_organization_header,
     validate_user_header,
 )
 
@@ -35,9 +36,11 @@ class RouterName(StrEnum):
     RouterName.AUTH: {
         "default": [
             Depends(get_session),
+            Depends(validate_organization_header),
         ],
         "auth": [
             Depends(get_session),
+            Depends(validate_organization_header),
         ],
     },
     RouterName.CHAT: {
@@ -45,94 +48,114 @@ class RouterName(StrEnum):
             Depends(get_session),
             Depends(validate_user_header),
             Depends(validate_chat_request),
+            Depends(validate_organization_header),
         ],
         "auth": [
             Depends(get_session),
             Depends(validate_chat_request),
             Depends(validate_authorization),
+            Depends(validate_organization_header),
         ],
     },
     RouterName.CONVERSATION: {
         "default": [
             Depends(get_session),
             Depends(validate_user_header),
+            Depends(validate_organization_header),
         ],
         "auth": [
             Depends(get_session),
             Depends(validate_authorization),
+            Depends(validate_organization_header),
         ],
     },
     RouterName.DEPLOYMENT: {
         "default": [
             Depends(get_session),
+            Depends(validate_organization_header),
         ],
         "auth": [
             Depends(get_session),
             Depends(validate_authorization),
+            Depends(validate_organization_header),
         ],
     },
     RouterName.EXPERIMENTAL_FEATURES: {
         "default": [
             Depends(get_session),
+            Depends(validate_organization_header),
         ],
         "auth": [
             Depends(get_session),
             Depends(validate_authorization),
+            Depends(validate_organization_header),
         ],
     },
     RouterName.TOOL: {
         "default": [
             Depends(get_session),
+            Depends(validate_organization_header),
         ],
         "auth": [
             Depends(get_session),
             Depends(validate_authorization),
+            Depends(validate_organization_header),
         ],
     },
     RouterName.USER: {
         "default": [
             Depends(get_session),
+            Depends(validate_organization_header),
         ],
         "auth": [
             # TODO: Remove auth only for create user endpoint
             Depends(get_session),
+            Depends(validate_organization_header),
         ],
     },
     RouterName.AGENT: {
         "default": [
             Depends(get_session),
+            Depends(validate_organization_header),
         ],
         "auth": [
             Depends(get_session),
             Depends(validate_authorization),
+            Depends(validate_organization_header),
         ],
     },
     RouterName.DEFAULT_AGENT: {
         "default": [
             Depends(get_session),
+            Depends(validate_organization_header),
         ],
         "auth": [
             Depends(get_session),
             Depends(validate_authorization),
+            Depends(validate_organization_header),
         ],
     },
     RouterName.SNAPSHOT: {
         "default": [
             Depends(get_session),
             Depends(validate_user_header),
+            Depends(validate_organization_header),
         ],
         "auth": [
             Depends(get_session),
             Depends(validate_authorization),
+            Depends(validate_organization_header),
         ],
     },
     RouterName.MODEL: {
         "default": [
             Depends(get_session),
+            Depends(validate_organization_header),
         ],
         "auth": [
             Depends(get_session),
             Depends(validate_authorization),
+            Depends(validate_organization_header),
         ],
     },
     RouterName.SCIM: {
diff --git a/src/backend/crud/organization.py b/src/backend/crud/organization.py
index 3403c88d89..b974f6780b 100644
--- a/src/backend/crud/organization.py
+++ b/src/backend/crud/organization.py
@@ -1,6 +1,6 @@
 from sqlalchemy.orm import Session
 
-from backend.database_models import Agent
+from backend.database_models.agent import Agent
 from backend.database_models.organization import Organization
 from backend.database_models.user import User, UserOrganizationAssociation
 from backend.schemas.organization import UpdateOrganization
diff --git a/src/backend/database_models/base.py b/src/backend/database_models/base.py
index 32223851d9..3245f14da1 100644
--- a/src/backend/database_models/base.py
+++ b/src/backend/database_models/base.py
@@ -1,7 +1,45 @@
+from enum import StrEnum
 from uuid import uuid4
 
 from sqlalchemy import DateTime, String, func
-from sqlalchemy.orm import DeclarativeBase, mapped_column
+from sqlalchemy.orm import DeclarativeBase, Query, mapped_column
+
+
+class FilterFields(StrEnum):
+    ORGANIZATION_ID = "organization_id"
+
+
+class CustomFilterQuery(Query):
+    """
+    Custom query class that filters by field if the entity has field
+    and the filter value is set.
+    """
+
+    ALLOWED_FILTER_FIELDS = [FilterFields.ORGANIZATION_ID]
+
+    def __new__(cls, *args, **kwargs):
+        from backend.services.context import GLOBAL_REQUEST_CONTEXT
+
+        request_ctx = GLOBAL_REQUEST_CONTEXT.get()
+        if request_ctx and request_ctx.use_global_filtering:
+            query = None
+            for field in cls.ALLOWED_FILTER_FIELDS:
+                if (
+                    args
+                    and hasattr(args[0][0], field)
+                    and hasattr(request_ctx, field)
+                    and getattr(request_ctx, field)
+                ):
+                    if query:
+                        query = query.filter_by(**{field: getattr(request_ctx, field)})
+                    else:
+                        query = Query(*args, **kwargs).filter_by(
+                            **{field: getattr(request_ctx, field)}
+                        )
+            if query:
+                return query
+
+        return object.__new__(cls)
 
 
 class Base(DeclarativeBase):
diff --git a/src/backend/database_models/database.py b/src/backend/database_models/database.py
index 93ef2122f8..7f6ec597f0 100644
--- a/src/backend/database_models/database.py
+++ b/src/backend/database_models/database.py
@@ -6,6 +6,7 @@
 from sqlalchemy.orm import Session
 
 from backend.config.settings import Settings
+from backend.database_models.base import CustomFilterQuery
 
 load_dotenv()
 
@@ -16,7 +17,7 @@
 
 
 def get_session() -> Generator[Session, Any, None]:
-    with Session(engine) as session:
+    with Session(engine, query_cls=CustomFilterQuery) as session:
         yield session
 
 
diff --git a/src/backend/routers/agent.py b/src/backend/routers/agent.py
index 8b57be96ff..97080077cc 100644
--- a/src/backend/routers/agent.py
+++ b/src/backend/routers/agent.py
@@ -193,6 +193,9 @@ async def list_agents(
     # TODO: get organization_id from user
     user_id = ctx.get_user_id()
     logger = ctx.get_logger()
+    # request organization_id is used for filtering agents instead of header Organization-Id if enabled
+    if organization_id:
+        ctx.without_global_filtering()
 
     try:
         agents = agent_crud.get_agents(
diff --git a/src/backend/routers/organization.py b/src/backend/routers/organization.py
index 274e992447..f1a14c2512 100644
--- a/src/backend/routers/organization.py
+++ b/src/backend/routers/organization.py
@@ -11,6 +11,7 @@
     UpdateOrganization,
 )
 from backend.schemas.context import Context
+from backend.schemas.user import User
 from backend.services.context import get_context
 from backend.services.request_validators import validate_organization_request
 
@@ -88,11 +89,11 @@ def get_organization(
         session (DBSessionDep): Database session.
 
     Returns:
-        ManagedTool: Tool with the given ID.
+        ManagedTool: Organization with the given ID.
     """
     organization = organization_crud.get_organization(session, organization_id)
     if not organization:
-        raise HTTPException(status_code=404, detail="Model not found")
+        raise HTTPException(status_code=404, detail="Organization not found")
     return organization
 
 
@@ -114,7 +115,7 @@ def delete_organization(
     """
     organization = organization_crud.get_organization(session, organization_id)
     if not organization:
-        raise HTTPException(status_code=404, detail="Tool not found")
+        raise HTTPException(status_code=404, detail="Organization not found")
     organization_crud.delete_organization(session, organization_id)
 
     return DeleteOrganization()
@@ -138,3 +139,24 @@ def list_organizations(
     """
     all_organizations = organization_crud.get_organizations(session)
     return all_organizations
+
+
+@router.get("/{organization_id}/users", response_model=list[User])
+def get_organization_users(
+    organization_id: str, session: DBSessionDep, ctx: Context = Depends(get_context)
+) -> list[User]:
+    """
+    Get organization users by ID.
+
+    Args:
+        organization_id (str): Organization ID.
+        session (DBSessionDep): Database session.
+
+    Returns:
+        list[User]: List of users in the organization
+    """
+    organization = organization_crud.get_organization(session, organization_id)
+    if not organization:
+        raise HTTPException(status_code=404, detail="Organization not found")
+
+    return organization.users
diff --git a/src/backend/schemas/context.py b/src/backend/schemas/context.py
index 8a7e0c9c30..eefde4142a 100644
--- a/src/backend/schemas/context.py
+++ b/src/backend/schemas/context.py
@@ -2,8 +2,10 @@
 
 from pydantic import BaseModel
 
+from backend.crud import organization as organization_crud
 from backend.crud import user as user_crud
 from backend.database_models.database import DBSessionDep
+from backend.schemas import Organization
 from backend.schemas.agent import Agent, AgentToolMetadata
 from backend.schemas.metrics import MetricsAgent, MetricsMessageType, MetricsUser
 from backend.schemas.user import User
@@ -28,6 +30,9 @@ class Context(BaseModel):
     agent_id: Optional[str] = None
     stream_start_ms: Optional[float] = None
     logger: Optional[Any] = None
+    organization_id: Optional[str] = None
+    organization: Optional[Organization] = None
+    use_global_filtering: Optional[bool] = False
 
     # Metrics
     metrics_user: Optional[MetricsUser] = None
@@ -125,6 +130,42 @@ def with_agent_id(self, agent_id: str) -> "Context":
         self.agent_id = agent_id
         return self
 
+    def with_organization_id(self, organization_id: str) -> "Context":
+        self.organization_id = organization_id
+        return self
+
+    def with_organization(
+        self,
+        session: DBSessionDep | None = None,
+        organization: Organization | None = None,
+    ) -> "Context":
+        if not organization and not session:
+            return self
+
+        if not organization:
+            organization = organization_crud.get_organization(
+                session, self.organization_id
+            )
+            organization = (
+                Organization.model_validate(organization) if organization else None
+            )
+
+        if organization:
+            self.organization = organization
+
+        return self
+
+    def with_global_filtering(self) -> "Context":
+        self.use_global_filtering = True
+        return self
+
+    def without_global_filtering(self) -> "Context":
+        self.use_global_filtering = False
+        return self
+
+    def get_organization(self):
+        return self.organization
+
     def get_stream_start_ms(self):
         return self.stream_start_ms
 
diff --git a/src/backend/services/context.py b/src/backend/services/context.py
index 48bcf99c78..75759f4f9f 100644
--- a/src/backend/services/context.py
+++ b/src/backend/services/context.py
@@ -1,3 +1,4 @@
+import contextvars
 import uuid
 
 from fastapi import Request
@@ -5,6 +6,8 @@
 
 from backend.schemas.context import Context
 
+GLOBAL_REQUEST_CONTEXT = contextvars.ContextVar("GLOBAL_REQUEST_CONTEXT", default=None)
+
 
 class ContextMiddleware:
     def __init__(self, app: ASGIApp):
@@ -41,13 +44,23 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send):
         agent_id = request.headers.get("Agent-Id")
         context.with_agent_id(agent_id)
 
+        organization_id = request.headers.get("Organization-Id", None)
+        context.with_organization_id(organization_id)
+
+        context.without_global_filtering()
+        if organization_id:
+            context.with_global_filtering()
+
         context.with_logger()
 
         # Set the context on the scope
         scope["context"] = context
+        GLOBAL_REQUEST_CONTEXT.set(context)
 
         await self.app(scope, receive, send)
 
+        # Clear the organization ID from the global context
+        GLOBAL_REQUEST_CONTEXT.set(None)
         # Clear the context after the request is complete
         del scope["context"]
 
diff --git a/src/backend/services/request_validators.py b/src/backend/services/request_validators.py
index 754e56c54c..c55c27b114 100644
--- a/src/backend/services/request_validators.py
+++ b/src/backend/services/request_validators.py
@@ -101,6 +101,26 @@ def validate_user_header(session: DBSessionDep, request: Request):
         raise HTTPException(status_code=401, detail="User not found.")
 
 
+def validate_organization_header(session: DBSessionDep, request: Request):
+    """
+    Validate that the request has the `Organization-Id` header, used for requests
+    that require an Organization.
+
+    Args:
+        request (Request): The request to validate
+
+    Raises:
+        HTTPException: If no `Organization-Id` header.
+
+    """
+
+    organization_id = request.headers.get("Organization-Id", None)
+    if organization_id:
+        organization = organization_crud.get_organization(session, organization_id)
+        if not organization:
+            raise HTTPException(status_code=404, detail=f"Organization ID {organization_id} not found.")
+
+
 def validate_deployment_header(request: Request, session: DBSessionDep):
     """
     Validate that the request has the `Deployment-Name` header, used for chat requests
diff --git a/src/backend/tests/unit/conftest.py b/src/backend/tests/unit/conftest.py
index 7486e7cd67..218fb19777 100644
--- a/src/backend/tests/unit/conftest.py
+++ b/src/backend/tests/unit/conftest.py
@@ -11,6 +11,7 @@
 
 from backend.config.deployments import AVAILABLE_MODEL_DEPLOYMENTS, ModelDeploymentName
 from backend.database_models import get_session
+from backend.database_models.base import CustomFilterQuery
 from backend.main import app, create_app
 from backend.schemas.deployment import Deployment
 from backend.schemas.organization import Organization
@@ -47,7 +48,7 @@ def session(engine: Any) -> Generator[Session, None, None]:
     # Begin the nested transaction
     transaction = connection.begin()
     # Use connection within the started transaction
-    session = Session(bind=connection)
+    session = Session(bind=connection, query_cls=CustomFilterQuery)
     # Run Alembic migrations
     alembic_cfg = Config("src/backend/alembic.ini")
     upgrade(alembic_cfg, "head")
@@ -107,7 +108,7 @@ def session_chat(engine_chat: Any) -> Generator[Session, None, None]:
     # Begin the nested transaction
     transaction = connection.begin()
     # Use connection within the started transaction
-    session = Session(bind=connection)
+    session = Session(bind=connection, query_cls=CustomFilterQuery)
     # Run Alembic migrations
     alembic_cfg = Config("src/backend/alembic.ini")
     upgrade(alembic_cfg, "head")
diff --git a/src/backend/tests/unit/factories/agent.py b/src/backend/tests/unit/factories/agent.py
index 022105dfa7..7d2ae9c277 100644
--- a/src/backend/tests/unit/factories/agent.py
+++ b/src/backend/tests/unit/factories/agent.py
@@ -12,6 +12,7 @@ class Meta:
 
     user = factory.SubFactory(UserFactory)
     user_id = factory.SelfAttribute("user.id")
+    organization_id = None
     name = factory.Faker("sentence")
     description = factory.Faker("sentence")
     preamble = factory.Faker("sentence")
diff --git a/src/backend/tests/unit/factories/conversation.py b/src/backend/tests/unit/factories/conversation.py
index 0471f2c8c4..6f064426a8 100644
--- a/src/backend/tests/unit/factories/conversation.py
+++ b/src/backend/tests/unit/factories/conversation.py
@@ -16,6 +16,7 @@ class Meta:
     updated_at = factory.Faker("date_time")
     text_messages = []
     agent_id = None
+    organization_id = None
 
 
 class ConversationFileAssociationFactory(BaseFactory):
diff --git a/src/backend/tests/unit/routers/test_agent.py b/src/backend/tests/unit/routers/test_agent.py
index ad2773662d..86652eb9b4 100644
--- a/src/backend/tests/unit/routers/test_agent.py
+++ b/src/backend/tests/unit/routers/test_agent.py
@@ -181,6 +181,75 @@ def test_list_agents(session_client: TestClient, session: Session, user) -> None
     assert len(response_agents) == 3
 
 
+def test_list_organization_agents(
+    session_client: TestClient,
+    session: Session,
+    user,
+) -> None:
+    session.query(Agent).delete()
+    organization = get_factory("Organization", session).create()
+    organization1 = get_factory("Organization", session).create()
+    for i in range(3):
+        _ = get_factory("Agent", session).create(
+            user=user,
+            organization_id=organization.id,
+            name=f"agent-{i}-{organization.id}",
+        )
+        _ = get_factory("Agent", session).create(
+            user=user, organization_id=organization1.id
+        )
+
+    response = session_client.get(
+        "/v1/agents", headers={"User-Id": user.id, "Organization-Id": organization.id}
+    )
+    assert response.status_code == 200
+    response_agents = response.json()
+    agents = sorted(response_agents, key=lambda x: x["name"])
+    for i in range(3):
+        assert agents[i]["name"] == f"agent-{i}-{organization.id}"
+
+
+def test_list_organization_agents_query_param(
+    session_client: TestClient,
+    session: Session,
+    user,
+) -> None:
+    session.query(Agent).delete()
+    organization = get_factory("Organization", session).create()
+    organization1 = get_factory("Organization", session).create()
+    for i in range(3):
+        _ = get_factory("Agent", session).create(
+            user=user, organization_id=organization.id
+        )
+        _ = get_factory("Agent", session).create(
+            user=user,
+            organization_id=organization1.id,
+            name=f"agent-{i}-{organization1.id}",
+        )
+
+    response = session_client.get(
+        f"/v1/agents?organization_id={organization1.id}",
+        headers={"User-Id": user.id, "Organization-Id": organization.id},
+    )
+    assert response.status_code == 200
+    response_agents = response.json()
+    agents = sorted(response_agents, key=lambda x: x["name"])
+    for i in range(3):
+        assert agents[i]["name"] == f"agent-{i}-{organization1.id}"
+
+
+def test_list_organization_agents_nonexistent_organization(
+    session_client: TestClient,
+    session: Session,
+    user,
+) -> None:
+    response = session_client.get(
+        "/v1/agents", headers={"User-Id": user.id, "Organization-Id": "123"}
+    )
+    assert response.status_code == 404
+    assert response.json() == {"detail": "Organization ID 123 not found."}
+
+
 def test_list_private_agents(
     session_client: TestClient, session: Session, user
 ) -> None:
diff --git a/src/backend/tests/unit/routers/test_chat.py b/src/backend/tests/unit/routers/test_chat.py
index 6b74fefe54..fdcead9310 100644
--- a/src/backend/tests/unit/routers/test_chat.py
+++ b/src/backend/tests/unit/routers/test_chat.py
@@ -136,7 +136,7 @@ def test_streaming_new_chat_metrics_with_agent(
 def test_streaming_new_chat_with_agent(
     session_client_chat: TestClient, session_chat: Session, user: User
 ):
-    agent = get_factory("Agent", session_chat).create(user=user)
+    agent = get_factory("Agent", session_chat).create(user=user, tools=[])
     deployment = get_factory("Deployment", session_chat).create()
     model = get_factory("Model", session_chat).create(deployment=deployment)
     get_factory("AgentDeploymentModel", session_chat).create(
@@ -153,7 +153,7 @@ def test_streaming_new_chat_with_agent(
             "Deployment-Name": agent.deployment,
         },
         params={"agent_id": agent.id},
-        json={"message": "Hello", "max_tokens": 10},
+        json={"message": "Hello", "max_tokens": 10, "agent_id": agent.id},
     )
     assert response.status_code == 200
     validate_chat_streaming_response(
@@ -165,7 +165,7 @@ def test_streaming_new_chat_with_agent(
 def test_streaming_new_chat_with_agent_existing_conversation(
     session_client_chat: TestClient, session_chat: Session, user: User
 ):
-    agent = get_factory("Agent", session_chat).create(user=user)
+    agent = get_factory("Agent", session_chat).create(user=user, tools=[])
     deployment = get_factory("Deployment", session_chat).create()
     model = get_factory("Model", session_chat).create(deployment=deployment)
     get_factory("AgentDeploymentModel", session_chat).create(
@@ -208,7 +208,7 @@ def test_streaming_new_chat_with_agent_existing_conversation(
             "Deployment-Name": agent.deployment,
         },
         params={"agent_id": agent.id},
-        json={"message": "Hello", "max_tokens": 10, "conversation_id": conversation.id},
+        json={"message": "Hello", "max_tokens": 10, "conversation_id": conversation.id, "agent_id": agent.id},
     )
 
     assert response.status_code == 200
@@ -253,7 +253,7 @@ def test_streaming_chat_with_existing_conversation_from_other_agent(
             "Deployment-Name": ModelDeploymentName.CoherePlatform,
         },
         params={"agent_id": agent.id},
-        json={"message": "Hello", "max_tokens": 10, "conversation_id": conversation.id},
+        json={"message": "Hello", "max_tokens": 10, "conversation_id": conversation.id, "agent_id": agent.id},
     )
 
     assert response.status_code == 404
diff --git a/src/backend/tests/unit/routers/test_conversation.py b/src/backend/tests/unit/routers/test_conversation.py
index 692fba0367..359d01ad41 100644
--- a/src/backend/tests/unit/routers/test_conversation.py
+++ b/src/backend/tests/unit/routers/test_conversation.py
@@ -161,6 +161,35 @@ def test_get_conversation_lists_message_files(
     assert response_conversation["messages"][0]["files"][0]["id"] == file.id
 
 
+def test_get_organization_conversation_list(
+    session_client: TestClient,
+    session: Session,
+    user: User,
+) -> None:
+    organization = get_factory("Organization", session).create(name="test org")
+    organization1 = get_factory("Organization", session).create(name="test org1")
+    for i in range(3):
+        get_factory("Conversation", session).create(
+            user_id=user.id,
+            organization_id=organization.id,
+            description=organization.id,
+        )
+        get_factory("Conversation", session).create(
+            user_id=user.id, organization_id=organization1.id
+        )
+
+    response = session_client.get(
+        "/v1/conversations",
+        headers={"User-Id": user.id, "Organization-Id": organization.id},
+    )
+    response_conversation = response.json()
+
+    assert response.status_code == 200
+    assert len(response_conversation) == 3
+    for conversation in response_conversation:
+        assert conversation["description"] == organization.id
+
+
 def test_fail_get_nonexistent_conversation(
     session_client: TestClient,
     session: Session,
diff --git a/src/backend/tests/unit/routers/test_organization.py b/src/backend/tests/unit/routers/test_organization.py
index 70fa4d310f..658321c992 100644
--- a/src/backend/tests/unit/routers/test_organization.py
+++ b/src/backend/tests/unit/routers/test_organization.py
@@ -67,3 +67,18 @@ def test_delete_organization(session_client: TestClient, session: Session) -> No
     response = session_client.delete(f"/v1/organizations/{organization.id}")
     assert response.status_code == 200
     assert response.json() == {}
+
+
+def test_list_organization_users(session_client: TestClient, session: Session) -> None:
+    organization = get_factory("Organization", session).create(name="test organization")
+    for i in range(5):
+        user = get_factory("User", session).create(fullname=f"test user {i}")
+        organization.users.append(user)
+
+    response = session_client.get(f"/v1/organizations/{organization.id}/users")
+    results = response.json()
+    assert response.status_code == 200
+    assert len(results) == 5
+    results = sorted(results, key=lambda x: x["fullname"])
+    for i, result in enumerate(results):
+        assert result["fullname"] == f"test user {i}"
diff --git a/src/backend/tools/google_drive/tool.py b/src/backend/tools/google_drive/tool.py
index 32e3de5372..74212bd208 100644
--- a/src/backend/tools/google_drive/tool.py
+++ b/src/backend/tools/google_drive/tool.py
@@ -74,7 +74,9 @@ async def call(self, parameters: dict, **kwargs: Any) -> List[Dict[str, Any]]:
                 },
             )
             if documents.error:
-                raise Exception(f"Error getting documents for {query} with {documents.error}")
+                raise Exception(
+                    f"Error getting documents for {query} with {documents.error}"
+                )
 
             hits = documents.result["hits"]
             chunks = sorted(

From 95d5639d33d779a51c2dc0cf7d3b6af7f0951309 Mon Sep 17 00:00:00 2001
From: Scott <146760070+scott-cohere@users.noreply.github.com>
Date: Thu, 22 Aug 2024 14:07:53 -0400
Subject: [PATCH 6/6] remove file path in tools (#733)

fix
---
 src/backend/tools/files.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/backend/tools/files.py b/src/backend/tools/files.py
index c153bdbc00..86f181ce91 100644
--- a/src/backend/tools/files.py
+++ b/src/backend/tools/files.py
@@ -126,7 +126,7 @@ async def call(self, parameters: dict, **kwargs: Any) -> List[Dict[str, Any]]:
                 {
                     "text": retrieved_file.file_content,
                     "title": retrieved_file.file_name,
-                    "url": retrieved_file.file_path,
+                    "url": retrieved_file.file_name,
                 }
             ]
 
@@ -181,7 +181,7 @@ async def call(
                     {
                         "text": file.file_content,
                         "title": file.file_name,
-                        "url": file.file_path,
+                        "url": file.file_name,
                     }
                 )
             return results