From 8804b39362825cad738fb15f1a37b8e302ecdbb3 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Mon, 29 Jan 2024 10:17:42 -0800
Subject: [PATCH 1/6] Render documents in UI

---
 backend/app/retrieval.py               |  14 +-
 frontend/src/components/Config.tsx     | 211 ++++++++++++-------------
 frontend/src/components/Document.tsx   |  96 +++++++++++
 frontend/src/components/FileUpload.tsx |   5 +-
 frontend/src/components/Message.tsx    |  20 ++-
 frontend/src/hooks/useChatList.ts      |   5 +-
 6 files changed, 231 insertions(+), 120 deletions(-)
 create mode 100644 frontend/src/components/Document.tsx

diff --git a/backend/app/retrieval.py b/backend/app/retrieval.py
index 42dd6093..df127f27 100644
--- a/backend/app/retrieval.py
+++ b/backend/app/retrieval.py
@@ -67,7 +67,7 @@ def _get_messages(messages):
         ] + chat_history
 
     @chain
-    def get_search_query(messages):
+    async def get_search_query(messages):
         convo = []
         for m in messages:
             if isinstance(m, AIMessage):
@@ -76,11 +76,11 @@ def get_search_query(messages):
             if isinstance(m, HumanMessage):
                 convo.append(f"Human: {m.content}")
         conversation = "\n".join(convo)
-        prompt = search_prompt.invoke({"conversation": conversation})
-        response = llm.invoke(prompt)
+        prompt = await search_prompt.ainvoke({"conversation": conversation})
+        response = await llm.ainvoke(prompt)
         return response.content
 
-    def invoke_retrieval(messages):
+    async def invoke_retrieval(messages):
         if len(messages) == 1:
             human_input = messages[-1].content
             return AIMessage(
@@ -93,7 +93,7 @@ def invoke_retrieval(messages):
                 },
             )
         else:
-            search_query = get_search_query.invoke(messages)
+            search_query = await get_search_query.ainvoke(messages)
             return AIMessage(
                 content="",
                 additional_kwargs={
@@ -104,10 +104,10 @@ def invoke_retrieval(messages):
                 },
             )
 
-    def retrieve(messages):
+    async def retrieve(messages):
         params = messages[-1].additional_kwargs["function_call"]
         query = json.loads(params["arguments"])["query"]
-        response = retriever.invoke(query)
+        response = await retriever.ainvoke(query)
         msg = LiberalFunctionMessage(name="retrieval", content=response)
         return msg
 
