mobees-ai/chatbot/main.py

import json
import os
import sys
import logging
from datetime import datetime, timedelta
from io import BytesIO
from django.utils.text import slugify
from django.conf import settings

import openai
from llama_index import VectorStoreIndex, Document, StorageContext, load_index_from_storage
from llama_index.llms import ChatMessage, MessageRole
from llama_index.prompts  import ChatPromptTemplate
# from langchain.prompts.chat import (ChatPromptTemplate, HumanMessagePromptTemplate, SystemMessagePromptTemplate)


from adm.constants import CTS
from adm.storage import MediaStorage
from adm.services import ParameterService

logger = logging.getLogger('dsp')

class Chatbot():
    def __init__(self) -> None:
        objParameter = ParameterService()
        openai.api_key = objParameter.getParameterByKey("OPENAI_API_KEY").value
        self.storage_path = CTS.CHATBOT_STORAGE_PATH
        media_storage = MediaStorage(self.storage_path)
        self.media_storage = media_storage

        # load index
        self.ai_storage_context = None
        self.ai_index = None
        self.base_dir = os.path.join(settings.BASE_DIR, 'storage')
        print(self.base_dir)
        try:
            self.ai_storage_context = StorageContext.from_defaults(persist_dir=self.base_dir)
            self.ai_index = load_index_from_storage(self.ai_storage_context)
        except BaseException as error:
            # exc_type, exc_obj, exc_tb = sys.exc_info()
            # fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1]
            # strErro = "ERRO: %s | %s | %s | %s" % (error, exc_type, fname, exc_tb.tb_lineno)
            logger.info('Chatbot() init: '+str(error))
            pass

    def train (self, chat=None):
        status = 400
        data = None

        logger.info('Chatbot Training started | Chat ? '+str(chat))

        docs = []
        if chat:
            # use chat to build up on existing store
            # save chat as doc in storage
            if not 'id' in chat:
                question = chat['chat'][0].split(':')[-1]
                chat['id'] = slugify(question)
            if not 'type' in chat:
                chat['type'] = 'instructions'
            if not 'category' in chat:
                chat['category'] = 'uncategorized'

            blob_name = chat['id']
            blob_ext = '.json'
            file_name = blob_name + blob_ext
            blob_dir = 'instructions/'
            blob_path = blob_dir + file_name
            text = json.dumps(chat, ensure_ascii=False)
            
            path_bucket = os.path.join(self.media_storage.location, self.storage_path, '/', blob_path)
            self.media_storage.save(path_bucket, BytesIO(text.encode("utf-8")))

            docs.append({
                'id_': blob_name,
                'text': text,
                'metadata': {
                    'filename': file_name, 
                    'category': chat['category']
                }
            })

        else:
            blobs = self.media_storage.list(self.storage_path + '/')
            for blob in blobs:
                text = blob.download_as_bytes().decode('utf-8')
                if len(text)>0:
                    obj = json.loads(text) if text.startswith('{') else text

                    if not 'category' in obj:  
                        obj['category'] = 'uncategorized'
                    
                    file_name = blob.name.split('/')[-1]

                    docs.append({
                        'id_': obj['id'],
                        'text': text,
                        'metadata': {
                            'filename': file_name, 
                            'category': obj['category']
                        }
                    })

        # parse documents
        ai_documents = []
        for doc in docs:
            ai_documents.append(Document(**doc))

        # create/add to index
        if self.ai_index is None:
            self.ai_index = VectorStoreIndex.from_documents(ai_documents, show_progress=True)
        else:
            for ai_doc in ai_documents:
                self.ai_index.insert(ai_doc)
        
        # save index to disk
        if ai_documents:
            self.ai_index.storage_context.persist(persist_dir=self.base_dir)
            status = 200
        
        logger.info('Chatbot Training finished | Status '+str(status))
        
        return status, data

    def chat (self, question=None):
        status = 400

        pre_prompt = [
            ChatMessage(
                role=MessageRole.SYSTEM,
                content=(
                "Act as an IT Support Level 1 for the company Mobees. Your name is Mel. You must answer the company's Drivers questions in Brazilian Portuguese, as they only speak this language. Mobees Drivers access support through a chatbot interface in the company's app, which is the only channel for support. If you're not sure how to answer, respond with only the following `<assignee:support>` (this will let our system know it should forward the issue to our support team and take it from there). "
                ),
            ),
            ChatMessage(
                role=MessageRole.USER,
                content=(
                    "Context information is below.\n"
                    "---------------------\n"
                    "{context_str}\n"
                    "---------------------\n"
                    "Given the context information and prior knowledge, "
                    "answer the question: {query_str}\n"
                )
            ),
        ]
        text_qa_template = ChatPromptTemplate(pre_prompt)
        # text_qa_template = Prompt.from_langchain_prompt(prompt_template)

        answer = self.ai_index.as_query_engine(
            text_qa_template=text_qa_template
        ).query(question)
        print(answer)

        status = 200
        data = { 'answer': answer.response }
        return status, data

        # chat_engine = self.ai_index.as_chat_engine(
        #     verbose=True
        # )
        # chat_engine.chat_repl()

instance = Chatbot()