utils.py

from pathlib import Path
from uuid import uuid4
from enum import Enum


class ChromaConstants(str, Enum):
    PUBLICATIONS_CHROMA_PATH = "chroma_data"
    STRUCTURED_COLLECTION = "structured"
    CITATION_COLLECTION = "citation"


def get_question_uuid():
    """Create a UUID and a unique directory for each question."""
    save_dir = Path("data") / "questions"
    save_dir.mkdir(exist_ok=True)

    there_is_collision = True
    while there_is_collision:
        question_id = uuid4()
        question_dir = save_dir / str(question_id)
        if not question_dir.is_dir():
            question_dir.mkdir()
            there_is_collision = False

    return question_id, question_dir


def save_generated_data(
    strategy_name,
    question,
    question_id,
    question_dir,
    refined_question,
    context,
    refined_context,
    answer,
    refined_answer,
    settings,
):
    """Save the questions and contexts retrieved."""

    # Process the list[Document] in context and refined_context to include
    # the page_content and the scores for each entry in the list
    context = [doc.page_content + "\n" + str(doc.metadata["score"]) for doc in context]
    refined_context = [
        doc.page_content + "\n" + str(doc.metadata["score"]) for doc in refined_context
    ]
    context = "\n\n".join(context)
    refined_context = "\n\n".join(refined_context)

    NUM_DIVISION_CHARS = 50
    log_file = question_dir / f'{question_id}_{strategy_name}_{settings["k"]}.log'
    with open(log_file, "w") as f:
        f.write(f"STRATEGY: {strategy_name}\n")
        f.write(f"QUESTION ID: {question_id}\n")
        f.write("SETTINGS:\n")
        for key, value in settings.items():
            f.write(f"\t{key} = {value}\n")
        f.write(f"ORIGINAL QUESTION: {question}\n")
        f.write(f"RETRIEVER QUESTION: {refined_question}\n")
        f.write("\n" + "-" * NUM_DIVISION_CHARS + "\n\n")
        f.write("ANSWER:\n")
        f.write(answer)
        f.write("\n\n" + "-" * NUM_DIVISION_CHARS + "\n\n")
        f.write("REFINED ANSWER:\n")
        f.write(refined_answer)
        f.write("\n\n" + "-" * NUM_DIVISION_CHARS + "\n\n")
        f.write("CONTEXT:\n")
        f.write(context)
        f.write("\n\n" + "-" * NUM_DIVISION_CHARS + "\n\n")
        f.write("REFINED CONTEXT:\n")
        f.write(refined_context)