utils/format.py

import json
from typing import Union

from utils.interview_messages import first_conversation_message


def extract_json(input_string: str) -> Union[bool, dict]:
    """String to Json function"""

    # First, ensure we remove json wrapper
    input_string = input_string.replace("```json", "```").replace("```", "")

    # Now, ensure we have stripped everything so it is just json
    input_string_formatted = input_string.lstrip("{").rstrip("}")

    # Ensure we do not have the weird \_ behaviour that models sometimes include
    input_string_formatted = input_string_formatted.replace("\_", "_")

    try:
        return True, json.loads(input_string_formatted)

    except json.JSONDecodeError as e:
        print(f"Error parsing JSON Output: {input_string}. Error: {e}")
        return False, {}


def format_chat_history_cohere(chat_history: list, background_info: dict) -> list:
    """Takes streamlit chat history, and converts to cohere format"""

    # TODO: Could use cohere to track history, maybe for the future
    new_output = [
        {
            "role": "USER",
            "message": first_conversation_message.replace(
                "<cv>", background_info["cv"]
            ),
        },
        {
            "role": "CHATBOT",
            "message": f"Thanks. Can you send me the job posting so I can reference that while I think of some questions to ask you?",
        },
        {
            "role": "USER",
            "message": f"Here is the contents of the job posting that I'm interested in: {background_info['job_posting']}",
        },
    ]

    # Lazy approach to format it correctly for cohere input!
    for item in chat_history:
        new_output.append(
            {
                "role": "USER" if item["role"] == "user" else "CHATBOT",
                "message": item["message"],
            }
        )

    return new_output


if __name__ == "__main__":
    example_json = """
```json
{
"dogs": "are blue?"
}
"""
    extract_json(example_json)