src/generation.py

import re
import json
import time

from src.config import (
    TEMPERATURE,
    LONG_ANSWER_USER_TEMPLATE,
    SENTENCE_SPLIT_SYSTEM,
    SENTENCE_SPLIT_USER_TEMPLATE,
    FACTOID_SYSTEM,
    FACTOID_USER_TEMPLATE,
    Q_QUESTIONS_SYSTEM,
    Q_QUESTIONS_USER_TEMPLATE,
    M_ANSWERS_SYSTEM,
    M_ANSWERS_USER_TEMPLATE,
)


def _chat(client, model_name, system_prompt, user_prompt, max_tokens, logprobs=False):
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": user_prompt})
    return client.chat.completions.create(
        model=model_name,
        messages=messages,
        max_tokens=int(max_tokens),
        temperature=TEMPERATURE,
        logprobs=logprobs,
        top_logprobs=1 if logprobs else None,
    )


def query_model(system_prompt, user_prompt, model_name, max_tokens, client):
    """Returns {"answer": str}."""
    resp = _chat(client, model_name, system_prompt, user_prompt, max_tokens, logprobs=False)
    return {"answer": resp.choices[0].message.content or ""}


def query_model_with_logprobs(question, model_name, max_tokens, client):
    """
    Generates a 3-sentence answer with logprobs captured.
    Returns (answer_str, token_logprobs) where token_logprobs is
    a list of {"token": str, "logprob": float}.
    """
    user_prompt = LONG_ANSWER_USER_TEMPLATE.format(question=question)
    resp = _chat(client, model_name, "", user_prompt, max_tokens, logprobs=True)
    answer = (resp.choices[0].message.content or "").strip()
    token_logprobs = []
    lp = resp.choices[0].logprobs
    if lp and lp.content:
        for t in lp.content:
            token_logprobs.append({"token": t.token, "logprob": t.logprob})
    return answer, token_logprobs


def _clean_json(raw):
    return re.sub(r"```json|```", "", raw).strip()


def split_answers_into_json_strings(answer, model_name, max_tokens, client):
    """
    Splits answer into sentences, then decomposes each into sentence_parts + factoids.
    Returns (master_json_output, sentences).
    """
    user_prompt = SENTENCE_SPLIT_USER_TEMPLATE.format(answer=answer)
    try:
        res = query_model(SENTENCE_SPLIT_SYSTEM, user_prompt, model_name, max_tokens, client)
        sentences = json.loads(_clean_json(res["answer"]))
    except Exception as e:
        print(f"Sentence split error: {e}")
        return [], []

    master_json_output = []
    for i, sentence in enumerate(sentences, 1):
        user_prompt2 = FACTOID_USER_TEMPLATE.format(sentence=sentence)
        try:
            res2 = query_model(FACTOID_SYSTEM, user_prompt2, model_name, max_tokens, client)
            parsed = json.loads(_clean_json(res2["answer"]))
            master_json_output.append(parsed)
        except Exception as e:
            print(f"Factoid extraction error for sentence {i}: {e}")
            master_json_output.append([])

    return master_json_output, sentences


def parse_json_to_dictionaries(master_json_output, original_sentences):
    sentence_part_to_factoids = {}
    factoid_to_sentence_part = {}
    sentence_parts = []
    factoids_list = []

    for sentence_data, original_sentence in zip(master_json_output, original_sentences):
        if not isinstance(sentence_data, list):
            print(f"WARNING: skipping non-list factoid data: {type(sentence_data).__name__}")
            continue
        for item in sentence_data[:2]:
            if not isinstance(item, dict):
                print(f"WARNING: skipping non-dict sentence part: {item!r}")
                continue
            s_part = str(item.get("sentence_part", "") or "").strip()
            extracted_factoids = item.get("factoids", [])
            if not isinstance(extracted_factoids, list):
                extracted_factoids = []

            if not s_part:
                continue
            if s_part not in original_sentence:
                print(f"WARNING: sentence_part not substring of original: '{s_part}'")

            constrained_factoids = [str(f).strip() for f in extracted_factoids if str(f).strip()][:2]
            if not constrained_factoids:
                continue

            sentence_parts.append(s_part)
            sentence_part_to_factoids[s_part] = constrained_factoids
            for factoid in constrained_factoids:
                factoids_list.append(factoid)
                factoid_to_sentence_part[factoid] = s_part

    return sentence_part_to_factoids, factoid_to_sentence_part, sentence_parts, factoids_list


def generate_Q_questions(factoid, original_answer, Q, model_name, max_tokens, client, enable_sleep=True):
    user_prompt = Q_QUESTIONS_USER_TEMPLATE.format(
        Q=Q, original_answer=original_answer, factoid=factoid
    )
    max_retries = 3
    for attempt in range(max_retries):
        try:
            res = query_model(Q_QUESTIONS_SYSTEM, user_prompt, model_name, max_tokens, client)
            questions = json.loads(_clean_json(res["answer"]))
            return [str(q).strip() for q in questions][:Q]
        except Exception as e:
            err = str(e).lower()
            if enable_sleep and any(x in err for x in ("429", "rate limit", "too many requests")):
                print(f"Rate limit, sleeping 20s (attempt {attempt+1}/{max_retries})")
                time.sleep(20)
            else:
                print(f"Q-question error: {e}")
                if attempt == max_retries - 1 or not enable_sleep:
                    return []
    return []


def generate_M_answers(question, M, model_name, max_tokens, client, enable_sleep=True):
    user_prompt = M_ANSWERS_USER_TEMPLATE.format(question=question)
    sub_answers = []
    for i in range(M):
        max_retries = 3
        for attempt in range(max_retries):
            try:
                res = query_model(M_ANSWERS_SYSTEM, user_prompt, model_name, max_tokens, client)
                sub_answers.append(res["answer"].strip())
                break
            except Exception as e:
                err = str(e).lower()
                if enable_sleep and any(x in err for x in ("429", "rate limit", "too many requests")):
                    print(f"Rate limit, sleeping 20s (attempt {attempt+1}/{max_retries})")
                    time.sleep(20)
                else:
                    print(f"M-answer error (call {i+1}): {e}")
                    if attempt == max_retries - 1 or not enable_sleep:
                        sub_answers.append("ERROR")
                        break
    return sub_answers