#!/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 ollama import chat
import threading

class OllamaNamespace(Namespace):

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

    def on_connect(self):
        pass

    def on_disconnect(self):
        pass

    def on_chat(self, data):
        self.chat_lock.acquire()
        try:
            response = chat(model='qwen', stream=True, messages=[
                {
                    'role': 'user',
                    'content': data,
                },
            ])
            res = ''
            for part in response:
                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')
        self.chat_lock.release()