util.py

from google import genai
from google.genai import types
import time
import json
import os
from datetime import datetime
import re
import json
import transformers as hf
import peft
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForCausalLM



ssi_prompt = '''

<|begin_of_text|><|start_header_id|>system<|end_header_id|>

You are an intelligent and knowledgeable assistant.<|eot_id|><|start_header_id|>user<|end_header_id|>

# Key Information Types

{information_types}

# Dialogue
{dialog}

Identify Key Information Values from the Dialogue using the Key Information Types. If there is Key Information that does not fit into any existing Key Information Types, create an appropriate new Information Type for the Value with a description.<|eot_id|><|start_header_id|>assistant<|end_header_id|>

# Key Information Values

'''.strip()+'\n\n'

schema_dict = {}

with open("data/multiwoz_schema.json", "r", encoding="utf-8") as f:
    schema_data = json.load(f)
    for service in schema_data:
        for slot in service["slots"]:
            schema_dict[slot['name']] = slot['description']
            if slot['is_categorical']:
                schema_dict[slot['name']] = schema_dict[slot['name']] + ". Possible values include: " + ", ".join(slot['possible_values'])

def get_description(slot):
    if slot in schema_dict:
        return schema_dict[slot]
    else:
        return slot


folder_of_this_file = os.path.dirname(os.path.abspath(__file__))
api_key = open(os.path.join(folder_of_this_file, ".config/api_key.txt"), "r").read().strip()
key_index = 0
client = genai.Client(api_key = api_key)

prompt_count = {}
input_word_count = {}
output_word_count = {}

# gemma-3-27b-it
# gemini-2.5-flash
# gpt-5-nano-2025-08-07

tokenizers = {}
models = {}

def prompt_LLM(prompt, model, count_usage=True, retries = 10, use = "unknown"):
    global input_word_count, output_word_count, prompt_count, client
    attempt = 0
    while True:
        try:            
            # print(model, " is generating response...")
            if model.startswith("gpt"):
                return prompt_OpenAI(prompt, model=model)
            if model.startswith("Qwen"):
                return prompt_Qwen(prompt, model=model)
            print(model, " is generating response...")
            response = client.models.generate_content(
                model=model.split("google/")[1],
                contents=prompt,
                config=types.GenerateContentConfig(
                    thinking_config=types.ThinkingConfig(thinking_level="minimal")
                ),
            )
            # print("Response: ", response)
            if count_usage:
                if model not in prompt_count:
                    prompt_count[model] = {}
                    input_word_count[model] = {}
                    output_word_count[model] = {}
                if use not in prompt_count[model]:
                    prompt_count[model][use] = 0
                    input_word_count[model][use] = 0
                    output_word_count[model][use] = 0
                input_word_count[model][use] += len(prompt.split())
                output_word_count[model][use] += len(response.text.strip().split())
                prompt_count[model][use] += 1
            if not response:
                print("empty response")
                raise Exception("empty response")
            if not response.text:
                print("empty response text")
                raise Exception("empty response text")
            return response.text.strip()
        except Exception as e:
            raise e
            # print(f"Error occurred. Retrying... (Attempt {attempt + 1}/{retries})")
            # print("Error details:", str(e))
            # time.sleep(60)
            
            # else:
            #     print(prompt)
            #     print("key_index:", key_index)
            #     print("date:", date)
            #     print(response)
            #     print(response.text)
            #     raise Exception(f"Failed to get response from LLM after {retries} retries: {str(e)}")

def prompt_OpenAI(prompt, model="gpt-5-nano-2025-08-07", MaxToken=50, outputs=1):
    client = OpenAI()

    # result = client.responses.create(
    #     model="gpt-5",
    #     input="Write a haiku about code.",
    #     reasoning={ "effort": "low" },
    #     text={ "verbosity": "low" },
    # )

    # print(result.output_text)

def prompt_Qwen(prompt, model="Qwen/Qwen3-30B-A3B-Instruct-2507", MaxToken=100, outputs=1):

    if model not in tokenizers:
        tokenizers[model] = AutoTokenizer.from_pretrained(model)
    tokenizer = tokenizers[model]

    if model not in models:
        models[model] = AutoModelForCausalLM.from_pretrained(model, device_map="auto", torch_dtype="auto")
    model = models[model]

    messages = [
        {"role": "user", "content": prompt},
    ]
    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device)

    outputs = model.generate(**inputs, max_new_tokens=MaxToken)
    return tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:-1])

