#!/usr/bin/env python
# encoding: utf-8
"""
@author: simin. Gong
@license: (C) Copyright 2013-2023 Spacemit Corporation Limited.
@contact: simin.gong@spacemit.com
@file: LLMNamespace.py
@date: 2024/11/28 09:25 
@desc:
"""

from flask_socketio.namespace import Namespace
from ai_tool.model_manager.ollama_model_manager import ModelManager # type: ignore
from ollama import chat
import threading


class OllamaNamespace(Namespace):

    def __init__(self, namespace=None):
        super(OllamaNamespace, self).__init__(namespace)
        self.chat_lock = threading.Lock()
        self.stop_chat = False

    def on_connect(self):
        pass

    def on_disconnect(self):
        pass

    def on_chat(self, data):
        self.stop_chat = False
        self.chat_lock.acquire()
        try:
            self.emit('chat_start')
            response = chat(model=self.model_name, stream=True, messages=[
                {
                    'role': 'user',
                    'content': data,
                },
            ])
            res = ''
            for part in response:
                if self.stop_chat:  # 检查停止对话标志位
                    self.emit('chat_over', 'stop')
                    break
                res += part['message']['content']
                self.emit('chat_response', res)
        except Exception as e:
            print(e)
            self.emit('chat_error', f'{e}')
        finally:
            self.emit('chat_over', 'over')
        self.chat_lock.release()
        
    def on_select_model(self, data):
        self.model_name = data 
        
    def on_stop_chat(self):
        self.stop_chat = True
        
    def on_get_local_model_list(self):
        model_manager = ModelManager()
        model_manager.update_local_model_list()
        self.emit('modellist_response', model_manager.local_model_list)