@@ -22,7 +22,46 @@ def format(self, chat_data: dict[str, Any], image_dir: str = 'images') -> str:
22
22
pass
23
23
24
24
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 :
26
65
"""Format the chat data into Markdown format.
27
66
28
67
Args:
@@ -37,56 +76,41 @@ def format(self, chat_data: dict[str, Any], image_dir: str | None = 'images') ->
37
76
for tab_index , tab in enumerate (chat_data ['tabs' ]):
38
77
bubbles = tab ['bubbles' ]
39
78
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 )
43
79
44
80
for bubble in bubbles :
81
+ # USER
45
82
if bubble ['type' ] == 'user' :
46
83
user_text = ["## User:\n \n " ]
47
84
85
+ # Selections
48
86
if "selections" in bubble and bubble ["selections" ]:
49
87
user_text .append (f"[selections] \n { "\n " .join ([s ["text" ] for s in bubble ['selections' ]])} " )
50
88
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 :
52
91
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 " )
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 " )
61
100
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 )
83
106
if user_text_text :
84
107
user_text .append (f"[text] \n { user_text_text } " )
85
108
86
109
user_text .append ("\n " )
87
110
88
111
if len (user_text ) > 2 :
89
112
formatted_chat .append ("\n " .join (user_text ))
113
+ # AI
90
114
elif bubble ['type' ] == 'ai' :
91
115
model_type = bubble .get ('modelType' , 'Unknown' )
92
116
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') ->
96
120
97
121
logger .success ("Chats formatted." )
98
122
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 } " ]
108
123
except Exception as e :
109
124
logger .error (f"Unexpected error: { e } . Full traceback: { traceback .format_exc ()} " )
110
- return [ f"Error: { e } " ]
125
+ return
111
126
112
127
class FileSaver (ABC ):
113
128
@abstractmethod
@@ -159,9 +174,10 @@ def export(self, chat_data: dict[str, Any], output_dir: str, image_dir: str) ->
159
174
try :
160
175
os .makedirs (output_dir , exist_ok = True )
161
176
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 )
165
181
except Exception as e :
166
182
logger .error (f"Failed to export chat data: { e } " )
167
183
0 commit comments