def show_token_usage():
    global input_word_count, output_word_count, prompt_count
    print("Token usage:")
    for key in prompt_count.keys():
        print("Model:", key)
        for use_key in prompt_count[key].keys():
            print("\tUse case:", use_key)
            print("\tInput word count:", input_word_count[key][use_key])
            print("\tOutput word count:", output_word_count[key][use_key])
            print("\tTotal prompts:", prompt_count[key][use_key])

nodes = {}

def generate_output_no_instruction(context, information, model = "gemma-3-27b-it"):

    context = "\n".join(context)
    agent_information = "\n".join(["The desired " + slot + " is " + ", ".join(information[slot]) for slot in information])
    
    prompt = f"""Agent information:
Agent information:
{agent_information}

Dialogue context:
{context}
    

Your task is to generate a natural, concise reply (1-3 sentences) that logically follows the conversation context.

Guidelines:
- Use the workflow as a reference to decide what information to request or provide next.
- Focus only on what is relevant to the current user message.
- Keep the response clear, helpful, and conversational.
- If you are supposed to perform some actions such as searching for some information, booking, etc., pretend you have performed the actions and make up the results. 
- If you need some information to answer the user, makes up a reasonable value for the information based on the dialogue context and use it in the response. 

Output the response that the agent should say to the user.
"""
    response = prompt_LLM(prompt, model = model)
    # print("Agent response:", response)
    return response

# def generate_output(instruction, context, information, model = "gemma-3-27b-it"):

#     context = "\n".join(context)
#     agent_information = "\n".join(["The desired " + slot + " is " + ", ".join(information[slot]) for slot in information])
#     prompt = f"""
# Your task is to find nodes corresponding to the actions that you should do to logically follow the workflow and the conversation context.

# Guidelines:
# - Use the workflow as a reference to decide what information to request or provide next.
# - Focus only on what is relevant to the current user message.
    
# Workflow:
# {instruction}

# Agent information:
# {agent_information}

# Dialogue context:
# {context}

# Output at most three node numbers, separated by commas. Only output the numbers. 

# Output format:
# [goal number]-[node number], [goal number]-[node number], ...
# Example output: 
# 1-3,2-5,2-7
# Example output (if no node is needed):
# None
#     """
#     chosen_nodes = prompt_LLM(prompt, model = model).split(",")
#     if len(chosen_nodes) == 1 and chosen_nodes[0].strip().lower() in ["none", "null", "no node"]:
#         chosen_nodes = []
#     else:
#         chosen_nodes = [node.strip() for node in chosen_nodes]
#     for node in chosen_nodes:
#         nodes[node] = nodes.get(node, 0) + 1
#     print("Nodes chosen:", chosen_nodes)
#     # Process each chosen node
#     print(chosen_nodes)
#     chosen_nodes = [node.split("N")[0]+node.split("N")[1] if "N" in node else node for node in chosen_nodes]
#     chosen_instructions = [instruction.split("Goal "+node.split("-")[0])[1].split("- N"+node.split("-")[1])[1].split("\n")[0].strip() for node in chosen_nodes]
#     print(chosen_instructions)
#     instruction_text = "\n".join(chosen_instructions)
    
#     prompt = f"""
# Workflow:
# Use node {chosen_nodes} to decide the next action.
# {instruction_text}

# Agent information:
# {agent_information}

# Dialogue context:
# {context}
    

# Your task is to generate a natural, concise reply (1-3 sentences) that logically follows the workflow and the conversation context.

# Guidelines:
# - Use the workflow as a reference to decide what information to request or provide next.
# - Focus only on what is relevant to the current user message.
# - Keep the response clear, helpful, and conversational.
# - If you are supposed to perform some actions such as searching for some information, booking, etc., pretend you have performed the actions and make up the results. 
# - If you need some information to answer the user, makes up a reasonable value for the information based on the dialogue context and use it in the response. 

# Output the response that the agent should say to the user.