diff --git a/frontend/src/components/Config.tsx b/frontend/src/components/Config.tsx
index d1dad7c5..d06efadc 100644
--- a/frontend/src/components/Config.tsx
+++ b/frontend/src/components/Config.tsx
@@ -377,7 +377,19 @@ export function Config(props: {
     </>
   );
   return (
-    <div className={props.className}>
+    <form
+      className={cn("flex flex-col", props.className)}
+      onSubmit={async (e) => {
+        e.preventDefault();
+        e.stopPropagation();
+        const form = e.target as HTMLFormElement;
+        const key = form.key.value;
+        if (!key) return;
+        setInflight(true);
+        await props.saveConfig(key, values!, files, isPublic);
+        setInflight(false);
+      }}
+    >
       {settings}
       {typeField && (
         <Types
@@ -400,113 +412,100 @@ export function Config(props: {
         </>
       )}
 
-      <form
-        className={cn("flex flex-col gap-8")}
-        onSubmit={async (e) => {
-          e.preventDefault();
-          e.stopPropagation();
-          const form = e.target as HTMLFormElement;
-          const key = form.key.value;
-          if (!key) return;
-          setInflight(true);
-          await props.saveConfig(key, values!, files, isPublic);
-          setInflight(false);
-        }}
-      >
-        {!props.config && typeSpec?.files && (
-          <FileUploadDropzone
-            state={dropzone}
-            files={files}
-            setFiles={setFiles}
-          />
+      {!props.config && typeSpec?.files && (
+        <FileUploadDropzone
+          state={dropzone}
+          files={files}
+          setFiles={setFiles}
+          className="mb-8"
+        />
+      )}
+      <div
+        className={cn(
+          "flex flex-col gap-8",
+          readonly && "opacity-50 cursor-not-allowed"
         )}
-        <div
-          className={cn(
-            "flex flex-col gap-8",
-            readonly && "opacity-50 cursor-not-allowed"
-          )}
-        >
-          {orderBy(
-            Object.entries(
-              props.configSchema?.properties.configurable.properties ?? {}
-            ),
-            ([key]) => ORDER.indexOf(last(key.split("/"))!)
-          ).map(([key, value]) => {
-            const title = value.title;
-            if (key.split("/")[0].includes("==")) {
-              const [parentKey, parentValue] = key.split("/")[0].split("==");
-              if (values?.configurable?.[parentKey] !== parentValue) {
-                return null;
-              }
-            } else {
-              return null;
-            }
-            if (
-              last(key.split("/")) === "retrieval_description" &&
-              !files.length
-            ) {
+      >
+        {orderBy(
+          Object.entries(
+            props.configSchema?.properties.configurable.properties ?? {}
+          ),
+          ([key]) => ORDER.indexOf(last(key.split("/"))!)
+        ).map(([key, value]) => {
+          const title = value.title;
+          if (key.split("/")[0].includes("==")) {
+            const [parentKey, parentValue] = key.split("/")[0].split("==");
+            if (values?.configurable?.[parentKey] !== parentValue) {
               return null;
             }
-            if (value.type === "string" && value.enum) {
-              return (
-                <SingleOptionField
-                  key={key}
-                  id={key}
-                  field={value}
-                  title={title}
-                  value={values?.configurable?.[key] as string}
-                  setValue={(value: string) =>
-                    setValues({
-                      ...values,
-                      configurable: { ...values!.configurable, [key]: value },
-                    })
-                  }
-                  readonly={readonly}
-                />
-              );
-            } else if (value.type === "string") {
-              return (
-                <StringField
-                  key={key}
-                  id={key}
-                  field={value}
-                  title={title}
-                  value={values?.configurable?.[key] as string}
-                  setValue={(value: string) =>
-                    setValues({
-                      ...values,
-                      configurable: { ...values!.configurable, [key]: value },
-                    })
-                  }
-                  readonly={readonly}
-                />
-              );
-            } else if (
-              value.type === "array" &&
-              value.items?.type === "string" &&
-              value.items?.enum
-            ) {
-              return (
-                <MultiOptionField
-                  key={key}
-                  id={key}
-                  field={value}
-                  title={title}
-                  value={values?.configurable?.[key] as string[]}
-                  setValue={(value: string[]) =>
-                    setValues({
-                      ...values,
-                      configurable: { ...values!.configurable, [key]: value },
-                    })
-                  }
-                  readonly={readonly}
-                  descriptions={TOOL_DESCRIPTIONS}
-                />
-              );
-            }
-          })}
-        </div>
-      </form>
-    </div>
+          } else {
+            return null;
+          }
+          if (
+            last(key.split("/")) === "retrieval_description" &&
+            !files.length
+          ) {
+            return null;
+          }
+          if (value.type === "string" && value.enum) {
+            return (
+              <SingleOptionField
+                key={key}
+                id={key}
+                field={value}
+                title={title}
+                value={values?.configurable?.[key] as string}
+                setValue={(value: string) =>
+                  setValues({
+                    ...values,
+                    configurable: { ...values!.configurable, [key]: value },
+                  })
+                }
+                readonly={readonly}
+              />
+            );
+          } else if (value.type === "string") {
+            return (
+              <StringField
+                key={key}
+                id={key}
+                field={value}
+                title={title}
+                value={values?.configurable?.[key] as string}
+                setValue={(value: string) =>
+                  setValues({
+                    ...values,
+                    configurable: { ...values!.configurable, [key]: value },
+                  })
+                }
+                readonly={readonly}
+              />
+            );
+          } else if (
+            value.type === "array" &&
+            value.items?.type === "string" &&
+            value.items?.enum
+          ) {
+            return (
+              <MultiOptionField
+                key={key}
+                id={key}
+                field={value}
+                title={title}
+                value={values?.configurable?.[key] as string[]}
+                setValue={(value: string[]) =>
+                  setValues({
+                    ...values,
+                    configurable: { ...values!.configurable, [key]: value },
+                  })
+                }
+                readonly={readonly}
+                descriptions={TOOL_DESCRIPTIONS}
+              />
+            );
+          }
+        })}
+      </div>
+    </form>
   );
 }
diff --git a/frontend/src/components/Document.tsx b/frontend/src/components/Document.tsx
new file mode 100644
index 00000000..4bb1b4c1
--- /dev/null
+++ b/frontend/src/components/Document.tsx
@@ -0,0 +1,96 @@
+import { useMemo, useState } from "react";
+import { cn } from "../utils/cn";
+
+export interface PageDocument {
+  page_content: string;
+  metadata: Record<string, unknown>;
+}
+
+function PageDocument(props: { document: PageDocument; className?: string }) {
+  const [open, setOpen] = useState(false);
+
+  const metadata = useMemo(() => {
+    return Object.keys(props.document.metadata)
+      .sort((a, b) => {
+        const aValue = JSON.stringify(props.document.metadata[a]);
+        const bValue = JSON.stringify(props.document.metadata[b]);
+
+        const aLines = aValue.split("\n");
+        const bLines = bValue.split("\n");
+
+        if (aLines.length !== bLines.length) {
+          return aLines.length - bLines.length;
+        }
+
+        return aValue.length - bValue.length;
+      })
+      .map((key) => {
+        const value = props.document.metadata[key];
+        return {
+          key,
+          value:
+            typeof value === "string" || typeof value === "number"
+              ? `${value}`
+              : JSON.stringify(value),
+        };
+      });
+  }, [props.document.metadata]);
+
+  if (!open) {
+    return (
+      <button
+        className={cn(
+          "flex min-w-0 max-w-full gap-4 overflow-hidden px-2 transition-colors hover:bg-divider-500/25 active:bg-divider-500/50",
+          props.className
+        )}
+        onClick={() => setOpen(true)}
+      >
+        <span className="min-w-0 flex-grow overflow-hidden text-ellipsis whitespace-nowrap text-left">
+          {props.document.page_content.trim().replace(/\n/g, " ")}
+        </span>
+      </button>
+    );
+  }
+
+  return (
+    <button
+      className={cn(
+        "flex gap-4 px-2 text-left transition-colors hover:bg-divider-500/25 active:bg-divider-500/50",
+        props.className
+      )}
+      onClick={() => setOpen(false)}
+    >
+      <span className="flex flex-col gap-4">
+        <span className="whitespace-pre-wrap">
+          {props.document.page_content}
+        </span>
+
+        <span className="flex flex-col flex-wrap items-start gap-2">
+          {metadata.map(({ key, value }, idx) => {
+            return (
+              <span
+                className="rounded-md bg-divider-500 p-1 px-1.5 text-xs"
+                key={idx}
+              >
+                <span className="mr-1.5 font-mono font-bold">{key}</span>
+                <span className="whitespace-pre-wrap">{value}</span>
+              </span>
+            );
+          })}
+        </span>
+      </span>
+    </button>
+  );
+}
+
+export function DocumentList(props: { documents: PageDocument[] }) {
+  return (
+    <div className="flex flex-col items-stretch gap-4 rounded-lg ring-1 ring-gray-300 overflow-hidden">
+      <div className="mx-2 grid divide-y empty:hidden">
+        {props.documents.map((document, idx) => (
+          <PageDocument document={document} key={idx} className="py-3" />
+        ))}
+      </div>
+    </div>
+  );
+}
diff --git a/frontend/src/components/FileUpload.tsx b/frontend/src/components/FileUpload.tsx
index 6bbac49d..e7276893 100644
--- a/frontend/src/components/FileUpload.tsx
+++ b/frontend/src/components/FileUpload.tsx
@@ -45,6 +45,7 @@ export function FileUploadDropzone(props: {
   state: DropzoneState;
   files: File[];
   setFiles: React.Dispatch<React.SetStateAction<File[]>>;
+  className?: string;
 }) {
   const { getRootProps, getInputProps, fileRejections } = props.state;
 
@@ -74,7 +75,7 @@ export function FileUploadDropzone(props: {
   );
 
   return (
-    <section className="">
+    <section className={props.className}>
       <aside>
         <Label id="files" title="Files" />
         <div className="prose">
@@ -86,7 +87,7 @@ export function FileUploadDropzone(props: {
         <p>
           Drag n' drop some files here, or click to select files.
           <br />
-          Accepted files: .txt, .csv, .html, .docx, .pdf. 
+          Accepted files: .txt, .csv, .html, .docx, .pdf.
           <br />
           No file should exceed 10 MB.
         </p>
diff --git a/frontend/src/components/Message.tsx b/frontend/src/components/Message.tsx
index 2df9ec9d..8d7d8bca 100644
--- a/frontend/src/components/Message.tsx
+++ b/frontend/src/components/Message.tsx
@@ -6,6 +6,7 @@ import { marked } from "marked";
 import DOMPurify from "dompurify";
 import { ChevronDownIcon } from "@heroicons/react/24/outline";
 import { LangSmithActions } from "./LangSmithActions";
+import { DocumentList } from "./Document";
 
 function tryJsonParse(value: string) {
   try {
@@ -91,19 +92,23 @@ export const Message = memo(function Message(
   props: MessageType & { runId?: string }
 ) {
   const [open, setOpen] = useState(false);
+  const contentIsDocuments =
+    ["function", "tool"].includes(props.type) &&
+    Array.isArray(props.content) &&
+    props.content.every((d) => !!d.page_content);
   return (
     <div className="flex flex-col mb-8">
       <div className="leading-6 flex flex-row">
         <div
           className={cn(
-            "font-medium text-sm text-gray-400 uppercase mr-2 mt-1 w-24 flex flex-col",
+            "font-medium text-sm text-gray-400 uppercase mr-2 mt-1 w-28 flex flex-col",
             props.type === "function" && "mt-2"
           )}
         >
-          {props.type}
+          {contentIsDocuments ? "Documents" : props.type}
         </div>
         <div className="flex-1">
-          {["function", "tool"].includes(props.type) && (
+          {["function", "tool"].includes(props.type) && !contentIsDocuments && (
             <Function
               call={false}
               name={props.name ?? props.additional_kwargs?.name}
@@ -127,7 +132,11 @@ export const Message = memo(function Message(
                 args={call.function?.arguments}
               />
             ))}
-          {(["function", "tool"].includes(props.type) ? open : true) ? (
+          {(
+            ["function", "tool"].includes(props.type) && !contentIsDocuments
+              ? open
+              : true
+          ) ? (
             typeof props.content === "string" ? (
               <div
                 className="text-gray-900 prose"
@@ -135,6 +144,9 @@ export const Message = memo(function Message(
                   __html: DOMPurify.sanitize(marked(props.content)).trim(),
                 }}
               />
+            ) : contentIsDocuments ? (
+              // eslint-disable-next-line @typescript-eslint/no-explicit-any
+              <DocumentList documents={props.content as any} />
             ) : (
               <div className="text-gray-900 prose">{str(props.content)}</div>
             )
diff --git a/frontend/src/hooks/useChatList.ts b/frontend/src/hooks/useChatList.ts
index 514c123e..caed0825 100644
--- a/frontend/src/hooks/useChatList.ts
+++ b/frontend/src/hooks/useChatList.ts
@@ -3,7 +3,10 @@ import orderBy from "lodash/orderBy";
 
 export interface Message {
   type: string;
-  content: string;
+  content:
+    | string
+    | { page_content: string; metadata: Record<string, object> }[]
+    | object;
   name?: string;
   additional_kwargs?: {
     name?: string;

From 4afb37f1a3331e5252a3f16af45e34b03ad0c53c Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Mon, 29 Jan 2024 11:31:59 -0800
Subject: [PATCH 2/6] Add chevron

---
 frontend/src/components/Document.tsx | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/frontend/src/components/Document.tsx b/frontend/src/components/Document.tsx
index 4bb1b4c1..f0b51756 100644
--- a/frontend/src/components/Document.tsx
+++ b/frontend/src/components/Document.tsx
@@ -1,5 +1,6 @@
 import { useMemo, useState } from "react";
 import { cn } from "../utils/cn";
+import { ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/24/outline";
 
 export interface PageDocument {
   page_content: string;
@@ -40,12 +41,13 @@ function PageDocument(props: { document: PageDocument; className?: string }) {
     return (
       <button
         className={cn(
-          "flex min-w-0 max-w-full gap-4 overflow-hidden px-2 transition-colors hover:bg-divider-500/25 active:bg-divider-500/50",
+          "flex items-start min-w-0 max-w-full gap-4 overflow-hidden px-2 transition-colors hover:bg-divider-500/25 active:bg-divider-500/50",
           props.className
         )}
         onClick={() => setOpen(true)}
       >
-        <span className="min-w-0 flex-grow overflow-hidden text-ellipsis whitespace-nowrap text-left">
+        <ChevronRightIcon className="mt-0.5 h-5 w-5" />
+        <span className="min-w-0 flex-grow basis-0 overflow-hidden text-ellipsis whitespace-nowrap text-left">
           {props.document.page_content.trim().replace(/\n/g, " ")}
         </span>
       </button>
@@ -55,13 +57,15 @@ function PageDocument(props: { document: PageDocument; className?: string }) {
   return (
     <button
       className={cn(
-        "flex gap-4 px-2 text-left transition-colors hover:bg-divider-500/25 active:bg-divider-500/50",
+        "flex items-start gap-4 px-2 text-left transition-colors hover:bg-divider-500/25 active:bg-divider-500/50",
         props.className
       )}
       onClick={() => setOpen(false)}
     >
-      <span className="flex flex-col gap-4">
-        <span className="whitespace-pre-wrap">
+      <ChevronDownIcon className="mt-0.5 h-5 w-5" />
+
+      <span className="flex flex-grow basis-0 flex-col gap-4">
+        <span className="whitespace-pre-line">
           {props.document.page_content}
         </span>
 

From 8d8e0e0559bc3b9e2ca51e19ea52bb6a9fcde04c Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Mon, 29 Jan 2024 11:36:57 -0800
Subject: [PATCH 3/6] More clickable

---
 frontend/src/components/Document.tsx | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/frontend/src/components/Document.tsx b/frontend/src/components/Document.tsx
index f0b51756..2863ab27 100644
--- a/frontend/src/components/Document.tsx
+++ b/frontend/src/components/Document.tsx
@@ -41,12 +41,12 @@ function PageDocument(props: { document: PageDocument; className?: string }) {
     return (
       <button
         className={cn(
-          "flex items-start min-w-0 max-w-full gap-4 overflow-hidden px-2 transition-colors hover:bg-divider-500/25 active:bg-divider-500/50",
+          "flex items-start min-w-0 max-w-full gap-4 overflow-hidden px-4 transition-colors hover:bg-gray-50/50 active:bg-gray-50",
           props.className
         )}
         onClick={() => setOpen(true)}
       >
-        <ChevronRightIcon className="mt-0.5 h-5 w-5" />
+        <ChevronRightIcon className="mt-1 h-4 w-4 text-gray-500" />
         <span className="min-w-0 flex-grow basis-0 overflow-hidden text-ellipsis whitespace-nowrap text-left">
           {props.document.page_content.trim().replace(/\n/g, " ")}
         </span>
@@ -57,12 +57,12 @@ function PageDocument(props: { document: PageDocument; className?: string }) {
   return (
     <button
       className={cn(
-        "flex items-start gap-4 px-2 text-left transition-colors hover:bg-divider-500/25 active:bg-divider-500/50",
+        "flex items-start gap-4 px-4 text-left transition-colors hover:bg-gray-50/50 active:bg-gray-50",
         props.className
       )}
       onClick={() => setOpen(false)}
     >
-      <ChevronDownIcon className="mt-0.5 h-5 w-5" />
+      <ChevronDownIcon className="mt-1 h-4 w-4 text-gray-500" />
 
       <span className="flex flex-grow basis-0 flex-col gap-4">
         <span className="whitespace-pre-line">
@@ -90,7 +90,7 @@ function PageDocument(props: { document: PageDocument; className?: string }) {
 export function DocumentList(props: { documents: PageDocument[] }) {
   return (
     <div className="flex flex-col items-stretch gap-4 rounded-lg ring-1 ring-gray-300 overflow-hidden">
-      <div className="mx-2 grid divide-y empty:hidden">
+      <div className="grid divide-y empty:hidden">
         {props.documents.map((document, idx) => (
           <PageDocument document={document} key={idx} className="py-3" />
         ))}

From 671dc720782cdfc0e8fdfa9078fe5dd5496b009f Mon Sep 17 00:00:00 2001
From: Harrison Chase <hw.chase.17@gmail.com>
Date: Mon, 29 Jan 2024 12:38:36 -0800
Subject: [PATCH 4/6] cr

---
 backend/app/agent.py                    |  6 +++---
 backend/app/agent_types/google_agent.py | 15 ++++++++++++--
 backend/app/agent_types/openai_agent.py | 21 ++++++++++++++-----
 backend/app/agent_types/xml_agent.py    | 27 ++++++++++++++++++-------
 backend/app/message_types.py            | 11 ++++++++++
 backend/app/retrieval.py                | 15 ++++++--------
 6 files changed, 69 insertions(+), 26 deletions(-)
 create mode 100644 backend/app/message_types.py

diff --git a/backend/app/agent.py b/backend/app/agent.py
index 23ea62a3..5df7c7ef 100644
--- a/backend/app/agent.py
+++ b/backend/app/agent.py
@@ -11,13 +11,15 @@
 from app.agent_types.google_agent import get_google_agent_executor
 from app.agent_types.openai_agent import get_openai_agent_executor
 from app.agent_types.xml_agent import get_xml_agent_executor
+from app.chatbot import get_chatbot_executor
 from app.checkpoint import RedisCheckpoint
 from app.llms import (
     get_anthropic_llm,
     get_google_llm,
-    get_openai_llm,
     get_mixtral_fireworks,
+    get_openai_llm,
 )
+from app.retrieval import get_retrieval_executor
 from app.tools import (
     RETRIEVAL_DESCRIPTION,
     TOOL_OPTIONS,
@@ -26,8 +28,6 @@
     get_retrieval_tool,
     get_retriever,
 )
-from app.chatbot import get_chatbot_executor
-from app.retrieval import get_retrieval_executor
 
 
 class AgentType(str, Enum):
diff --git a/backend/app/agent_types/google_agent.py b/backend/app/agent_types/google_agent.py
index daf8752e..499dfbf2 100644
--- a/backend/app/agent_types/google_agent.py
+++ b/backend/app/agent_types/google_agent.py
@@ -9,6 +9,8 @@
 from langgraph.graph.message import MessageGraph
 from langgraph.prebuilt import ToolExecutor, ToolInvocation
 
+from app.message_types import LiberalFunctionMessage
+
 
 def get_google_agent_executor(
     tools: list[BaseTool],
@@ -17,7 +19,16 @@ def get_google_agent_executor(
     checkpoint: BaseCheckpointSaver,
 ):
     def _get_messages(messages):
-        return [SystemMessage(content=system_message)] + messages
+        msgs = []
+        for m in messages:
+            if isinstance(m, LiberalFunctionMessage):
+                _dict = m.dict()
+                _dict["content"] = str(_dict["content"])
+                m_c = FunctionMessage(**_dict)
+                msgs.append(m_c)
+            else:
+                msgs.append(m)
+        return [SystemMessage(content=system_message)] + msgs
 
     if tools:
         llm_with_tools = llm.bind(functions=tools)
@@ -51,7 +62,7 @@ async def call_tool(messages):
         # We call the tool_executor and get back a response
         response = await tool_executor.ainvoke(action)
         # We use the response to create a FunctionMessage
-        function_message = FunctionMessage(content=str(response), name=action.tool)
+        function_message = LiberalFunctionMessage(content=response, name=action.tool)
         # We return a list, because this will get added to the existing list
         return function_message
 
diff --git a/backend/app/agent_types/openai_agent.py b/backend/app/agent_types/openai_agent.py
index c4bf05d0..93ed1215 100644
--- a/backend/app/agent_types/openai_agent.py
+++ b/backend/app/agent_types/openai_agent.py
@@ -1,15 +1,16 @@
 import json
 
-from langchain.schema.messages import ToolMessage
 from langchain.tools import BaseTool
 from langchain.tools.render import format_tool_to_openai_tool
 from langchain_core.language_models.base import LanguageModelLike
-from langchain_core.messages import SystemMessage
+from langchain_core.messages import SystemMessage, ToolMessage
 from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.graph import END
 from langgraph.graph.message import MessageGraph
 from langgraph.prebuilt import ToolExecutor, ToolInvocation
 
+from app.message_types import LiberalToolMessage
+
 
 def get_openai_agent_executor(
     tools: list[BaseTool],
@@ -18,7 +19,17 @@ def get_openai_agent_executor(
     checkpoint: BaseCheckpointSaver,
 ):
     def _get_messages(messages):
-        return [SystemMessage(content=system_message)] + messages
+        msgs = []
+        for m in messages:
+            if isinstance(m, LiberalToolMessage):
+                _dict = m.dict()
+                _dict["content"] = str(_dict["content"])
+                m_c = ToolMessage(**_dict)
+                msgs.append(m_c)
+            else:
+                msgs.append(m)
+
+        return [SystemMessage(content=system_message)] + msgs
 
     if tools:
         llm_with_tools = llm.bind(tools=[format_tool_to_openai_tool(t) for t in tools])
@@ -58,9 +69,9 @@ async def call_tool(messages):
         responses = await tool_executor.abatch(actions)
         # We use the response to create a ToolMessage
         tool_messages = [
-            ToolMessage(
+            LiberalToolMessage(
                 tool_call_id=tool_call["id"],
-                content=json.dumps(response),
+                content=response,
                 additional_kwargs={"name": tool_call["function"]["name"]},
             )
             for tool_call, response in zip(
diff --git a/backend/app/agent_types/xml_agent.py b/backend/app/agent_types/xml_agent.py
index 2a1654cc..6e194f76 100644
--- a/backend/app/agent_types/xml_agent.py
+++ b/backend/app/agent_types/xml_agent.py
@@ -1,14 +1,19 @@
-from langchain.schema.messages import FunctionMessage
 from langchain.tools import BaseTool
 from langchain.tools.render import render_text_description
 from langchain_core.language_models.base import LanguageModelLike
-from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+from langchain_core.messages import (
+    AIMessage,
+    FunctionMessage,
+    HumanMessage,
+    SystemMessage,
+)
 from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.graph import END
 from langgraph.graph.message import MessageGraph
 from langgraph.prebuilt import ToolExecutor, ToolInvocation
 
 from app.agent_types.prompts import xml_template
+from app.message_types import LiberalFunctionMessage
 
 
 def _collapse_messages(messages):
@@ -39,6 +44,11 @@ def construct_chat_history(messages):
                 collapsed_messages.append(_collapse_messages(temp_messages))
                 temp_messages = []
             collapsed_messages.append(message)
+        elif isinstance(message, LiberalFunctionMessage):
+            _dict = message.dict()
+            _dict["content"] = str(_dict["content"])
+            m_c = FunctionMessage(**_dict)
+            temp_messages.append(m_c)
         else:
             temp_messages.append(message)
 
@@ -61,7 +71,7 @@ def get_xml_agent_executor(
         tool_names=", ".join([t.name for t in tools]),
     )
 
-    llm_with_stop = llm.bind(stop=["</tool_input>"])
+    llm_with_stop = llm.bind(stop=["</tool_input>", "<observation>"])
 
     def _get_messages(messages):
         return [
@@ -87,9 +97,12 @@ async def call_tool(messages):
         # We construct an ToolInvocation from the function_call
         tool, tool_input = last_message.content.split("</tool>")
         _tool = tool.split("<tool>")[1]
-        _tool_input = tool_input.split("<tool_input>")[1]
-        if "</tool_input>" in _tool_input:
-            _tool_input = _tool_input.split("</tool_input>")[0]
+        if "<tool_input>" not in _tool:
+            _tool_input = ""
+        else:
+            _tool_input = tool_input.split("<tool_input>")[1]
+            if "</tool_input>" in _tool_input:
+                _tool_input = _tool_input.split("</tool_input>")[0]
         action = ToolInvocation(
             tool=_tool,
             tool_input=_tool_input,
@@ -97,7 +110,7 @@ async def call_tool(messages):
         # We call the tool_executor and get back a response
         response = await tool_executor.ainvoke(action)
         # We use the response to create a FunctionMessage
-        function_message = FunctionMessage(content=str(response), name=action.tool)
+        function_message = LiberalFunctionMessage(content=response, name=action.tool)
         # We return a list, because this will get added to the existing list
         return function_message
 
diff --git a/backend/app/message_types.py b/backend/app/message_types.py
new file mode 100644
index 00000000..c38e66e4
--- /dev/null
+++ b/backend/app/message_types.py
@@ -0,0 +1,11 @@
+from typing import Any
+
+from langchain_core.messages import FunctionMessage, ToolMessage
+
+
+class LiberalFunctionMessage(FunctionMessage):
+    content: Any
+
+
+class LiberalToolMessage(ToolMessage):
+    content: Any
diff --git a/backend/app/retrieval.py b/backend/app/retrieval.py
index df127f27..972aa836 100644
--- a/backend/app/retrieval.py
+++ b/backend/app/retrieval.py
@@ -1,24 +1,21 @@
 import json
 from typing import Any
+
 from langchain_core.language_models.base import LanguageModelLike
 from langchain_core.messages import (
-    SystemMessage,
-    HumanMessage,
     AIMessage,
     FunctionMessage,
+    HumanMessage,
+    SystemMessage,
 )
-from langchain_core.runnables import chain
+from langchain_core.prompts import PromptTemplate
 from langchain_core.retrievers import BaseRetriever
+from langchain_core.runnables import chain
 from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.graph import END
 from langgraph.graph.message import MessageGraph
-from langchain_core.prompts import PromptTemplate
-
-
-# This is sadly needed to allow for arbitrary typed content
-class LiberalFunctionMessage(FunctionMessage):
-    content: Any
 
+from app.message_types import LiberalFunctionMessage
 
 search_prompt = PromptTemplate.from_template(
     """Given the conversation below, come up with a search query to look up.

From 932fb4a7c12c84b49c26951110b0eeb76a1ccbcc Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Mon, 29 Jan 2024 12:39:59 -0800
Subject: [PATCH 5/6] Clearer docs rendering

---
 frontend/src/components/Document.tsx | 2 +-
 frontend/src/components/Message.tsx  | 8 ++++----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/frontend/src/components/Document.tsx b/frontend/src/components/Document.tsx
index 2863ab27..39e84a72 100644
--- a/frontend/src/components/Document.tsx
+++ b/frontend/src/components/Document.tsx
@@ -89,7 +89,7 @@ function PageDocument(props: { document: PageDocument; className?: string }) {
 
 export function DocumentList(props: { documents: PageDocument[] }) {
   return (
-    <div className="flex flex-col items-stretch gap-4 rounded-lg ring-1 ring-gray-300 overflow-hidden">
+    <div className="flex flex-col items-stretch gap-4 rounded-lg ring-1 ring-gray-300 overflow-hidden my-2">
       <div className="grid divide-y empty:hidden">
         {props.documents.map((document, idx) => (
           <PageDocument document={document} key={idx} className="py-3" />
diff --git a/frontend/src/components/Message.tsx b/frontend/src/components/Message.tsx
index 8d7d8bca..24d65ea5 100644
--- a/frontend/src/components/Message.tsx
+++ b/frontend/src/components/Message.tsx
@@ -35,7 +35,7 @@ function Function(props: {
           {props.name}
         </span>
       )}
-      {!props.call && (
+      {!props.call && props.setOpen && (
         <span
           className={cn(
             "inline-flex items-center rounded-md bg-gray-50 px-2 py-1 text-sm font-medium text-gray-600 ring-1 ring-inset ring-gray-500/10 cursor-pointer relative top-1",
@@ -105,15 +105,15 @@ export const Message = memo(function Message(
             props.type === "function" && "mt-2"
           )}
         >
-          {contentIsDocuments ? "Documents" : props.type}
+          {props.type}
         </div>
         <div className="flex-1">
-          {["function", "tool"].includes(props.type) && !contentIsDocuments && (
+          {["function", "tool"].includes(props.type) && (
             <Function
               call={false}
               name={props.name ?? props.additional_kwargs?.name}
               open={open}
-              setOpen={setOpen}
+              setOpen={contentIsDocuments ? undefined : setOpen}
             />
           )}
           {props.additional_kwargs?.function_call && (

From 25c3c4d9f09436085daf958a0fa23362b1a66a36 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Mon, 29 Jan 2024 13:33:05 -0800
Subject: [PATCH 6/6] Lint

---
 backend/app/retrieval.py | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/backend/app/retrieval.py b/backend/app/retrieval.py
index 972aa836..4a713129 100644
--- a/backend/app/retrieval.py
+++ b/backend/app/retrieval.py
@@ -1,13 +1,7 @@
 import json
-from typing import Any
 
 from langchain_core.language_models.base import LanguageModelLike
-from langchain_core.messages import (
-    AIMessage,
-    FunctionMessage,
-    HumanMessage,
-    SystemMessage,
-)
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 from langchain_core.prompts import PromptTemplate
 from langchain_core.retrievers import BaseRetriever
 from langchain_core.runnables import chain