ask_candid/utils.py

from typing import Any

from langchain_core.documents import Document


def html_format_docs_chat(docs: list[Document]) -> str:
    """Formats Candid sources

    Parameters
    ----------
    docs : list[Document]
        Retrieved documents for context

    Returns
    -------
    str
        Formatted HTML
    """

    html = ""
    if docs:
        docs_html = []
        for doc in docs:
            s_name = doc.metadata.get("source", "Source")
            s_url = doc.metadata.get("url", "URL")

            s_html = (
                "<span class='source-item'>"
                f"<a href='{s_url}' target='_blank' rel='noreferrer' class='ssearch-source'>"
                f"{doc.metadata['title']} &vert; {s_name}</a></span>"
            )

            docs_html.append(s_html)

        html = f"<h2>Candid Resources</h2><div id='ssearch-sources'>{'<br>'.join(docs_html)}</div>"
    return html


def format_chat_ag_response(chatbot: list[Any]) -> list[Any]:
    """If we called retriever, we appended sources as as one more message. Here we concatinate HTML of sources
        with the AI response
    Returns:
        _type_: updated chatbot message as HTML
    """
    sources = ""
    if chatbot:
        title = (chatbot[-1].get("metadata") or {}).get("title", None)
        if title == "Sources HTML":
            sources = chatbot[-1]["content"]
            chatbot.pop(-1)
            chatbot[-1]["content"] = chatbot[-1]["content"] + sources
    return chatbot