#     """
#     response = prompt_LLM(prompt, model = model)
#     # print("Agent response:", response)
#     # response = response.replace("[value_city]", "[value_name]").replace("[attraction id]", "[attraction_id]").replace("[city name]","[value_name]").replace("[value_parking]", "[parking]").replace("[value_internet]", "[internet]").replace("[city_name]", "[value_name]")
#     return response
def generate_output(instruction, context, information, model = "gemma-3-27b-it"):

    context = "\n".join(context)
    agent_information = "\n".join(["The desired " + slot + " is " + ", ".join(information[slot]) for slot in information])
    prompt = f"""
Workflow:
{instruction}

Agent information:
{agent_information}

Dialogue context:
{context}
    

Your task is to generate a natural, concise reply (1-3 sentences) that logically follows the workflow and the conversation context.

Guidelines:
- Use the workflow as a reference to decide what information to request or provide next.
- Focus only on what is relevant to the current user message.
- Keep the response clear, helpful, and conversational.
- If you are supposed to perform some actions such as searching for some information, booking, etc., pretend you have performed the actions and make up the results. 
- If you need some information to answer the user, makes up a reasonable value for the information based on the dialogue context and use it in the response. 

Output the response that the agent should say to the user.

    """
    response = prompt_LLM(prompt, model = model)
    # print("Agent response:", response)
    # response = response.replace("[value_city]", "[value_name]").replace("[attraction id]", "[attraction_id]").replace("[city name]","[value_name]").replace("[value_parking]", "[parking]").replace("[value_internet]", "[internet]").replace("[city_name]", "[value_name]")
    return response

def check_nodes():
    print("Node usage statistics:")
    print(nodes)




adapter_id = 'jdfinch/ssi_dots_lora'
from transformers.utils import logging
logging.set_verbosity_error()

# Load PEFT config to get the base model
peft_config = peft.PeftConfig.from_pretrained(adapter_id)
base_model = peft_config.base_model_name_or_path

# Load base model
model = hf.AutoModelForCausalLM.from_pretrained(base_model, device_map="auto", torch_dtype="auto")
tokenizer = hf.AutoTokenizer.from_pretrained(base_model)

# Load adapter on top
model = peft.PeftModel.from_pretrained(model, adapter_id)

print("Model and adapter loaded.")

def extract_goals(text):
    text = text.split("# key information values")[-1].strip()
    # print("Extracting goals from text:")
    # print(text)
    goals = {}
    current_domain = None
    buffer = []

    # Regex patterns
    domain_pattern = re.compile(r'^##\s*(\w+)')
    slot_pattern = re.compile(r"\*\s*([^:]+):\s*(.+)")

    for line in text.splitlines():
        line = line.strip()

        # Detect domain (goal)
        domain_match = domain_pattern.match(line)
        if domain_match:
            current_domain = domain_match.group(1)
            continue

        # Detect slot-value pairs
        slot_match = slot_pattern.match(line)
        if slot_match:
            if not current_domain:
                slot, value = slot_match.groups()
                buffer.append((slot.strip(), value.strip()))
                if "-" in slot:
                    current_domain = slot.split("-")[0]
                    for buffered_slot, buffered_value in buffer:
                        # if buffered_slot in goals:
                            # print(goals)
                            # print(text)
                            # raise ValueError(f"Duplicate slot '{buffered_slot}' found.")
                        goals[buffered_slot] = buffered_value.strip()
                    buffer = []
            else:
                slot, value = slot_match.groups()
                if "-" not in slot:
                    slot = current_domain + "-" + slot.strip()
                # if slot in goals:
                #     print(goals)
                #     print(text)
                #     raise ValueError(f"Duplicate slot '{slot}' found.")
                # print("add slot", slot, "with value", value, "to domain", current_domain)
                goals[slot] = value.strip()

    return goals

def schema_induction(dialog, user_information):
    # print("-"*40)
    # print("doing schema induction...")

    information_types = ""
    current_domain = None
    for slot, _ in user_information["slots"].items():
        goal = slot.split("-")[0]
        if goal != current_domain:
            current_domain = goal
            information_types += f"## {current_domain}\n"
        information_types += f"* {slot}: {get_description(slot)}\n"

    for request in user_information["requests"]:
        goal = request.split("-")[0]
        if goal != current_domain:
            current_domain = goal
            information_types += f"## {current_domain}\n"
        information_types += f"* {request}: {get_description(request)}\n"

    # Revise prompt
    prompt = ssi_prompt.format(
        information_types=information_types,
        dialog=dialog
    )

    # print("Prompt for schema induction:")
    # print(prompt)

    # Tokenize prompt
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    # Generate
    output = model.generate(
        **inputs,
        max_new_tokens=1000,
        do_sample=False,
        top_p=None,
        pad_token_id=tokenizer.eos_token_id
    )

    # Decode and print only new generation
    generated = tokenizer.decode(output[0], skip_special_tokens=False)
    # print("Generated schema:")
    # print(generated)

    generated = generated.lower()
    
    goals = extract_goals(generated)
    return goals