From e20958d17be60106fe02e29d63e1102efeb62ce2 Mon Sep 17 00:00:00 2001
From: sreyakumar <121137643+sreyakumar@users.noreply.github.com>
Date: Thu, 24 Oct 2024 12:26:57 -0700
Subject: [PATCH] edited generate_for_whole_db

---
 GAMER_workbook.ipynb                    | 4 ++--
 src/metadata_chatbot/agents/workflow.py | 7 +++----
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/GAMER_workbook.ipynb b/GAMER_workbook.ipynb
index e93d605..239b751 100644
--- a/GAMER_workbook.ipynb
+++ b/GAMER_workbook.ipynb
@@ -189,13 +189,13 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Based on the retrieved information, the genotype for subject 675387 is wt/wt (wild-type).\n"
+      "Based on the retrieved information, the genotype of subject 675387 is wt/wt (wild-type).\n"
      ]
     }
    ],
    "source": [
     "from metadata_chatbot.agents.GAMER import GAMER\n",
-    "query = \"What is the genotype for subject 675387?\"\n",
+    "query = \"What are all the assets using mouse 675387\"\n",
     "\n",
     "model = GAMER()\n",
     "result = model.invoke(query)\n",
diff --git a/src/metadata_chatbot/agents/workflow.py b/src/metadata_chatbot/agents/workflow.py
index cdf940e..81fd512 100644
--- a/src/metadata_chatbot/agents/workflow.py
+++ b/src/metadata_chatbot/agents/workflow.py
@@ -64,9 +64,8 @@ def generate_for_whole_db(state):
 
     logging.info("Generating answer...")
 
-    documents_list = db_surveyor.invoke({'query': query, 'chat_history': chat_history, 'agent_scratchpad': []})
-    documents_list = [str(item) for item in documents_list]
-    documents = "\n\n".join(documents_list)
+    documents_dict = db_surveyor.invoke({'query': query, 'chat_history': chat_history, 'agent_scratchpad': []})
+    documents = documents_dict['output'][0]['text']
     return {"query": query, "documents": documents}
 
 def filter_generator(state):
@@ -191,7 +190,7 @@ def generate(state):
 
 app = workflow.compile()
 
-# query = "Can you give me a timeline of events for subject 675387?"
+# query = "What are all the assets using mouse 675387"
 
 # inputs = {"query" : query}
 # answer = app.invoke(inputs)