a general purpose chat client made with Flet
import flet as ft
from flet_core.types import AppView

from settings.settings import APP_CONFIG
from stores.chat_store import ChatStore
from stores.user_store import UserStore
from stores.ai_store import AIStore
from views.chat import ChatView


def main(page: ft.Page):
    # add storages
    page.user_store = UserStore(page=page)
    page.chat_store = ChatStore(page=page)
    page.ai_store = AIStore(page=page)

    page.title = APP_CONFIG["title"]
    page.theme_mode = ft.ThemeMode.DARK

    # FIRST APP LAUNCH
    # 1. create new chat
    first_chat = page.chat_store.new_chat()
    # 2. start chatting
    page.views.append(
        ChatView(
            page=page,
            chat=first_chat,
            user=page.user_store.user,
        )
    )

    # first render
    page.update()


ft.app(
    name=APP_CONFIG["name"],
    view=AppView.FLET_APP,
    assets_dir=APP_CONFIG["assets_dir"],
    upload_dir=APP_CONFIG["upload_dir"],
    use_color_emoji=True,
    target=main,
),

app.py
yudame/flet_chat_app
import flet as ft
import json
import os

class AppBar(ft.AppBar):
    def __init__(self, page):
        super().__init__()
        self.page = page
        self.leading = ft.Icon(ft.icons.HEALTH_AND_SAFETY_ROUNDED)
        self.leading_width = 60
        self.title = self.get_title()
        self.center_title = False
        self.bgcolor = ft.colors.BLACK54
        self.actions = self.get_actions()

    def get_title(self):
        topics = list(json.loads(os.environ.get("AI_CONFIG", "{}")).get("chat_topics", {}).keys())
        return ft.Row(
            [
                ft.TextButton(
                    topic_category.title(), on_click=lambda _: print(f"talk about {str(topic_category)}")
                ) for topic_category in topics
            ]
        )

    def get_actions(self):
        def get_actions():
            return (
                [
                    # ft.IconButton(ft.icons.WB_SUNNY_OUTLINED),  # to change theme light/dark
                    ft.IconButton(
                        ft.icons.RESTART_ALT_ROUNDED,
                        on_click=lambda x: print("reset app"),
                    ),
                ],
            )

components/appbar.py
yudame/flet_chat_app
import flet as ft

from stores.chat_store import Chat, Message
from stores.user_store import User


class ChatHistory(ft.Container):
    messages_list = ft.ListView(expand=1, spacing=10, padding=20, auto_scroll=True)

    def __init__(self, page: ft.Page, chat: Chat, user: User):
        super().__init__()
        self.page = page
        self.chat = chat
        self.user = user
        self.load_messages()

        self.content = self.messages_list
        self.expand = 1

    def load_messages(self) -> None:
        # init with first message if empty
        if not len(self.chat.messages):
            self.chat.add_message(
                author_user=self.page.ai_store.ai_user,
                message_text=self.page.ai_store.ai_role.first_message,
            )
        # add all messages to the screen
        for message in self.chat.messages:
            self.add_message(message)
        self.page.update()

    def add_message(self, message: Message, update_page: bool = True) -> None:
        if message.author is self.user:
            self.messages_list.controls.append(
                ft.Text(
                    message.message,
                    text_align=ft.TextAlign.RIGHT,
                    color="blue",
                )
            )
        else:
            self.messages_list.controls.append(
                ft.Text(
                    message.message,
                    text_align=ft.TextAlign.LEFT,
                )
            )

        if update_page:
            self.page.update()

components/chat_history.py
yudame/flet_chat_app
import flet as ft

from components.chat_history import ChatHistory
from stores.chat_store import Message


class ChatInput(ft.Container):
    text_field = ft.TextField(
        multiline=True,
        min_lines=1,
        max_lines=7,
        expand=True,
        border_color=ft.colors.TRANSPARENT,
        filled=True,
        border_radius=ft.BorderRadius(8, 8, 8, 8),
        shift_enter=True,
    )

    send_button = ft.IconButton(ft.icons.SEND_ROUNDED)

    def __init__(self, page: ft.Page, chat_history: ChatHistory):
        super().__init__()
        self.page = page

        self.rail_button = ft.IconButton(ft.icons.MENU_ROUNDED)
        self.chat_history_container = chat_history
        self.author_user = chat_history.user

        self.text_field.on_submit = self.send_message
        self.send_button.on_click = self.send_message
        self.content = ft.Row(
            spacing=4,
            controls=[
                self.rail_button,
                self.text_field,
                self.send_button,
            ],
        )

    def send_message(self, e: ft.ControlEvent) -> None:
        user_message_text: str = self.text_field.value
        self.text_field.value = ""
        self.lock_input()

        # add to local datastore
        user_message: Message = self.chat_history_container.chat.add_message(
            author_user=self.author_user, message_text=user_message_text
        )
        # display on screen
        self.chat_history_container.add_message(user_message)

        # get response text from ai
        self.page.ai_store.get_next_message(chat=self.chat_history_container.chat)
        # display ai response on screen
        if self.chat_history_container.chat.messages[-1].author is not self.author_user:
            ai_message: Message = self.chat_history_container.chat.messages[-1]
            self.chat_history_container.add_message(ai_message)

        # unlock the user
        self.unlock_input()

    def lock_input(self, page_update=True):
        self.text_field.disabled = True
        self.send_button.disabled = True
        if page_update:
            self.page.update()

    def unlock_input(self, page_update=True):
        self.text_field.disabled = False
        self.send_button.disabled = False
        if page_update:
            self.page.update()

