117 lines
3.5 KiB
Python
117 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
from typing import Self, Any, Sequence, Optional, Callable
|
|
from threading import Thread
|
|
from requests import post as Post, Response
|
|
from json import loads as json_decode
|
|
from re import compile as re_compile, Pattern as REPattern, IGNORECASE as RE_IGNORECASE
|
|
from websockets.sync import server as WebSocketServer
|
|
from Utils.Utils import Utils
|
|
from Utils.Check import Check
|
|
from Drivers.FilesDriver import FilesDriver
|
|
|
|
class PayloadModel:
|
|
def __init__(self:Self) -> None:
|
|
self.model:str = "gemma"
|
|
self.stream:bool = True
|
|
|
|
class AIChat:
|
|
|
|
DEFAULT_SETTINGS:dict[str, Any|None] = {
|
|
"autostart" : True,
|
|
"ai_model" : "gemma",
|
|
"ai_stream" : True
|
|
}
|
|
|
|
def __init__(self:Self, inputs:Optional[dict[str, Any|None]|Sequence[Any|None]] = None) -> None:
|
|
|
|
self.__inputs:dict[str, Any|None] = self.get_dictionary(inputs, overwrite = True)
|
|
self.__working:bool = False
|
|
self.__started:bool = False
|
|
self.__thread:Thread|None = None
|
|
self.__payload:PayloadModel = PayloadModel()
|
|
|
|
self.files:FilesDriver = FilesDriver(self)
|
|
|
|
def start(self:Self, callback:Optional[Callable[[bool], bool]] = None) -> bool:
|
|
|
|
end:Callable[[bool], bool] = lambda ok: callback(ok) if callable(callback) else ok
|
|
|
|
if self.__started:
|
|
return end(False)
|
|
self.__started = True
|
|
|
|
self.__working = True
|
|
self.__payload.model = self.get("ai_model", self.__payload.model)
|
|
self.__thread = Thread(target = self.__listener)
|
|
self.__thread.start()
|
|
|
|
return end(True)
|
|
|
|
def close(self:Self, callback:Optional[Callable[[bool], bool]] = None) -> bool:
|
|
|
|
end:Callable[[bool], bool] = lambda ok: callback(ok) if callable(callback) else ok
|
|
|
|
if not self.__started:
|
|
return end(False)
|
|
self.__started = False
|
|
|
|
self.__working = False
|
|
|
|
return end(True)
|
|
|
|
def load_json(self:Self, data:Any|None, only_dictionary:bool = True) -> list[list[Any|None]|dict[str, Any|None]]:
|
|
|
|
items:list[list[Any|None]|dict[str, Any|None]] = []
|
|
|
|
if Check.is_string(data):
|
|
items.extends(self.files.load_json(data))
|
|
elif Utils.is_array(data):
|
|
if only_dictionary:
|
|
|
|
item:Any|None
|
|
|
|
for item in data:
|
|
items.extend(self.load_json(item, only_dictionary))
|
|
|
|
else:
|
|
items.append(data)
|
|
elif Utils.is_dictionary(data):
|
|
items.append(data)
|
|
|
|
return items
|
|
|
|
def send(self:Self, message:str) -> str:
|
|
try:
|
|
|
|
response:Response
|
|
|
|
with Post('http://localhost:11434/api/chat', json = {**self.PAYLOAD, "prompt" : message}) as response:
|
|
|
|
line:bytes
|
|
|
|
for line in response.iter_lines():
|
|
if line:
|
|
|
|
chunk:dict[str, Any|None] = json_decode(line)
|
|
|
|
print(chunk.get("response", ""), end = "", flush = True)
|
|
|
|
if chunk.get("done"):
|
|
break
|
|
|
|
except Exception as exception:
|
|
print(f"An error occurred while sending the message: {exception}")
|
|
|
|
def __listener(self:Self) -> None:
|
|
while self.__working:
|
|
|
|
user_input:str = input('> ')
|
|
|
|
if user_input in ("close", "exit", "quit"):
|
|
self.__working = False
|
|
else:
|
|
self.send(user_input)
|
|
|
|
ai_chat = AIChat() |