api.py

import os
import pickle
from datetime import datetime
from pathlib import Path

from fishaudio import FishAudio
from fishaudio.utils import save
from openai import OpenAI
from tenacity import retry, stop_after_attempt, wait_exponential

try:
    from dotenv import load_dotenv

    env_path = Path(os.path.dirname(__file__)) / ".env"
    if env_path.exists():
        load_dotenv(env_path)
except ImportError:
    pass

# params
INFO_audio_ID_dict = {
    "matsu": os.getenv("INFO_audio_matsu"),
    "ken": os.getenv("INFO_audio_ken"),
    "mana": os.getenv("INFO_audio_mana"),
    "cai": os.getenv("INFO_audio_cai"),
    "ren": os.getenv("INFO_audio_ren"),
}
# print(INFO_audio_ID_dict)

client = OpenAI(api_key=os.getenv("gpt"))
#
audio_client = FishAudio(api_key=os.getenv("audio"))

PKL_FILE = "./resource/knowledge_data.pkl"


def load_pkl(file_path):
    if os.path.exists(file_path):
        with open(file_path, "rb") as f:
            return pickle.load(f)
    else:
        return {}


def save_pkl(file_name, data):
    sorted_data = dict(
        sorted(
            data.items(),
            key=lambda x: datetime.strptime("-".join(x[0].split("-")[0:3]), "%Y-%m-%d"),
        )
    )
    with open(file_name, "wb") as f:
        pickle.dump(sorted_data, f)


@retry(
    stop=stop_after_attempt(5),
    wait=wait_exponential(multiplier=1, min=2, max=10),
)
def generate_item(prompt, sys_prompt, model="gpt-5.4-mini", temperature=1.0):
    # print("prompt=", prompt)
    # print("temperature=", temperature)
    response = client.chat.completions.create(
        model=model,
        temperature=temperature,
        messages=[
            {
                "role": "system",
                "content": sys_prompt,
            },
            {"role": "user", "content": prompt},
        ],
    )
    return response.choices[0].message.content.strip()


@retry(
    stop=stop_after_attempt(5),
    wait=wait_exponential(multiplier=1, min=2, max=10),
)
def get_embedding(text):
    response = client.embeddings.create(input=text, model="text-embedding-3-small")
    return [emb.embedding for emb in response.data]


@retry(
    stop=stop_after_attempt(5),
    wait=wait_exponential(multiplier=1, min=2, max=10),
)
def audio_text(text, audio_path, speaker_id):
    # TODO change speaker
    # print(speaker_id, ":", INFO_audio_ID_dict[speaker_id])
    audio = audio_client.tts.convert(
        text=text,
        reference_id=INFO_audio_ID_dict[speaker_id],
        format="mp3",
        model="s2-pro",
    )
    save(audio, audio_path)

    # with open(audio_path, "wb") as f:
    #     for chunk in audio_client.tts(
    #         TTSRequest(
    #             reference_id=INFO_audio_ID_dict[speaker_id], text=text, latency="normal"
    #         ),
    #         backend="s1",
    #     ):
    #         f.write(chunk)
    return audio_path