models/openai/role_models.py

import logging
import pandas as pd
from models.custom_parsers import CustomStringOutputParser
from utils.app_utils import get_random_name
from langchain.chains import ConversationChain
from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate


ISSUE_MAPPING = {
    "anxiety": "issue_Anxiety",
    "suicide": "issue_Suicide",
    "safety_planning": "issue_Suicide",
    "GCT": "issue_Gral",
}

EN_TEXTER_TEMPLATE_ = """The following is a conversation between you and a crisis counselor.
{current_issue}
You are able to reply with what the character should say. You are able to reply with your character's dialogue inside and nothing else. Do not write explanations.
Do not disclose your name unless asked.
Current conversation:
{history}
helper: {input}
texter:"""

SP_TEXTER_TEMPLATE_ = """La siguiente es una conversacion contigo y un consejero de crisis
{current_issue}
Puedes responder como lo haria tu personaje. Puedes responder como si fueras tu personaje y nada mas. No escribas explicaciones
No reveles tu nombre a menos que te lo pregunten
Conversacion Actual:
{history}
helper: {input}
texter:"""

CURRENT_ISSUE_MAPPING = {
    "issue_Suicide-en": "Your character, {texter_name}, has suicidal thoughts. Your character has a plan to end his life and has all the means and requirements to do so. {seed}",
    "issue_Anxiety-en": "Your character, {texter_name}, is experiencing anxiety. Your character has suicide thoughts but no plan. {seed}",
    "issue_Suicide-es": "Tu personaje, {texter_name}, tiene pensamientos suicidas. Tu personaje tiene un plan para terminar con su vida y tiene todos los medios y requerimientos para hacerlo. {seed}",
    "issue_Anxiety-es": "Tu personaje, {texter_name}, experimenta ansiedad. Tu personaje tiene pensamientos suicidas pero ningun plan. {seed}",
    "issue_Gral-en": "Your character {texter_name} is experiencing a mental health crisis. {seed}",
    "issue_Gral-es": "Tu personaje {texter_name} esta experimentando una crisis de salud mental. {seed}",
}

def get_template_role_models(issue: str, language: str, texter_name: str = "", seed="") -> str:
    """_summary_

    Args:
        issue (str): Issue for template, current options are ['issue_Suicide','issue_Anxiety']
        language (str): Language for the template, current options are ['en','es']
        texter_name (str): texter to apply to template, defaults to None

    Returns:
        str: template
    """
    current_issue = CURRENT_ISSUE_MAPPING.get(
        f"{issue}-{language}", CURRENT_ISSUE_MAPPING[f"issue_Gral-{language}"]
    )
    default_name = get_random_name()
    current_issue = current_issue.format(
        texter_name=default_name if not texter_name else texter_name,
        seed = seed
    )

    if language == "en":
        template = EN_TEXTER_TEMPLATE_.format(current_issue=current_issue, history="{history}", input="{input}")
    elif language == "es":
        template = SP_TEXTER_TEMPLATE_.format(current_issue=current_issue, history="{history}", input="{input}")

    return template

def get_role_chain(template, memory, temperature=0.8):

    PROMPT = PromptTemplate(
        input_variables=['history', 'input'],
        template=template
    )
    llm = OpenAI(
        temperature=temperature, 
        max_tokens=150,
    )
    llm_chain = ConversationChain(
        llm=llm,
        prompt=PROMPT,
        memory=memory,
        output_parser=CustomStringOutputParser()
    )
    logging.debug(f"loaded GPT3.5 model")
    return llm_chain, "helper:"