utils.py

import json
from collections import defaultdict
import openai
import re
from config import CFG_APP
from text_embedder import SentenceTransformersTextEmbedder
from datetime import datetime
import tiktoken

doc_metadata = json.load(open(CFG_APP.DOC_METADATA_PATH, "r"))
# Embedding Model
if "sentence-transformers" in CFG_APP.EMBEDDING_MODEL:
    text_embedder = SentenceTransformersTextEmbedder(
        model_name=CFG_APP.EMBEDDING_MODEL,
        paragraphs_path=CFG_APP.DATA_FOLDER,
        device=CFG_APP.DEVICE,
        load_existing_index=True,
    )
else:
    raise ValueError("Embedding model not found !")


# Util Functions
def retrieve_doc_metadata(doc_metadata, doc_id):
    for meta in doc_metadata:
        if meta["id"] == doc_id:
            return meta


def get_reformulation_prompt(query: str) -> list:
    return [
        {
            "role": "user",
            "content": f"""{CFG_APP.REFORMULATION_PROMPT}
            ---
            query: {query}
            standalone question: """,
        }
    ]

def get_hyde_prompt(query: str) -> list:
    return [
        {
            "role": "user",
            "content": f"""{CFG_APP.HYDE_PROMPT}
            ---
            query: {query}
            output: """,
        }
    ]


def make_pairs(lst):
    """From a list of even lenght, make tupple pairs
    Args:
        lst (list): a list of even lenght
    Returns:
        list: the list as tupple pairs
    """
    assert not (l := len(lst) % 2), f"your list is of lenght {l} which is not even"
    return [(lst[i], lst[i + 1]) for i in range(0, len(lst), 2)]


def make_html_source(paragraph, meta_doc, i):
    content = paragraph["content"]
    meta_paragraph = paragraph["meta"]
    return f"""
<div class="card" id="document-{i}">
    <div class="card-content">
        <h2>Doc {i} - {meta_doc['short_name']} - Page {meta_paragraph['page_number']}</h2>
        <p>{content}</p>
    </div>
    <div class="card-footer">
        <span>{meta_doc['short_name']}</span>
        <a href="{meta_doc['url']}#page={meta_paragraph['page_number']}" target="_blank" class="pdf-link">
            <span role="img" aria-label="Open PDF">🔗</span>
        </a>
    </div>
</div>
"""


def preprocess_message(text: str, docs_url: dict) -> str:
    return re.sub(
        r"\[doc (\d+)\]",
        lambda match: f'<a href="{docs_url[match.group(1)]}" target="_blank" class="pdf-link">{match.group(0)}</a>',
        text,
    )


def parse_glossary(query):
    file = "glossary.json"
    glossary = json.load(open(file, "r"))
    words_query = query.split(" ")
    for i, word in enumerate(words_query):
        for key in glossary.keys():
            if word.lower() == key.lower():
                words_query[i] = words_query[i] + f" ({glossary[key]})"
    return " ".join(words_query)


def num_tokens_from_string(string: str, encoding_name: str) -> int:
    encoding = tiktoken.encoding_for_model(encoding_name)
    num_tokens = len(encoding.encode(string))
    return num_tokens