components/chat_input.py
yudame/flet_chat_app
from dataclasses import dataclass
from enum import Enum
from typing import Dict, List, Optional, Union


@dataclass(frozen=True, slots=True)
class AIRole:
    name: str
    title: str
    intro: str
    first_message: str


RecursiveDict = Dict[str, Union[str, "RecursiveDict"]]


class LLMStepName(Enum):
    PRACTICE = "practice"
    RESPONSE = "response"


@dataclass(frozen=True, slots=True)
class LLMStep:
    name: Union[str, LLMStepName]
    instruction: str
    cheat_sheet: Optional[Union[str, RecursiveDict]]


class AIConfig:
    __slots__ = ["openai_api", "ai_role", "llm_steps"]

    def __init__(
        self,
        openai_api: Dict[str, str],
        ai_role: AIRole,
        llm_steps: List[LLMStep],
    ) -> object:
        self.openai_api = openai_api
        self.ai_role = ai_role
        self.llm_steps = llm_steps

settings/ai_datatypes.py
yudame/flet_chat_app
from typing import Dict

APP_CONFIG: Dict = {
    "name": "Not A Therapist App by Official",
    "title": "AI Chat Demo",
    "assets_dir": "assets",
    "upload_dir": "assets/uploads",
}

settings/settings.py
yudame/flet_chat_app
from typing import List, Optional

from icecream import ic
from langchain.chat_models import ChatOpenAI

from settings.ai_datatypes import LLMStep, LLMStepName
from settings.ai_config import AI_CONFIG
from stores.base_store import BaseStore
from stores.chat_store import Chat
from stores.user_store import User
from utils import count_tokens, dict_to_cheat_sheet


