dam_helper/llm_manager/llm_parser.py

from form.form import work_categories
from local_storage.entities import SavedDetails
from enums.enums import Questions


class LlmParser:

    @classmethod
    def parse_verification_prompt_answers(cls, llm_answer:str, details: list[SavedDetails] | None)\
            -> dict[Questions, str | None]:
        print(f"llm answer: {llm_answer}")
        answers = {}
        for d in details:
            answers.update(d.to_answers())

        unanswered_questions = sorted([k.value for k in Questions.values() if k not in answers])
        unanswered_questions = iter(unanswered_questions)
        i = 0
        question_id = next(unanswered_questions)
        lines = [l for l in llm_answer.split("\n") if len(l.strip()) > 0 and not l.strip().endswith("?")]
        while i < len(lines):
            if question_id is None:
                break
            line = lines[i].strip()
            if len(line) == 0:
                i += 1
            elif "null" in lines[i]:
                answers[Questions(question_id)] = None
                i += 1
                question_id = next(unanswered_questions, None)
            elif ":" in lines[i]:
                answers[Questions(question_id)] = line.split(":")[1]
                i += 1
                question_id = next(unanswered_questions, None)
            else:
                answers[Questions(question_id)] = line
                i += 1
                question_id = next(unanswered_questions, None)
        return answers

    @classmethod
    def parse_get_categories_answer(cls, category_answer) -> list[str]:
        categories = []
        for category in category_answer.split(";"):
            categories.extend([k for k, v in work_categories.items() if category.lower() in v.lower()])
        return categories