ChatGPT Антиспам. Часть 1

import os
import environ
import openai
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.text_splitter import MarkdownHeaderTextSplitter

env = environ.Env()
environ.Env.read_env()


class AntiSpam:
    def __init__(self):
        self.dev = False
        self.db = None
        openai.api_key = env('OPENAI_API_KEY')
        self.prj_dir = os.getcwd()  # Путь до каталога с проектом
        self.spam_base = 'spam_base.txt'  # Текстовая база спам сообщений
        self.embedding = OpenAIEmbeddings()  # Инициализирум модель эмбеддингов
        self.load_embeddings()
        self.role_system = 'Ты антиспам'  # Роль ChatGPT
        self.temperature = 0
        self.model = 'gpt-3.5-turbo'

    @staticmethod
    def load_txt_file(file_name):
        f = open(file_name, 'r', encoding="utf8")
        file_text = f.read()
        return file_text

    def generate_embeddings(self, chunks):
        # Создадим индексную базу
        self.db = FAISS.from_documents(chunks, self.embedding)
        # Сохраняем базу локально
        self.db.save_local(folder_path=self.prj_dir, index_name="faiss_index")

    def load_embeddings(self):
        try:
            self.db = FAISS.load_local(folder_path=self.prj_dir, index_name="faiss_index", embeddings=self.embeddings)
        except Exception as e:
            self.process_spam_db()
            if self.dev:
                print('Ошибка загрузки базы FAISS', e.__str__())

    def process_spam_db(self):
        spam_text = self.load_txt_file(self.spam_base)

        headers_to_split_on = [
            ("#", "Header 1"),
        ]
        markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
        chunks = markdown_splitter.split_text(text=spam_text)

        if self.dev:
            print("Общее количество чанков: ", len(chunks))

        self.generate_embeddings(chunks)

    def find_examples(self, message):
        # Поиск релевантных отрезков из базы знаний
        docs = self.db.similarity_search(message, k=4)

        db_messages = ""
        for i, doc in enumerate(docs):
            db_messages = db_messages + f"Пример спам сообщения №{i + 1}\n{doc.page_content}\n"

        if self.dev:
            print('db_messages:\n', db_messages)

        return db_messages

    def check_message(self, message):
        try:
            examples = self.find_examples(message)

            role_user = f"Ответь на вопрос является ли данное сообщение спамом на основе представленных примеров спам " \
                        f"сообщений. Ответь коротко, да или нет, не придумывай ничего от себя. " \
                        f"Примеры спам сообщений: {examples}\n\n" \
                        f"Текст сообщения которое надо проверить на спам: \n{message}"

            gpt_messages = [
                {"role": "system", "content": self.role_system},
                {"role": "user", "content": role_user}
            ]
            completion = openai.ChatCompletion.create(
                model=self.model,
                messages=gpt_messages,
                temperature=self.temperature
            )

            if self.dev:
                print(completion)
                print("-----------------")

            return completion.choices[0].message.content
        except Exception as e:
            print('Ошибка обработки запроса ChatGPT. ', e.__str__())

        return None


if __name__ == '__main__':
    anti_spam = AntiSpam()
    anti_spam.dev = True

    chat_message = anti_spam.load_txt_file('message.txt')
    spam_message = anti_spam.load_txt_file('spam.txt')

    res = anti_spam.check_message(spam_message)
    print("Тест spam_message: ", res)

    res = anti_spam.check_message(chat_message)
    print("Тест chat_message: ", res)