src/chains.py - main - GreyCat/knowledge-assistant-bot

src/chains.py
46 строк2 KB
Vlad
Create chains.py
08 сен 2025, 21:39
08 сен 2025, 21:39a99a4c3
import os
from dotenv import load_dotenv
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI

load_dotenv()
api_key = os.getenv("OPENROUTER_API_KEY")

llm = ChatOpenAI(
    model="moonshotai/kimi-k2:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=api_key,
    temperature=0.3
)

# === Промпты ===
stuff_template = """Контекст:\n{context}\n\nВопрос: {question}\n\nОтвет:"""
refine_template = """Ответ:\n{existing_answer}\n\nДоп. контекст:\n{context}\n\nУточни ответ на вопрос: {question}\n\nОбновлённый ответ:"""

stuff_prompt = PromptTemplate.from_template(stuff_template)
refine_prompt = PromptTemplate.from_template(refine_template)

# === Цепочки ===
def create_stuff_chain():
    def chain(question, docs):
        context = "\n\n".join([doc.page_content for doc in docs])
        response = (stuff_prompt | llm).invoke({"context": context, "question": question})
        return response.content if hasattr(response, 'content') else str(response)
    return chain

def create_refine_chain():
    def chain(question, docs):
        if not docs:
            return "Нет информации для ответа."
        # начальный ответ
        initial = (stuff_prompt | llm).invoke({"context": docs[0].page_content, "question": question})
        ans = initial.content if hasattr(initial, 'content') else str(initial)
        # уточняем остальными
        for doc in docs[1:]:
            refined = (refine_prompt | llm).invoke({"existing_answer": ans, "context": doc.page_content, "question": question})
            ans = refined.content if hasattr(refined, 'content') else str(refined)
        return ans
    return chain

stuff_chain = create_stuff_chain()
refine_chain = create_refine_chain()