import logging
import re
from typing import Any, List, Optional

import langchain
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain_openai import ChatOpenAI

from init_openai import init_openai

logger = logging.getLogger("SoCloverAI")
init_openai()
model_name = "gpt-4-1106-preview"


def set_trial(trial: int) -> None:
    langchain.llm_cache.inner_cache.set_trial(trial)


def dump_cache_stats_since_last_call() -> None:
    logger.info(langchain.llm_cache.get_cache_stats_summary())
    langchain.llm_cache.clear_cache_stats()


def create_llm_model(temperature: float, model_name: str) -> ChatOpenAI:
    # mypy seems confused about the model_name parameter:
    #   Unexpected keyword argument "model_name" for "ChatOpenAI"
    result = ChatOpenAI(temperature=temperature, model_name=model_name)  # type: ignore
    return result


async def predict(temperature: float, template: str, **kwargs: Any) -> List[str]:
    prompt = PromptTemplate(
        template=template.strip(), input_variables=["word0", "word1"]
    )
    llm = create_llm_model(temperature, model_name)
    chain = LLMChain(llm=llm, prompt=prompt, verbose=False)
    output = await chain.apredict(**kwargs)
    logger.debug(output)
    predictions = parse_candidates(output)
    best = parse_best(output)
    if best:
        predictions = [best] + predictions

    strip_chars = ' \t"'
    predictions = [prediction.strip(strip_chars) for prediction in predictions]
    predictions = [prediction for prediction in predictions if prediction]

    # remove duplicates while preserving order
    seen = set()
    unique_predictions = list()
    for prediction in predictions:
        if prediction not in seen:
            unique_predictions.append(prediction)
            seen.add(prediction)
    predictions = unique_predictions
    return predictions


def parse_candidates(output: str) -> List[str]:
    result = []
    for line in output.splitlines():
        if not line.startswith("Candidates:"):
            continue
        candidates_str = line[len("Candidates: ") :]
        candidates = candidates_str.split(",")
        candidates = [candidate.strip() for candidate in candidates]
        result += candidates
    return result


pattern = re.compile(r"Best: (.*)")


def parse_best(output: str) -> Optional[str]:
    match = pattern.search(output)
    if match:
        return match.group(1)
    split_output = output.split()
    if len(split_output) == 1:
        logger.info(f"Invalid output format: {output}")
        return split_output[0]
    logger.info(f"Invalid output: {output}")
    return None

langchain.llm_cache.get_cache_stats_summary langchain_openai.ChatOpenAI langchain.llm_cache.inner_cache.set_trial langchain.llm_cache.clear_cache_stats langchain.chains.LLMChain
[((252, 283), 'logging.getLogger', 'logging.getLogger', (['"""SoCloverAI"""'], {}), "('SoCloverAI')\n", (269, 283), False, 'import logging\n'), ((284, 297), 'init_openai.init_openai', 'init_openai', ([], {}), '()\n', (295, 297), False, 'from init_openai import init_openai\n'), ((2273, 2297), 're.compile', 're.compile', (['"""Best: (.*)"""'], {}), "('Best: (.*)')\n", (2283, 2297), False, 'import re\n'), ((373, 421), 'langchain.llm_cache.inner_cache.set_trial', 'langchain.llm_cache.inner_cache.set_trial', (['trial'], {}), '(trial)\n', (414, 421), False, 'import langchain\n'), ((539, 578), 'langchain.llm_cache.clear_cache_stats', 'langchain.llm_cache.clear_cache_stats', ([], {}), '()\n', (576, 578), False, 'import langchain\n'), ((791, 849), 'langchain_openai.ChatOpenAI', 'ChatOpenAI', ([], {'temperature': 'temperature', 'model_name': 'model_name'}), '(temperature=temperature, model_name=model_name)\n', (801, 849), False, 'from langchain_openai import ChatOpenAI\n'), ((1137, 1184), 'langchain.chains.LLMChain', 'LLMChain', ([], {'llm': 'llm', 'prompt': 'prompt', 'verbose': '(False)'}), '(llm=llm, prompt=prompt, verbose=False)\n', (1145, 1184), False, 'from langchain.chains import LLMChain\n'), ((488, 533), 'langchain.llm_cache.get_cache_stats_summary', 'langchain.llm_cache.get_cache_stats_summary', ([], {}), '()\n', (531, 533), False, 'import langchain\n')]