tools

src/tools/simple_chat/tool.py

@@ -0,0 +1,230 @@
+from typing import Dict, Callable, Awaitable
+from nicegui import ui, binding, app
+from tools.base_tool import BaseTool, BasePage
+from utils import ollama
+from typing import Literal, List, Optional, Any
+from datetime import datetime
+from dataclasses import dataclass
+from niceguiasyncelement import AsyncColumn
+import json
+
+
+class SimpleChatTool(BaseTool):
+    @property
+    def name(self) -> str:
+        return "Simple Chat"
+
+    @property
+    def description(self) -> str:
+        return "Simple Chat Tool example"
+
+    @property
+    def icon(self) -> str:
+        return "chat"
+
+    @property
+    def enabled(self) -> bool:
+        """Enable/disable this tool (set to False to hide from menu and disable routes)"""
+        return True  # Set to False to disable this tool
+
+    @property
+    def routes(self) -> Dict[str, Callable[[], Awaitable]]:
+        """Define the routes for this tool"""
+        return {
+            '': lambda: MainPage.create(self)
+        }
+
+
+@binding.bindable_dataclass
+class LLMMessage():
+    name: str
+    role: Literal['system', 'user', 'assistant']
+    content: str
+    stamp: datetime
+
+
+class ChatMessageComponent(LLMMessage, ui.column):
+    """Custom chat message component that supports data binding"""
+
+    def __init__(self, role: Literal['user', 'assistant'], name: str, content: str, timestamp: Optional[str] = None):
+        # Initialize LLMMessage
+        LLMMessage.__init__(self, name=name, role=role, content=content, stamp=datetime.now())
+        # Initialize ui.column
+        ui.column.__init__(self)
+
+        self.timestamp = timestamp or datetime.now().strftime('%H:%M:%S')
+        self.is_user = role == 'user'
+
+        self.classes('w-full mb-2')
+
+        with self:
+            # Message container with proper alignment
+            container_classes = 'w-full flex'
+            if self.is_user:
+                container_classes += ' justify-end'
+            else:
+                container_classes += ' justify-start'
+
+            with ui.row().classes(container_classes):
+                # Message bubble
+                bubble_classes = 'max-w-xs lg:max-w-md px-4 py-2 rounded-lg'
+                if self.is_user:
+                    bubble_classes += ' bg-cyan-600 text-white ml-auto'
+                else:
+                    bubble_classes += ' bg-gray-700 text-white mr-auto'
+
+                with ui.column().classes(bubble_classes):
+                    # Message content - bind to the content property
+                    ui.markdown().classes('text-sm whitespace-pre-wrap').bind_content_from(self, 'content')
+
+                    # Timestamp and role
+                    with ui.row().classes('items-center gap-2 mt-1'):
+                        ui.label(self.role.title()).classes('text-xs opacity-75 font-medium')
+                        ui.label(self.timestamp).classes('text-xs opacity-60')
+
+    def to_dict(self) -> dict:
+        return {
+            'name': self.name,
+            'role': self.role,
+            'content': self.content,
+            'stamp': self.stamp.isoformat(),
+            'timestamp': self.timestamp
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> 'ChatMessageComponent':
+        message = cls(
+            role=data['role'],
+            name=data['name'],
+            content=data['content'],
+            timestamp=data.get('timestamp')
+        )
+        message.name = data['name']
+        message.stamp = datetime.fromisoformat(data['stamp'])
+        return message
+
+
+class MainPage(BasePage):
+    """Main page of the example tool"""
+
+    chat_container: ui.scroll_area
+    messages_container: ui.column
+    model_selector: ui.select
+    chat_input: ui.input
+    history: List[ChatMessageComponent]
+
+    is_responding = binding.BindableProperty()
+
+    auto_scroll: bool
+    auto_scroll_timer: ui.timer
+
+    async def content(self):
+        self.history = []
+
+        self.auto_scroll = True
+        self.is_responding = False
+
+        model_options = [model['name'] for model in await ollama.available_models()]
+
+        with ui.row().classes('w-full'):
+            ui.label('Simple Chat').classes('text-2xl font-bold text-white mb-4')
+            ui.space()
+            ui.button('Clear Chat', on_click=self.clear_history)
+            self.model_selector = ui.select(model_options, label='Model', value=model_options[0]).props('outlined dense')
+            self.model_selector.bind_value(app.storage.user, f'{self.tool.name}_selected_model')
+
+        # Main chat layout - full width and height
+        with ui.column().classes('w-full gap-4 h-full'):
+            # Chat messages area - takes all available space
+            with ui.card().classes('w-full flex-1 p-0'):
+                self.chat_container = ui.scroll_area(on_scroll=lambda e: self.on_scroll_event(e.vertical_percentage)).classes('w-full').style('height: 100%; min-height: 500px')
+                self.chat_container.on('wheel', self.wheel_callback, ['deltaY'])
+                with self.chat_container:
+                    self.messages_container = ui.column().classes('w-full p-4 gap-2')  # Container for messages
+
+            # Input area at the bottom - fixed height
+            with ui.card().classes('w-full p-4').style('flex-shrink: 0'):
+                with ui.row().classes('w-full gap-2 items-center'):
+                    self.chat_input = ui.input(placeholder='Type your message...').classes('flex-1').props('outlined dense')
+                    self.chat_input.bind_enabled_from(self, 'is_responding', backward=lambda x: not x)
+                    self.chat_input.on('keydown.enter', self.send_message)
+                    ui.button(icon='send', on_click=self.send_message).props('color=primary').bind_enabled_from(self, 'is_responding', backward=lambda x: not x)
+
+        # Add example messages
+        await self.load_chat_history()
+
+        self.auto_scroll_timer = ui.timer(0.1, lambda: self.chat_container.scroll_to(percent=1))
+
+    async def add_message(self, role: Literal['user', 'assistant'], name: str, content: str):
+        with self.messages_container:
+            message_component = ChatMessageComponent(role, name, content)
+
+        self.history.append(message_component)
+
+    async def send_message(self):
+        """Send a user message from the input field"""
+        if self.chat_input.value and self.chat_input.value.strip():
+            user_message = self.chat_input.value.strip()
+            self.chat_input.value = ''  # Clear the input
+            await self.add_message('user', 'User', user_message)
+
+        # create data dict
+        self.is_responding = True
+        data = {
+            'model': self.model_selector.value,
+            'messages': []
+        }
+        for mes in self.history:
+            data['messages'].append({'role': mes.role, 'content': mes.content})
+
+        # create new empty message object for the response
+        await self.add_message('assistant', 'Assistant', '')
+
+        # generate streaming response
+        try:
+            async for chunk in ollama.stream_chat(data):
+                if chunk.strip():
+                    # Parse the JSON chunk and extract content
+                    try:
+                        chunk_data = json.loads(chunk)
+                        if 'message' in chunk_data and 'content' in chunk_data['message']:
+                            content = chunk_data['message']['content']
+                            self.history[-1].content += content
+                    except json.JSONDecodeError:
+                        pass  # Skip malformed chunks
+        except Exception as e:
+            ui.notify(f'Error: {str(e)}', type='negative')
+        finally:
+            await self.save_chat_history()
+            self.is_responding = False
+
+    async def load_chat_history(self):
+        if f'{self.tool.name}_history' in app.storage.user:
+            for mes in app.storage.user[f'{self.tool.name}_history']:
+
+                with self.messages_container:
+                    message_component = ChatMessageComponent.from_dict(mes)
+
+                self.history.append(message_component)
+
+    async def save_chat_history(self):
+        app.storage.user[f'{self.tool.name}_history'] = []
+        for mes in self.history:
+            app.storage.user[f'{self.tool.name}_history'].append(mes.to_dict())
+
+    async def clear_history(self):
+        app.storage.user[f'{self.tool.name}_history'] = []
+        self.history = []
+        self.messages_container.clear()
+        ...
+
+    def wheel_callback(self, event_data):
+        delta_y = event_data.args['deltaY']
+        if delta_y < 0:
+            if self.auto_scroll_timer.active:
+                self.auto_scroll_timer.deactivate()
+
+    def on_scroll_event(self, vertical_percentage):
+        if vertical_percentage == 1:
+            if not self.auto_scroll_timer.active:
+                self.auto_scroll_timer.activate()