app/core/chatbot/chains.py

"""
LangChain chain definitions.
"""

from app.core.chatbot.llm import classify_query_llm, construct_query_llm, response_llm
from app.core.chatbot.prompt_templates import (
    classification_prompt,
    search_query_prompt,
    general_prompt,
    news_verification_prompt,
)

# Chain for classifying if a query needs web search
classification_chain = classification_prompt | construct_query_llm

# Chain for extracting optimized search queries
search_query_chain = search_query_prompt | construct_query_llm

# Chain for general responses (no search)
general_chain = general_prompt | response_llm

# Chain for news verification responses (with search results)
news_chain = news_verification_prompt | response_llm