src/agents/highlight_explain_agent/func.py

from typing import TypedDict, AnyStr

from .prompt import highlight_explain_chain, highlight_explain_question_generate_chain


async def highlight_explain(
    domain,
    question,
    highlight_terms,
    before_highlight_paragraph,
    after_highlight_paragraph,
    language,
):
    adjacent_paragraphs = (
        before_highlight_paragraph
        + "**"
        + highlight_terms
        + "**"
        + after_highlight_paragraph
    )
    response = await highlight_explain_chain.ainvoke(
        {
            "domain": domain,
            "highlight_terms": highlight_terms,
            "adjacent_paragraphs": adjacent_paragraphs,
            "question": question,
            "language": language,
        }
    )
    return response.explanation


async def highlight_explain_question_generate(
    domain,
    question,
    highlight_terms,
    before_highlight_paragraph,
    after_highlight_paragraph,
    language,
):
    adjacent_paragraphs = (
        before_highlight_paragraph
        + "**"
        + highlight_terms
        + "**"
        + after_highlight_paragraph
    )
    response = await highlight_explain_question_generate_chain.ainvoke(
        {
            "domain": domain,
            "highlight_terms": highlight_terms,
            "adjacent_paragraphs": adjacent_paragraphs,
            "question": question,
            "language": language,
        }
    )
    return response.questions