From f5707e95dfb5231cbf95d4d910c8fa28c4e371fe Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Thu, 18 Jan 2024 12:01:19 -0800
Subject: [PATCH] Make retrieval tool description configurable

---
 .../packages/gizmo-agent/gizmo_agent/main.py  | 17 +++++++++--
 .../packages/gizmo-agent/gizmo_agent/tools.py |  6 ++--
 frontend/src/components/Config.tsx            | 29 ++++++++++++++++++-
 3 files changed, 46 insertions(+), 6 deletions(-)

diff --git a/backend/packages/gizmo-agent/gizmo_agent/main.py b/backend/packages/gizmo-agent/gizmo_agent/main.py
index 8399ef782..647e34ce3 100644
--- a/backend/packages/gizmo-agent/gizmo_agent/main.py
+++ b/backend/packages/gizmo-agent/gizmo_agent/main.py
@@ -17,7 +17,13 @@
     get_openai_function_agent,
     # get_xml_agent,
 )
-from gizmo_agent.tools import TOOL_OPTIONS, TOOLS, AvailableTools, get_retrieval_tool
+from gizmo_agent.tools import (
+    RETRIEVAL_DESCRIPTION,
+    TOOL_OPTIONS,
+    TOOLS,
+    AvailableTools,
+    get_retrieval_tool,
+)
 
 DEFAULT_SYSTEM_MESSAGE = "You are a helpful assistant."
 
@@ -26,6 +32,7 @@ class ConfigurableAgent(RunnableBinding):
     tools: Sequence[str]
     agent: GizmoAgentType
     system_message: str = DEFAULT_SYSTEM_MESSAGE
+    retrieval_description: str = RETRIEVAL_DESCRIPTION
     assistant_id: Optional[str] = None
     user_id: Optional[str] = None
 
@@ -36,6 +43,7 @@ def __init__(
         agent: GizmoAgentType = GizmoAgentType.GPT_35_TURBO,
         system_message: str = DEFAULT_SYSTEM_MESSAGE,
         assistant_id: Optional[str] = None,
+        retrieval_description: str = RETRIEVAL_DESCRIPTION,
         kwargs: Optional[Mapping[str, Any]] = None,
         config: Optional[Mapping[str, Any]] = None,
         **others: Any,
@@ -48,7 +56,7 @@ def __init__(
                     raise ValueError(
                         "assistant_id must be provided if Retrieval tool is used"
                     )
-                _tools.append(get_retrieval_tool(assistant_id))
+                _tools.append(get_retrieval_tool(assistant_id, retrieval_description))
             else:
                 _tools.append(TOOLS[_tool]())
         if agent == GizmoAgentType.GPT_35_TURBO:
@@ -70,6 +78,7 @@ def __init__(
             tools=tools,
             agent=agent,
             system_message=system_message,
+            retrieval_description=retrieval_description,
             bound=agent_executor,
             kwargs=kwargs or {},
             config=config or {},
@@ -110,6 +119,7 @@ class AgentOutput(BaseModel):
         agent=GizmoAgentType.GPT_35_TURBO,
         tools=[],
         system_message=DEFAULT_SYSTEM_MESSAGE,
+        retrieval_description=RETRIEVAL_DESCRIPTION,
         assistant_id=None,
     )
     .configurable_fields(
@@ -124,6 +134,9 @@ class AgentOutput(BaseModel):
             options=TOOL_OPTIONS,
             default=[],
         ),
+        retrieval_description=ConfigurableField(
+            id="retrieval_description", name="Retrieval Description"
+        ),
     )
     .configurable_alternatives(
         ConfigurableField(id="type", name="Bot Type"),
diff --git a/backend/packages/gizmo-agent/gizmo_agent/tools.py b/backend/packages/gizmo-agent/gizmo_agent/tools.py
index cdefd966a..29525dce5 100644
--- a/backend/packages/gizmo-agent/gizmo_agent/tools.py
+++ b/backend/packages/gizmo-agent/gizmo_agent/tools.py
@@ -25,17 +25,17 @@ class PythonREPLInput(BaseModel):
     query: str = Field(description="python command to run")
 
 
-RETRIEVER_DESCRIPTION = """Can be used to look up information that was uploaded to this assistant.
+RETRIEVAL_DESCRIPTION = """Can be used to look up information that was uploaded to this assistant.
 If the user is referencing particular files, that is often a good hint that information may be here."""
 
 
-def get_retrieval_tool(assistant_id: str):
+def get_retrieval_tool(assistant_id: str, description: str):
     return create_retriever_tool(
         vstore.as_retriever(
             search_kwargs={"filter": RedisFilter.tag("namespace") == assistant_id}
         ),
         "Retriever",
-        RETRIEVER_DESCRIPTION,
+        description,
     )
 
 
diff --git a/frontend/src/components/Config.tsx b/frontend/src/components/Config.tsx
index aa0bbf391..9cf55bd49 100644
--- a/frontend/src/components/Config.tsx
+++ b/frontend/src/components/Config.tsx
@@ -161,7 +161,9 @@ function MultiOptionField(props: {
 
 function PublicLink(props: { assistantId: string }) {
   const currentLink = window.location.href;
-  const link = currentLink.includes('shared_id=') ? currentLink : currentLink + "?shared_id=" + props.assistantId;
+  const link = currentLink.includes("shared_id=")
+    ? currentLink
+    : currentLink + "?shared_id=" + props.assistantId;
   return (
     <div className="flex rounded-md shadow-sm mb-4">
       <button
@@ -320,6 +322,30 @@ export function Config(props: {
                   readonly={readonly}
                 />
               );
+            } else if (
+              key === "type==agent/retrieval_description" &&
+              (
+                values?.configurable?.["type==agent/tools"] as
+                  | string[]
+                  | undefined
+              )?.includes("Retrieval")
+            ) {
+              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 (key === "type==agent/tools") {
               return (
                 <MultiOptionField
@@ -347,6 +373,7 @@ export function Config(props: {
               setFiles={setFiles}
             />
           )}
+          {}
           <SingleOptionField
             id="public"
             field={{