main.py

import ollama
import logging
import json

from datetime import datetime

from yaspin import yaspin

from speak.speak import say

CHAT_MODEL = 'mistral'
IMAGE_MODEL = 'llava'

SYSTEM_PROMPT_PATH = "system_promt.txt"


def add_to_calendar(title: str, date: datetime = datetime.now()):
    print(f"adding -- {title} -- to calendar")


class TaskEngine:
    
    def __init__(self):
        self.task_list = [
            "add_to_calendar"
        ]

    def is_valid_task(self, task) -> bool:
        return task in self.task_list

    def run_task(self, task: str, parameters: dict):
        if self.is_valid_task(task):
            if task == "add_to_calendar":
                add_to_calendar(parameters["title"])
            
            # exec(task, parameters)


class Message:
    role: str
    content: str

    def __init__(self, role, content):
        self.role = role
        self.content = content
        
        self.log = logging.getLogger(__name__)

    @yaspin(text='handle ai answer')
    def handle_content(self) -> str:
        try:
            item: dict = json.loads(self.content)
    
            if item["action"] == "command":
                self.run_command(item)
                return item["comment"]
            
            if item["action"] == "query":
                return "running query"
            
            if item["action"] == "answer":
                return item["answer"]
            
            if item["action"] == "clarify":
                return item["question"]
            
            return self.content
        except ValueError as error:
            print(error)
            self.log.error("LLM return non json content")
    
    def run_command(self, item: dict):
        self.log.info("Starting task engine")
        TaskEngine().run_task(item["target"], item["parameters"])
    
    def query(self, item: dict):
        pass
        
    def get_role(self) -> str:
        return self.role

    def get_content(self) -> str:
        return self.content

    def is_fromatted(self) -> bool:
        return "{" and "}" in self.content

    def dict(self):
        return {'role': self.role, 'content': self.content}


class ChatSession:
    
    def __init__(self, session_name):
        self.chat_history = []
        self.sessionName = session_name
        
        self.log = logging.getLogger(__name__)
        
    def get_histroy(self) -> list:
        return [message.dict() for message in self.chat_history]

    def add_user_message(self, message: str):
        self.chat_history.append(Message('user', message))
        
    def add_system_message(self, message: str):
        self.chat_history.append(Message('system', message))
        
    def add_assistant_message(self, message: str):
        self.chat_history.append(Message('assistant', message))
        
    def pretty_print_histroy(self):        
        for message in self.chat_history:
            print("\n")
            print("\033[91m {}\033[00m" .format(message.get_role() + ": "))
            print(message.get_content())


class ChatService:
    def __init__(self, session: ChatSession, system_message: str,  model: str = CHAT_MODEL):
        self.initialized = False
        self.model = model
        self.session = session
        
        self.log = logging.getLogger(__name__)
        
        self.initialize_chat(system_message)
    
    @yaspin(text='init chat service')
    def initialize_chat(self, system_message):
        self.session.add_system_message(system_message)
        
        message_history = self.session.get_histroy()
        response = ollama.chat(model=self.model, messages=message_history)
        self.session.add_assistant_message(response['message']['content'])
        
        self.test()
        self.initialized = True
        self.log.info("Chat service initialized successfully")
    
    @yaspin(text='talking to ai')
    def chat(self, message: str,) -> str:
        self.session.add_user_message(message)
        
        message_history = self.session.get_histroy()
        response = ollama.chat(model=self.model, messages=message_history)
        
        self.session.add_assistant_message(response['message']['content'])
        
        return self.session.chat_history[-1].handle_content()
    
    def test(self):
        pass


def build_system_promt() -> str:
    file = open(SYSTEM_PROMPT_PATH, 'r')
    promt = file.read()
    
    promt += f"\n I am located in Bavaria, Germany. Current time stamp is: {datetime.now()}"
    promt += "\n Your reponse should be the JSON and no other text."
    
    return promt
    
    
def main():    
    log = logging.getLogger(__name__)
    
    system_promt = build_system_promt()
    
    session = ChatSession("main")
    chat = ChatService(session, system_promt)

    while True:
        user_input = input(">")
        
        if user_input == "history()":
            session.pretty_print_histroy()
            
        else:
            answer = chat.chat(user_input)
            print(answer)
            say(answer)

          
if __name__ == "__main__":
    try:
        main()
    except (KeyboardInterrupt, EOFError):
        ...