Skip to content

Commit b73ee82

Browse files
committed
bugfixes, refact: removed buggy returns as error was written to md, formatting was called twice
1 parent 54adbe2 commit b73ee82

File tree

2 files changed

+67
-50
lines changed

2 files changed

+67
-50
lines changed

chat.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -56,8 +56,6 @@ def export(
5656

5757
# Format the chat data
5858
formatter = MarkdownChatFormatter()
59-
formatted_chats = formatter.format(chat_data_dict, image_dir)
60-
6159
if output_dir:
6260
# Save the chat data
6361
saver = MarkdownFileSaver()
@@ -66,10 +64,12 @@ def export(
6664
success_message = f"Chat data has been successfully exported to {output_dir}"
6765
logger.info(success_message)
6866
else:
67+
formatted_chats = formatter.format(chat_data_dict, image_dir)
6968
# Print the chat data to the command line using markdown
7069
for formatted_data in formatted_chats:
7170
console.print(Markdown(formatted_data))
7271
logger.info("Chat data has been successfully printed to the command line")
72+
7373
except KeyError as e:
7474
error_message = f"KeyError: {e}. The chat data structure is not as expected. Please check the database content."
7575
logger.error(error_message)
@@ -193,7 +193,8 @@ def discover(
193193
if results:
194194
for db_path, result in results:
195195
console.print(Markdown("---"))
196-
console.print(f"DATABASE: {os.path.join(os.path.basename(os.path.dirname(db_path)), os.path.basename(db_path))}\n")
196+
folder_rel_path = os.path.join(os.path.basename(os.path.dirname(db_path)), os.path.basename(db_path))
197+
console.print(f"DATABASE: [link=file://{os.path.dirname(db_path).replace(' ', '%20')}]'{db_path}'[/link]\n")
197198
console.print(Markdown(result))
198199
console.print('\n\n')
199200
else:

src/export.py

Lines changed: 63 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,46 @@ def format(self, chat_data: dict[str, Any], image_dir: str = 'images') -> str:
2222
pass
2323

2424
class MarkdownChatFormatter(ChatFormatter):
25-
def format(self, chat_data: dict[str, Any], image_dir: str | None = 'images') -> str:
25+
def _extract_text_from_user_bubble(self, bubble: dict) -> str:
26+
try:
27+
if "delegate" in bubble:
28+
if bubble["delegate"]:
29+
user_text_text = bubble['delegate']["a"]
30+
else:
31+
user_text_text = ""
32+
elif "text" in bubble:
33+
if bubble["text"]:
34+
user_text_text = bubble['text']
35+
else:
36+
user_text_text = ""
37+
elif "initText" in bubble:
38+
if bubble["initText"]:
39+
try:
40+
user_text_text = json.loads(bubble["initText"])['root']['children'][0]['children'][0]['text']
41+
except Exception as e:
42+
user_text_text = "[ERROR: no user text found]"
43+
logger.error(f"Couldn't find user text entry in one of the bubbles. Error: {e}")
44+
logger.debug(f"Bubble:\n{json.dumps(bubble, indent=4)}")
45+
else:
46+
user_text_text = ""
47+
elif "rawText" in bubble:
48+
if bubble["rawText"]:
49+
user_text_text = bubble['text']
50+
else:
51+
user_text_text = ""
52+
else:
53+
user_text_text = "[ERROR: no user text found]"
54+
logger.error(f"Couldn't find user text entry in one of the bubbles.")
55+
logger.debug(f"Bubble:\n{json.dumps(bubble, indent=4)}")
56+
57+
except Exception as e:
58+
user_text_text = "[ERROR: no user text found]"
59+
logger.error(f"Couldn't find user text entry in one of the bubbles. Error: {e}")
60+
logger.debug(f"Bubble:\n{json.dumps(bubble, indent=4)}")
61+
62+
return user_text_text
63+
64+
def format(self, chat_data: dict[str, Any], image_dir: str | None = 'images') -> str | None:
2665
"""Format the chat data into Markdown format.
2766
2867
Args:
@@ -37,56 +76,41 @@ def format(self, chat_data: dict[str, Any], image_dir: str | None = 'images') ->
3776
for tab_index, tab in enumerate(chat_data['tabs']):
3877
bubbles = tab['bubbles']
3978
formatted_chat = [f"# Chat Transcript - Tab {tab_index + 1}\n"]
40-
tab_image_dir = os.path.join(image_dir, f"tab_{tab_index + 1}") if image_dir else None
41-
if tab_image_dir is not None:
42-
os.makedirs(tab_image_dir, exist_ok=True)
4379

4480
for bubble in bubbles:
81+
# USER
4582
if bubble['type'] == 'user':
4683
user_text = ["## User:\n\n"]
4784

85+
# Selections
4886
if "selections" in bubble and bubble["selections"]:
4987
user_text.append(f"[selections] \n{"\n".join([s["text"] for s in bubble['selections']])}")
5088

51-
if 'image' in bubble and tab_image_dir is not None:
89+
# Images
90+
if 'image' in bubble and image_dir is not None:
5291
image_path = bubble['image']['path']
53-
image_filename = os.path.basename(image_path)
54-
new_image_path = os.path.join(tab_image_dir, image_filename)
55-
shutil.copy(image_path, new_image_path)
56-
user_text.append(f"[image] \n![User Image]({new_image_path})")
57-
58-
if "delegate" in bubble:
59-
if bubble["delegate"]:
60-
user_text_text = bubble['delegate']["a"]
92+
if os.path.exists(image_path):
93+
image_filename = os.path.basename(image_path)
94+
new_image_path = os.path.join(tab_image_dir, image_filename)
95+
tab_image_dir = os.path.join(image_dir, f"tab_{tab_index + 1}") if image_dir else None
96+
if tab_image_dir is not None:
97+
os.makedirs(tab_image_dir, exist_ok=True)
98+
shutil.copy(image_path, new_image_path)
99+
user_text.append(f"[image] \n![User Image]({new_image_path})")
61100
else:
62-
user_text_text = ""
63-
elif "text" in bubble:
64-
if bubble["text"]:
65-
user_text_text = bubble['text']
66-
else:
67-
user_text_text = ""
68-
elif "initText" in bubble:
69-
if bubble["initText"]:
70-
user_text_text = bubble['rawText']
71-
user_text_text = json.loads(bubble["initText"])['root']['children'][0]['children'][0]['text']
72-
else:
73-
user_text_text = ""
74-
elif "rawText" in bubble:
75-
if bubble["rawText"]:
76-
user_text_text = bubble['text']
77-
else:
78-
user_text_text = ""
79-
else:
80-
user_text_text = "[ERROR: no user text found]"
81-
logger.error(f"Couldn't find user text entry in one of the bubbles.")
82-
logger.debug(f"Bubble:\n{json.dumps(bubble, indent=4)}")
101+
logger.error(f"Image file {image_path} not found for tab {tab_index + 1}.")
102+
user_text.append(f"[image] \n![User Image]()")
103+
104+
# Text
105+
user_text_text = self._extract_text_from_user_bubble(bubble)
83106
if user_text_text:
84107
user_text.append(f"[text] \n{user_text_text}")
85108

86109
user_text.append("\n")
87110

88111
if len(user_text) > 2:
89112
formatted_chat.append("\n".join(user_text))
113+
# AI
90114
elif bubble['type'] == 'ai':
91115
model_type = bubble.get('modelType', 'Unknown')
92116
raw_text = re.sub(r'```python:[^\n]+', '```python', bubble['rawText'])
@@ -96,18 +120,9 @@ def format(self, chat_data: dict[str, Any], image_dir: str | None = 'images') ->
96120

97121
logger.success("Chats formatted.")
98122
return formatted_chats
99-
except KeyError as e:
100-
logger.error(f"KeyError: Missing key {e} in chat data. Ensure the data structure is correct. Full error: {e}")
101-
return [f"Error: Missing key {e}"]
102-
except FileNotFoundError as e:
103-
logger.error(f"FileNotFoundError: {e}. Ensure the image paths are correct. Full error: {e}")
104-
return [f"Error: {e}"]
105-
except json.JSONDecodeError as e:
106-
logger.error(f"JSONDecodeError: {e}. Ensure the JSON data is correctly formatted. Full error: {e}")
107-
return [f"Error: {e}"]
108123
except Exception as e:
109124
logger.error(f"Unexpected error: {e}. Full traceback: {traceback.format_exc()}")
110-
return [f"Error: {e}"]
125+
return
111126

112127
class FileSaver(ABC):
113128
@abstractmethod
@@ -159,9 +174,10 @@ def export(self, chat_data: dict[str, Any], output_dir: str, image_dir: str) ->
159174
try:
160175
os.makedirs(output_dir, exist_ok=True)
161176
formatted_chats = self.formatter.format(chat_data, image_dir)
162-
for tab_index, formatted_data in enumerate(formatted_chats):
163-
tab_file_path = os.path.join(output_dir, f"tab_{tab_index + 1}.md")
164-
self.saver.save(formatted_data, tab_file_path)
177+
if formatted_chats is not None:
178+
for tab_index, formatted_data in enumerate(formatted_chats):
179+
tab_file_path = os.path.join(output_dir, f"tab_{tab_index + 1}.md")
180+
self.saver.save(formatted_data, tab_file_path)
165181
except Exception as e:
166182
logger.error(f"Failed to export chat data: {e}")
167183

0 commit comments

Comments
 (0)