def chat(
    query: str,
    history: list,
    query_mode : str,
    threshold: float = CFG_APP.THRESHOLD,
    k_total: int = CFG_APP.K_TOTAL,
) -> tuple:
    """retrieve relevant documents in the document store then query gpt-turbo
    Args:
        query (str): user message.
        history (list, optional): history of the conversation. Defaults to [system_template].
        report_type (str, optional): should be "All available" or "IPCC only". Defaults to "All available".
        threshold (float, optional): similarity threshold, don't increase more than 0.568. Defaults to 0.56.
    Yields:
        tuple: chat gradio format, chat openai format, sources used.
    """
    if query_mode == 'Reformulation':

        reformulated_query = openai.ChatCompletion.create(
            model=CFG_APP.MODEL_NAME,
            messages=get_reformulation_prompt(parse_glossary(query)),
            temperature=0,
            max_tokens=CFG_APP.MAX_TOKENS_REF_QUESTION,
        )

    else :

        reformulated_query = openai.ChatCompletion.create(
            model=CFG_APP.MODEL_NAME,
            messages=get_hyde_prompt(parse_glossary(query)),
            temperature=0,
            max_tokens=CFG_APP.MAX_TOKENS_REF_QUESTION,
        )

    reformulated_query = reformulated_query["choices"][0]["message"]["content"]
    if len(reformulated_query.split("\n")) == 2:
        reformulated_query, language = reformulated_query.split("\n")
        language = language.split(":")[1].strip()
    else:
        reformulated_query = reformulated_query.split("\n")[0]
        language = "English"

    sources, scores = text_embedder.retrieve_faiss(
        reformulated_query,
        k_total=k_total,
        threshold=threshold,
    )
    if CFG_APP.DEBUG == True:
        print("Scores : \n", scores)

    messages = history + [{"role": "user", "content": query}]

    if query_mode == 'HYDE' :
        reformulated_query = reformulated_query.split("?")[0] + '?'

    docs_url = defaultdict(str)

    if len(sources) > 0:
        docs_string = []
        docs_html = []

        num_tokens = num_tokens_from_string(CFG_APP.SOURCES_PROMPT, CFG_APP.MODEL_NAME)

        for i, data in enumerate(sources, 1):
            meta_doc = retrieve_doc_metadata(doc_metadata, data["meta"]["document_id"])
            doc_content = f"📃 Doc {i}: \n{data['content']}"
            num_tokens_doc = num_tokens_from_string(doc_content, CFG_APP.MODEL_NAME)
            if num_tokens + num_tokens_doc > CFG_APP.MAX_TOKENS_API:
                break
            num_tokens += num_tokens_doc
            docs_string.append(doc_content)
            docs_html.append(make_html_source(data, meta_doc, i))

            url_doc = f'<a href="{meta_doc["url"]}#page={data["meta"]["page_number"]}" target="_blank" class="pdf-link">'
            docs_url[i] = url_doc

        docs_string = "\n\n".join(
            [f"Query used for retrieval:\n{reformulated_query}"] + docs_string
        )
        docs_html = "\n\n".join(
            [f"Query used for retrieval:\n{reformulated_query}"] + docs_html
        )
        messages.append(
            {
                "role": "system",
                "content": f"{CFG_APP.SOURCES_PROMPT}\n\n{docs_string}\n\nAnswer in {language}:",
            }
        )

        if CFG_APP.DEBUG == True:
            print(f" 👨‍💻 question asked by the user : {query}")
            print(f" 🕛 time : {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

            print(" 🔌 messages sent to the API :")
            api_messages = [
                {"role": "system", "content": CFG_APP.INIT_PROMPT},
                {"role": "user", "content": reformulated_query},
                {
                    "role": "system",
                    "content": f"{CFG_APP.SOURCES_PROMPT}\n\n{docs_string}\n\nAnswer in {language}:",
                },
            ]
            for message in api_messages:
                print(
                    f"length : {len(message['content'])}, content : {message['content']}"
                )

        response = openai.ChatCompletion.create(
            model=CFG_APP.MODEL_NAME,
            messages=[
                {"role": "system", "content": CFG_APP.INIT_PROMPT},
                {"role": "user", "content": reformulated_query},
                {
                    "role": "system",
                    "content": f"{CFG_APP.SOURCES_PROMPT}\n\n{docs_string}\n\nAnswer in {language}:",
                },
            ],
            temperature=0,  # deterministic
            stream=True,
            max_tokens=CFG_APP.MAX_TOKENS_ANSWER,
        )
        complete_response = ""
        messages.pop()
        messages.append({"role": "assistant", "content": complete_response})
        for chunk in response:
            chunk_message = chunk["choices"][0]["delta"].get("content")
            if chunk_message:
                complete_response += chunk_message
                complete_response = preprocess_message(complete_response, docs_url)
                messages[-1]["content"] = complete_response
                gradio_format = make_pairs([a["content"] for a in messages[1:]])
                yield gradio_format, messages, docs_html

    else:
        docs_string = "⚠️ No relevant passages found in this report"
        complete_response = "**⚠️ No relevant passages found in this report, you may want to ask a more specific question.**"
        messages.append({"role": "assistant", "content": complete_response})
        gradio_format = make_pairs([a["content"] for a in messages[1:]])
        yield gradio_format, messages, docs_string