1- import glob
21import json
32import logging
4- import os
53import re
6- from datetime import datetime
4+ from datetime import datetime , timezone
5+ from pathlib import Path
76from random import randint
87from typing import Union
98
1312
1413from askui import VisionAgent
1514from askui .chat .click_recorder import ClickRecorder
15+ from askui .chat .exceptions import FunctionExecutionError , InvalidFunctionError
1616from askui .models import ModelName
1717from askui .reporting import Reporter
1818from askui .utils .image_utils import base64_to_image , draw_point_on_image
2323)
2424
2525
26- CHAT_SESSIONS_DIR_PATH = "./chat/sessions"
27- CHAT_IMAGES_DIR_PATH = "./chat/images"
26+ CHAT_SESSIONS_DIR_PATH = Path ( "./chat/sessions" )
27+ CHAT_IMAGES_DIR_PATH = Path ( "./chat/images" )
2828
2929click_recorder = ClickRecorder ()
3030
3131
32- def setup_chat_dirs ():
33- os . makedirs (CHAT_SESSIONS_DIR_PATH , exist_ok = True )
34- os . makedirs (CHAT_IMAGES_DIR_PATH , exist_ok = True )
32+ def setup_chat_dirs () -> None :
33+ Path . mkdir (CHAT_SESSIONS_DIR_PATH , parents = True , exist_ok = True )
34+ Path . mkdir (CHAT_IMAGES_DIR_PATH , parents = True , exist_ok = True )
3535
3636
37- def get_session_id_from_path (path ):
38- return os .path .splitext (os .path .basename (path ))[0 ]
37+ def get_session_id_from_path (path : str ) -> str :
38+ """Get session ID from file path."""
39+ return Path (path ).stem
3940
4041
41- def load_chat_history (session_id ) :
42- messages = []
43- session_path = os . path . join ( CHAT_SESSIONS_DIR_PATH , f" { session_id } .jsonl" )
44- if os . path . exists ( session_path ):
45- with open ( session_path , "r" ) as f :
46- for line in f :
47- messages .append (json .loads (line ))
42+ def load_chat_history (session_id : str ) -> list [ dict ] :
43+ """Load chat history for a given session ID."""
44+ messages : list [ dict ] = []
45+ session_path = CHAT_SESSIONS_DIR_PATH / f" { session_id } .jsonl"
46+ if session_path . exists () :
47+ with session_path . open ( "r" ) as f :
48+ messages .extend (json .loads (line ) for line in f )
4849 return messages
4950
5051
@@ -60,7 +61,8 @@ def load_chat_history(session_id):
6061
6162
6263def get_image (img_b64_str_or_path : str ) -> Image .Image :
63- if os .path .isfile (img_b64_str_or_path ):
64+ """Get image from base64 string or file path."""
65+ if Path (img_b64_str_or_path ).is_file ():
6466 return Image .open (img_b64_str_or_path )
6567 return base64_to_image (img_b64_str_or_path )
6668
@@ -75,7 +77,7 @@ def write_message(
7577 | list [str ]
7678 | list [Image .Image ]
7779 | None = None ,
78- ):
80+ ) -> None :
7981 _role = ROLE_MAP .get (role .lower (), UNKNOWN_ROLE )
8082 avatar = None if _role != UNKNOWN_ROLE else "❔"
8183 with st .chat_message (_role , avatar = avatar ):
@@ -96,10 +98,11 @@ def write_message(
9698
9799
98100def save_image (image : Image .Image ) -> str :
99- timestamp = datetime .now ().strftime ("%Y%m%d_%H%M%S_%f" )
100- image_path = os .path .join (CHAT_IMAGES_DIR_PATH , f"image_{ timestamp } .png" )
101+ """Save image to disk and return path."""
102+ timestamp = datetime .now (tz = timezone .utc ).strftime ("%Y%m%d_%H%M%S_%f" )
103+ image_path = CHAT_IMAGES_DIR_PATH / f"image_{ timestamp } .png"
101104 image .save (image_path )
102- return image_path
105+ return str ( image_path )
103106
104107
105108class Message (TypedDict ):
@@ -127,18 +130,15 @@ def add_message(
127130 _images = image
128131 else :
129132 _images = [image ]
130- for img in _images :
131- image_paths .append (save_image (img ))
133+ image_paths .extend (save_image (img ) for img in _images )
132134 message = Message (
133135 role = role ,
134136 content = content ,
135- timestamp = datetime .now ().isoformat (),
137+ timestamp = datetime .now (tz = timezone . utc ).isoformat (),
136138 image = image_paths ,
137139 )
138140 write_message (** message )
139- with open (
140- os .path .join (CHAT_SESSIONS_DIR_PATH , f"{ self ._session_id } .jsonl" ), "a"
141- ) as f :
141+ with (CHAT_SESSIONS_DIR_PATH / f"{ self ._session_id } .jsonl" ).open ("a" ) as f :
142142 json .dump (message , f )
143143 f .write ("\n " )
144144
@@ -147,17 +147,18 @@ def generate(self) -> None:
147147 pass
148148
149149
150- def get_available_sessions ():
151- session_files = glob .glob (os .path .join (CHAT_SESSIONS_DIR_PATH , "*.jsonl" ))
150+ def get_available_sessions () -> list [str ]:
151+ """Get list of available session IDs."""
152+ session_files = list (CHAT_SESSIONS_DIR_PATH .glob ("*.jsonl" ))
152153 return sorted ([get_session_id_from_path (f ) for f in session_files ], reverse = True )
153154
154155
155156def create_new_session () -> str :
156- timestamp = datetime .now ().strftime ("%Y%m%d%H%M%S%f" )
157+ """Create a new chat session."""
158+ timestamp = datetime .now (tz = timezone .utc ).strftime ("%Y%m%d%H%M%S%f" )
157159 random_suffix = f"{ randint (100 , 999 )} "
158160 session_id = f"{ timestamp } { random_suffix } "
159- with open (os .path .join (CHAT_SESSIONS_DIR_PATH , f"{ session_id } .jsonl" ), "w" ) as f :
160- pass
161+ (CHAT_SESSIONS_DIR_PATH / f"{ session_id } .jsonl" ).touch ()
161162 return session_id
162163
163164
@@ -200,7 +201,7 @@ def paint_crosshair(
200201"""
201202
202203
203- def rerun ():
204+ def rerun () -> None :
204205 st .markdown ("### Re-running..." )
205206 with VisionAgent (
206207 log_level = logging .DEBUG ,
@@ -220,9 +221,8 @@ def rerun():
220221 r"mouse\((\d+),\s*(\d+)\)" , message ["content" ]
221222 ):
222223 if not screenshot :
223- raise ValueError (
224- "Screenshot is required to paint crosshair"
225- )
224+ error_msg = "Screenshot is required to paint crosshair"
225+ raise ValueError (error_msg ) # noqa: TRY301
226226 x , y = map (int , match .groups ())
227227 screenshot_with_crosshair = paint_crosshair (
228228 screenshot , (x , y )
@@ -235,7 +235,7 @@ def rerun():
235235 write_message (
236236 message ["role" ],
237237 f"Move mouse to { element_description } " ,
238- datetime .now ().isoformat (),
238+ datetime .now (tz = timezone . utc ).isoformat (),
239239 image = screenshot_with_crosshair ,
240240 )
241241 agent .mouse_move (
@@ -246,17 +246,17 @@ def rerun():
246246 write_message (
247247 message ["role" ],
248248 message ["content" ],
249- datetime .now ().isoformat (),
249+ datetime .now (tz = timezone . utc ).isoformat (),
250250 message .get ("image" ),
251251 )
252252 func_call = f"agent.tools.os.{ message ['content' ]} "
253253 eval (func_call )
254254 except json .JSONDecodeError :
255255 continue
256256 except AttributeError :
257- st .write (f"Invalid function: { message [' content' ] } " )
258- except Exception as e :
259- st .write (f"Error executing { message [' content' ] } : { str ( e ) } " )
257+ st .write (str ( InvalidFunctionError ( message [" content" ])) )
258+ except Exception as e : # noqa: BLE001 - We want to catch all other exceptions here
259+ st .write (str ( FunctionExecutionError ( message [" content" ], e )) )
260260
261261
262262setup_chat_dirs ()
0 commit comments