class AIStore(BaseStore):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # self.page is provided by the BaseStore class
        self._validate_config()
        self.planning_llm = ChatOpenAI(
            openai_api_key=self.openai_api_key,
            openai_organization=self.openai_api_org,
            model_name=self.open_ai_planning_model,
            temperature=0.1,
        )
        self.chat_llm = ChatOpenAI(
            openai_api_key=self.openai_api_key,
            openai_organization=self.openai_api_org,
            model_name=self.openai_api_chat_model,
            temperature=0.7,
        )
        self.ai_role = AI_CONFIG.ai_role
        self.ai_user = User(id=AI_CONFIG.ai_role.name, name=AI_CONFIG.ai_role.name)
        self.system_user = User(id="system", name="system")
        self.conversation_start_message_list = [
            {"role": "system", "content": AI_CONFIG.ai_role.intro},
        ]

    def _validate_config(self) -> None:
        assert (
            AI_CONFIG.llm_steps
        ), "openai_api_config.llm_steps is missing required key 'response'"

        self.openai_api_key: str = AI_CONFIG.openai_api.get("api_key")
        self.openai_api_org: str = AI_CONFIG.openai_api.get("org_id")
        default_model: str = AI_CONFIG.openai_api.get("default_model")
        self.openai_api_chat_model: str = (
            AI_CONFIG.openai_api.get("chat_model") or default_model
        )
        self.open_ai_planning_model: str = (
            AI_CONFIG.openai_api.get("planning_model") or default_model
        )

    def prompt(
        self,
        chat: Chat,
        system_prompt: str,
        mode: str = "chat",
        max_tokens: Optional[int] = 90,  # 90 is about 3 long sentences
    ) -> str:
        messages_list: List = []
        # load summary if exists
        if chat.summary:
            messages_list.append(
                {
                    "role": "system",
                    "content": f"Here is a summary of conversation until now: {chat.summary}",
                }
            )
        else:
            messages_list += self.conversation_start_message_list
        # add all message history until now
        messages_list += chat.get_history_as_message_list()
        # add next step
        messages_list.append({"role": "system", "content": system_prompt})

        # convert to text
        full_message_text: str = "\n".join(
            [f"{message['role']}: {message['content']}" for message in messages_list]
        )
        total_tokens = count_tokens(full_message_text, self.openai_api_chat_model)
        if total_tokens > 2048:
            # todo: Handle token overflow
            return "Token limit exceeded"

        match mode:
            case "planning":
                ai_text_response = self.planning_llm.predict(
                    full_message_text, max_tokens=max_tokens
                )
            case "chat":
                ai_text_response = self.chat_llm.predict(
                    full_message_text, max_tokens=max_tokens
                )
            case _:
                raise ValueError(f"Invalid mode: {mode}")

        ai_text_response = ai_text_response.split(f"{AI_CONFIG.ai_role.name}: ")[-1]
        ai_text_response = ai_text_response.split(f"{AI_CONFIG.ai_role.title}: ")[-1]
        return ai_text_response

    def _generate_system_prompt(self, llm_step: LLMStep) -> str:
        system_prompt = f'Now you, {self.ai_user.name}, take a moment to do a "{llm_step.name}" step. {llm_step.instruction}'
        if llm_step.cheat_sheet:
            cheat_sheet_str = (
                llm_step.cheat_sheet
                if isinstance(llm_step.cheat_sheet, str)
                else dict_to_cheat_sheet(llm_step.cheat_sheet)
            )
            system_prompt += f"Use this reference cheat sheet:\n\n{cheat_sheet_str}\n\n"
        if llm_step.name not in (LLMStepName.PRACTICE, LLMStepName.RESPONSE):
            system_prompt += (
                "Be concise. Write a note to yourself using only a few words."
            )
        return system_prompt

    def _execute_planning_step(
        self, chat: Chat, llm_step: LLMStep, max_tokens: int
    ) -> None:
        system_prompt = self._generate_system_prompt(llm_step)

        ai_text_response = self.prompt(
            chat=chat,
            system_prompt=system_prompt,
            mode="planning",
            max_tokens=max_tokens,
        )

        ic(ai_text_response)
        chat.add_message(
            self.system_user, f"your notes on {llm_step.name}:\n{ai_text_response}\n"
        )

    def _take_planning_steps(self, chat: Chat):
        for llm_step in AI_CONFIG.llm_steps:
            if llm_step.name == LLMStepName.RESPONSE:  # should be the last one
                continue
            max_tokens = 120 if llm_step.name == LLMStepName.PRACTICE else 30
            self._execute_planning_step(chat, llm_step, max_tokens)

    def get_next_message(self, chat: Chat):
        self._take_planning_steps(chat=chat)

        response_step = next(
            (step for step in AI_CONFIG.llm_steps if step.name == LLMStepName.RESPONSE),
            None,
        )
        final_instruction = (
            response_step.instruction
            if response_step
            else "take a deep breath and prepare your best response"
        )

        ai_text_response = self.prompt(
            chat=chat, system_prompt=final_instruction, mode="chat"
        )

        chat.add_message(self.ai_user, ai_text_response)
        return ai_text_response

stores/ai_store.py
yudame/flet_chat_app
from abc import ABC
import flet as ft


class BaseStore(ABC):
    def __init__(self, page: ft.Page, *args, **kwargs):
        self.page: ft.Page = page

stores/base_store.py
yudame/flet_chat_app
from dataclasses import asdict, dataclass
import json
from typing import List, Set, Dict, Optional, Tuple, AbstractSet
from datetime import datetime
import flet as ft
import uuid

from stores.base_store import BaseStore
from stores.user_store import User


@dataclass(frozen=True, slots=True)
class Message:
    author: User
    message: str
    timestamp: datetime = datetime.utcnow()

    @property
    def serialized(self):
        return json.dumps(asdict(self))


class Chat:
    __slots__ = ["page", "id", "member_user_ids", "messages", "summary"]

    def __init__(self, page: ft.Page):
        self.page: ft.Page = page
        self.id: str = str(uuid.uuid4())
        self.member_user_ids: Set[str] = set()
        self.messages: List[Message] = []
        self.summary: str = ""

    def add_message(
        self,
        author_user: User,
        message_text: str,
        timestamp: Optional[datetime] = None,
    ) -> Message:
        message = Message(
            author=author_user,
            message=message_text,
            timestamp=timestamp or datetime.utcnow(),
        )
        self.member_user_ids.add(str(author_user.id))
        self.messages.append(message)
        return message

    def get_history_as_message_list(self) -> List[Dict[str, str]]:
        return [{"role": m.author.name, "content": m.message} for m in self.messages]

    def get_history_as_text(self) -> str:
        return "\n".join([f"{m.author.name}: {m.message}" for m in self.messages])

    def get_users(self) -> AbstractSet[str]:
        return frozenset(self.member_user_ids)

    def get_new_ai_message(self) -> Message:
        return self.add_message(author_user=self.page.ai_store.ai_user, message_text="")

    @property
    def serialized(self):
        return json.dumps(
            {
                "id": self.id,
                "member_user_ids": [user.serialized for user in self.member_user_ids],
                "messages": [message.serialized for message in self.messages],
            }
        )


