From 291b8f0ca6fc5bde4b17d0854c02a14c2718cbe4 Mon Sep 17 00:00:00 2001 From: jdbocarsly Date: Thu, 7 Nov 2024 10:48:14 -0600 Subject: [PATCH] add all data from collection to a chat --- pydatalab/src/pydatalab/apps/chat/blocks.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/pydatalab/src/pydatalab/apps/chat/blocks.py b/pydatalab/src/pydatalab/apps/chat/blocks.py index 949fb24a0..26b5b2fff 100644 --- a/pydatalab/src/pydatalab/apps/chat/blocks.py +++ b/pydatalab/src/pydatalab/apps/chat/blocks.py @@ -116,7 +116,11 @@ def plot_functions(self): def render(self): if not self.data.get("messages"): if (item_id := self.data.get("item_id")) is not None: - info_json = self._prepare_item_json_for_chat(item_id) + info_json, collections = self._prepare_item_json_for_chat(item_id) + if len(collections): + info_json = self._prepare_collection_json_for_chat( + collections[0]["collection_id"] + ) elif (collection_id := self.data.get("collection_id")) is not None: info_json = self._prepare_collection_json_for_chat(collection_id) else: @@ -129,7 +133,7 @@ def render(self): }, { "role": "user", - "content": f"""Here is the JSON data for the current item(s): {info_json}. + "content": f"""Here is the JSON data for the current item, {item_id}. If the item is in a collection, the data for the rest of the items in the collection are also given: {info_json}. Start with a friendly introduction and give me a one sentence summary of what this is (not detailed, no information about specific masses). """, }, ] @@ -315,7 +319,7 @@ def _prepare_item_json_for_chat(self, item_id: str): .replace(r"\n", " ") ) - return item_info_json + return item_info_json, item_info.get("collections", []) def _prepare_collection_json_for_chat(self, collection_id: str): from pydatalab.routes.v0_1.collections import get_collection @@ -325,8 +329,8 @@ def _prepare_collection_json_for_chat(self, collection_id: str): raise RuntimeError(f"Attempt to get collection data for {collection_id} failed.") children = collection_data["child_items"] - return ( - "[" - + ",".join([self._prepare_item_json_for_chat(child["item_id"]) for child in children]) - + "]" - ) + jsons = [] + for child in children: + jsons.append(self._prepare_item_json_for_chat(child["item_id"])[0]) + + return "[" + ",".join(jsons) + "]"