@@ -22,7 +22,46 @@ def format(self, chat_data: dict[str, Any], image_dir: str = 'images') -> str:
2222 pass
2323
2424class 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 " )
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 " )
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
112127class 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