53 lines
1.4 KiB
Python
53 lines
1.4 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
from typing import Self, Any
|
|
from threading import Thread
|
|
from requests import post as Post, Response
|
|
from json import loads as json_decode
|
|
|
|
class AIChat:
|
|
|
|
PAYLOAD:dict[str, str|bool] = {
|
|
"model" : "gemma",
|
|
"stream" : True
|
|
}
|
|
|
|
def __init__(self:Self) -> None:
|
|
self.__working:bool = True
|
|
self.__thread:Thread = Thread(target=self.__listener)
|
|
self.__thread.start()
|
|
|
|
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() |