ИИ МЕНЕДЖЕР НА БАЗЕ GIGACHAT

import os
import environ
from langchain.schema import HumanMessage, SystemMessage
from langchain.chat_models.gigachat import GigaChat
from langchain.text_splitter import MarkdownHeaderTextSplitter
from langchain.vectorstores import Chroma
from langchain_community.embeddings import GigaChatEmbeddings

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


class Chat:
    def __init__(self):
        self.chroma_db = None
        self.dev = False
        self.llm = GigaChat(credentials=env('AUTH'), verify_ssl_certs=False)
        self.embeddings = GigaChatEmbeddings(credentials=env('AUTH'), verify_ssl_certs=False)
        self.prj_dir = os.getcwd()  # Путь до каталога с проектом
        self.persist_directory = self.prj_dir + "\data" # Каталог хранения векторной базы


class BaseLLM(Chat):
    def __init__(self):
        super().__init__()
        self.all_chunks = []        self.product_base = 'base_ai.txt'  # Текстовая база товаров

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

    def load_product_base(self):
        load_txt = self.load_txt_file('base_ai.txt')

        headers_to_split_on = [
            ("#", "Header 1"),
        ]
        markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
        self.all_chunks = markdown_splitter.split_text(text=load_txt)

    def generate_embeddings(self):
        self.load_product_base()
        chunks = []        count = 0
        all_count = 0
        for chunk in self.all_chunks:
            all_count = all_count + 1
            count = count + 1
            chunks.append(chunk)
            if count == 10:
                self.save_embeddings(chunks)
                chunks = []                count = 0
                print(f"Всего обработано {all_count}")

        if len(chunks) > 0:
            self.save_embeddings(chunks)

    def save_embeddings(self, chunks):
        self.chroma_db = Chroma.from_documents(
            chunks,
            self.embeddings,
            persist_directory=self.persist_directory,
            collection_name="lc_chroma_store")
        self.chroma_db.persist()


class ProductManager(Chat):
    def __init__(self):
        super().__init__()
        self.load_base()
        self.messages = []        self.role_system = "Ты менеджер магазина автозапчастей. Отвечаешь на вопросы клиентов в чате. " \
                           "Посмотри на примеры товаров и сформируй ответ указав наименования товаров и их цены."

    def load_base(self):
        self.chroma_db = Chroma(persist_directory=self.persist_directory, embedding_function=self.embeddings,
                                collection_name="lc_chroma_store")

        collection = self.chroma_db.get()
        if self.dev:
            print(f"Всего в базе: {len(collection['ids'])}")

        if len(collection['ids']) == 0:
            admin = BaseLLM()
            admin.generate_embeddings()

    def find_examples(self, message):
        # Поиск релевантных отрезков из базы знаний
        docs = self.chroma_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 send_message(self, question):
        examples = self.find_examples(question)

        if len(self.messages) == 0:
            self.messages.append(SystemMessage(content=f"{self.role_system} Список товаров: ${examples}"))

        self.messages.append(HumanMessage(content=question))
        res = self.llm(self.messages)
        # Сохраняем контекст разговора
        self.messages.append(res)

        return res.content


manager = ProductManager()

while True:
    # Ввод пользователя
    user_input = input("Введите сообщение: ")
    if not user_input:
        break
    print("Менеджер: ", manager.send_message(user_input))