Optimize docs, naming, UX etc.

README.md

@@ -377,7 +377,7 @@ Example:
 from askui import locators as loc
 
 password_textfield_label = loc.Text("Password")
-password_textfield = loc.Class("textfield").right_of(password_textfield_label)
+password_textfield = loc.Element("textfield").right_of(password_textfield_label)
 
 agent.click(password_textfield)
 agent.type("********")
@@ -414,14 +414,13 @@ Instead of taking a screenshot, you can analyze specific images:
 
 ```python
 from PIL import Image
-from askui.utils.image_utils import ImageSource
 
 # From PIL Image
 image = Image.open("screenshot.png")
-result = agent.get("What's in this image?", ImageSource(image))
+result = agent.get("What's in this image?", image)
 
 # From file path
-result = agent.get("What's in this image?", ImageSource("screenshot.png"))
+result = agent.get("What's in this image?", "screenshot.png")
 ```
 
 #### Using response schemas

src/askui/__init__.py

@@ -3,16 +3,19 @@
 __version__ = "0.2.4"
 
 from .agent import VisionAgent
-from .models.types import JsonSchemaBase
+from .models.router import ModelRouter
+from .models.types.response_schemas import ResponseSchema, ResponseSchemaBase
 from .tools.toolbox import AgentToolbox
 from .tools.agent_os import AgentOs, ModifierKey, PcKey
 
 
 __all__ = [
     "AgentOs",
     "AgentToolbox",
-    "JsonSchemaBase",
+    "ModelRouter",
     "ModifierKey",
     "PcKey",
+    "ResponseSchema",
+    "ResponseSchemaBase",
     "VisionAgent",
 ]

src/askui/chat/__main__.py

@@ -70,16 +70,21 @@ def write_message(
     role: str,
     content: str | dict | list,
     timestamp: str,
-    image: Image.Image |str | None = None,
+    image: Image.Image | str | list[str | Image.Image] | list[str] | list[Image.Image] | None = None,
 ):
     _role = ROLE_MAP.get(role.lower(), UNKNOWN_ROLE)
     avatar = None if _role != UNKNOWN_ROLE else "❔"
     with st.chat_message(_role, avatar=avatar):
         st.markdown(f"*{timestamp}* - **{role}**\n\n")
         st.markdown(json.dumps(content, indent=2) if isinstance(content, (dict, list)) else content)
         if image:
-            img = get_image(image) if isinstance(image, str) else image
-            st.image(img)
+            if isinstance(image, list):
+                for img in image:
+                    img = get_image(img) if isinstance(img, str) else img
+                    st.image(img)
+            else:
+                img = get_image(image) if isinstance(image, str) else image
+                st.image(img)
 
 
 def save_image(image: Image.Image) -> str:
@@ -93,21 +98,29 @@ class Message(TypedDict):
     role: str
     content: str | dict | list
     timestamp: str
-    image: str | None
+    image: str | list[str] | None
 
 
 class ChatHistoryAppender(Reporter):
     def __init__(self, session_id: str) -> None:
         self._session_id = session_id
 
     @override
-    def add_message(self, role: str, content: Union[str, dict, list], image: Image.Image | None = None) -> None:
-        image_path = save_image(image) if image else None
+    def add_message(self, role: str, content: Union[str, dict, list], image: Image.Image | list[Image.Image] | None = None) -> None:
+        image_paths: list[str] = []
+        if image is None:
+            _images = []
+        elif isinstance(image, list):
+            _images = image
+        else:
+            _images = [image]
+        for img in _images:
+            image_paths.append(save_image(img))
         message = Message(
             role=role,
             content=content,
             timestamp=datetime.now().isoformat(),
-            image=image_path,
+            image=image_paths,
         )
         write_message(**message)
         with open(

src/askui/locators/__init__.py

@@ -1,9 +1,9 @@
-from askui.locators.locators import AiElement, Class, Description, Image, Text
+from askui.locators.locators import AiElement, Element, Prompt, Image, Text
 
 __all__ = [
     "AiElement",
-    "Class",
-    "Description",
+    "Element",
+    "Prompt",
     "Image",
     "Text",
 ]