app/gigachat_assistant.py - main - cppv/hackathon13

app/gigachat_assistant.py
93 строки4 KB
cppv
hackathon13
25 май 2025, 11:54
25 май 2025, 11:548af8c2d
from langchain_gigachat import GigaChat
from langchain.schema import HumanMessage, SystemMessage, AIMessage
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
import os

class GigaChatAssistant:
    # класс для работы с гигачатом
    
    def __init__(self, credentials, vectorstore):
        # тут инициализируем все что нужно
        self.credentials = credentials
        self.vectorstore = vectorstore
        self.chat_model = self._initialize_chat_model()
        self.qa_chain = self._initialize_qa_chain()
        
    def _initialize_chat_model(self):
        # создаем модель чата
        chat_model = GigaChat(
            credentials=self.credentials["auth_key"],  # ключ авторизации
            verify_ssl_certs=False
        )
        return chat_model
    
    def _initialize_qa_chain(self):
        # делаем цепочку для вопросов-ответов
        
        # промпт для модели
        template = """
        Ты - ассистент, который отвечает на вопросы по библиотеке LangChain.
        Используй только информацию из предоставленного контекста для ответа на вопрос.
        Если ты не знаешь ответа, просто скажи, что не можешь найти информацию по этому вопросу.
        Не придумывай информацию, которой нет в контексте.
        
        Контекст:
        {context}
        
        Вопрос: {question}
        
        Ответ:
        """
        
        prompt = PromptTemplate(
            template=template,
            input_variables=["context", "question"]
        )
        
        # собираем все в кучу
        qa_chain = RetrievalQA.from_chain_type(
            llm=self.chat_model,
            chain_type="stuff",
            retriever=self.vectorstore.as_retriever(search_kwargs={"k": 5}),
            chain_type_kwargs={"prompt": prompt}
        )
        
        return qa_chain
    
    def answer_question(self, question):
        # отвечаем на вопрос через поиск в документации
        try:
            response = self.qa_chain.run(question)
            return response
        except Exception as e:
            return f"Произошла ошибка при обработке вопроса: {str(e)}"
    
    def direct_chat(self, question, history=None):
        # просто болтаем с моделью напрямую
        try:
            messages = []
            
            # первым делом говорим модели кто она
            system_message = SystemMessage(content="""
            Ты - ассистент, который отвечает на вопросы по библиотеке LangChain.
            Отвечай кратко и по существу.
            """)
            messages.append(system_message)
            
            # добавляем прошлые сообщения если есть
            if history:
                for msg in history:
                    if msg['role'] == 'user':
                        messages.append(HumanMessage(content=msg['content']))
                    elif msg['role'] == 'assistant':
                        messages.append(AIMessage(content=msg['content']))
            
            # добавляем новый вопрос
            messages.append(HumanMessage(content=question))
            
            # получаем ответ
            response = self.chat_model.invoke(messages)
            return response.content
        except Exception as e:
            return f"Произошла ошибка при обработке вопроса: {str(e)}"