class ChatStore(BaseStore):
    __slots__ = ["chats"]

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        # self.page is provided by the BaseStore class

        self.chats: Dict[str, Chat] = dict()
        # load chats from client storage
        chats_serialized: str = self.page.client_storage.get("user_contacts")
        if chats_serialized:
            self.chats: Dict[str, Chat] = {
                chat_dict["id"]: Chat(**chat_dict)
                for chat_dict in json.loads(chats_serialized)
            }

    def new_chat(self) -> Chat:
        chat = Chat(page=self.page)
        self.add_chat(chat)
        return chat

    def add_chat(self, chat: Chat) -> None:
        self.chats[chat.id] = chat
        # self.page.client_storage.set(
        #     "chats", json.dumps({chat.id: chat.serialized for chat in self.chats})
        # )

    def get_chat(self, chat_id: str) -> Chat:
        try:
            return self.chats[chat_id]
        except KeyError:
            raise ValueError(f"No chat found for id: {chat_id}")

    def get_chat_ids(self) -> List[str]:
        return list(self.chats.keys())

stores/chat_store.py
yudame/flet_chat_app
from dataclasses import asdict, dataclass, replace
from datetime import datetime
import json
import logging
from typing import Optional, Set
import uuid

from stores.base_store import BaseStore


@dataclass(slots=True)
class User:
    id: str
    name: str = "Human"
    handle: Optional[str] = None
    created_at: datetime = datetime.utcnow()

    @property
    def serialized(self):
        return json.dumps(asdict(self))


class UserStore(BaseStore):
    __slots__ = ["page", "user", "contacts"]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # self.page is provided by the BaseStore class

        # load user info from client storage
        self.user = None
        user_serialized: str = self.page.client_storage.get("user")
        if user_serialized:
            try:
                self.user = User(**json.loads())
                self.page.client_storage.set("user", self.user.serialized)
            except:
                logging.exception(
                    f"Failed to load user. user_serialized: {user_serialized}"
                )
        if not self.user:
            self.user = User(id=str(uuid.uuid4()))

        # load contacts from client storage
        self.user_contacts: Set[User] = set()
        user_contacts_serialized: str = self.page.client_storage.get("user_contacts")
        if user_contacts_serialized:
            self.user_contacts: Set[User] = set(
                [
                    User(**user_dict)
                    for user_dict in json.loads(user_contacts_serialized)
                ]
            )

    def update_user_attr(self, attr, value):
        self.user = replace(self.user, **{attr: value})
        self.page.client_storage.set("user", self.user.serialized)
        self.page.update()

stores/user_store.py
yudame/flet_chat_app
import tiktoken


def dict_to_cheat_sheet(data: dict, depth: int = 1) -> str:
    """Convert a nested dictionary to a markdown cheat sheet"""
    markdown_str = ""

    for key, value in data.items():
        if depth == 1:
            markdown_str += f"# {key}\n\n"
        else:
            markdown_str += f"{'  ' * (depth - 2)}- {key}\n"

        if isinstance(value, dict):
            markdown_str += dict_to_cheat_sheet(value, depth + 1)
        else:
            markdown_str += f"{'  ' * (depth - 1)}- {value}\n"

    return markdown_str


def print_markdown(md_str: str):
    for line in md_str.split("\n"):
        if line.startswith("#"):
            print("-" * 50)
        print(line)


def count_tokens(text: str, model_name: str) -> int:
    encoding = tiktoken.encoding_for_model(model_name)
    token_count = len(encoding.encode(text))
    return token_count


def strip_text_fragments(text: str, fragments: set[str]) -> str:
    for fragment in fragments:
        prefix = f"{fragment}:"
        if text.startswith(prefix):
            return text[len(prefix) :].lstrip()

    return text

utils.py
yudame/flet_chat_app
import flet as ft

from components.appbar import AppBar
from components.chat_history import ChatHistory
from components.chat_input import ChatInput
from stores.chat_store import Chat
from stores.user_store import User


class ChatView(ft.View):
    def __init__(self, page: ft.Page, chat: Chat, user: User):
        super().__init__()
        self.route = "/chat"
        self.chat = chat
        self.user = user

        # COMPONENTS
        chat_history = ChatHistory(page=page, chat=chat, user=user)
        chat_input = ChatInput(page=page, chat_history=chat_history)

        # UI
        self.appbar = AppBar(page=page)
        self.controls = [
            chat_history,
            chat_input,
        ]

views/chat.py
yudame/flet_chat_app