updated to include KB selector

app.py

@@ -128,7 +128,7 @@ def build_talk_agent(session_id: str, message: str) -> Agent:
 
 # Creates an agent for each run
 # The agent uses local memory, which it differentiates between by session_hash.
-def build_agent(session_id: str, message: str) -> Agent:
+def build_agent(session_id: str, message: str, kbs:str) -> Agent:
 
     create_thread_id(session_id)
 
@@ -159,7 +159,9 @@ def build_agent(session_id: str, message: str) -> Agent:
 
     query_client = StructureRunTool(
         name="QueryResumeSearcher",
-        description="Use it to search for a candidate with the query.",
+        description=f"""Use it to search for a candidate with the query.
+        Add this as another argument after the input: {kbs}
+        """,
         driver=GriptapeCloudStructureRunDriver(
             structure_id=os.getenv("GT_STRUCTURE_ID"),
             api_key=os.getenv("GT_CLOUD_API_KEY"),
@@ -186,15 +188,16 @@ def build_agent(session_id: str, message: str) -> Agent:
     )
 
 
-def send_message(message: str, history, request: gr.Request) -> Any:
+def send_message(message: str, history, knowledge_bases, request: gr.Request) -> Any:
     if request:
         session_hash = request.session_hash
-        agent = build_agent(session_hash, message)
+        agent = build_agent(session_hash, message, str(knowledge_bases))
     response = agent.run(message)
     return response.output.value
 
-
-demo = gr.ChatInterface(fn=send_message)
+with gr.Blocks() as demo:
+    knowledge_bases = gr.CheckboxGroup(choices=["skills","demographics","linked_in","showreels"])
+    chatbot = gr.ChatInterface(fn=send_message, additional_inputs=knowledge_bases)
 demo.launch(auth=(os.environ.get("GRADIO_USERNAME"), os.environ.get("GRADIO_PASSWORD")))
 # demo.launch(share=True)