JetBrains-Research/commit-message-editing-visualization

1import pickle2import time3 4from grazie.api.client.chat.prompt import ChatPrompt5from grazie.api.client.endpoints import GrazieApiGatewayUrls6from grazie.api.client.gateway import AuthType, GrazieAgent, GrazieApiGatewayClient7from grazie.api.client.profiles import LLMProfile8 9import config10 11client = GrazieApiGatewayClient(12    grazie_agent=GrazieAgent("grazie-toolformers", "v1.0"),13    url=GrazieApiGatewayUrls.STAGING,14    auth_type=AuthType.APPLICATION,15    grazie_jwt_token=config.GRAZIE_API_JWT_TOKEN,16)17 18LLM_CACHE_FILE = config.CACHE_DIR / f"{config.LLM_MODEL}.cache.pkl"19LLM_CACHE = {}20LLM_CACHE_USED = {}21 22if not LLM_CACHE_FILE.exists():23    with open(LLM_CACHE_FILE, "wb") as file:24        pickle.dump(obj=LLM_CACHE, file=file)25 26with open(LLM_CACHE_FILE, "rb") as file:27    LLM_CACHE = pickle.load(file=file)28 29 30def llm_request(prompt):31    output = None32 33    while output is None:34        try:35            output = client.chat(36                chat=ChatPrompt().add_system("You are a helpful assistant.").add_user(prompt),37                profile=LLMProfile(config.LLM_MODEL),38            ).content39        except Exception:40            time.sleep(config.GRAZIE_TIMEOUT_SEC)41 42    assert output is not None43 44    return output45 46 47def generate_for_prompt(prompt):48    if prompt not in LLM_CACHE:49        LLM_CACHE[prompt] = []50 51    if prompt not in LLM_CACHE_USED:52        LLM_CACHE_USED[prompt] = 053 54    while LLM_CACHE_USED[prompt] >= len(LLM_CACHE[prompt]):55        new_response = llm_request(prompt)56        LLM_CACHE[prompt].append(new_response)57 58        with open(LLM_CACHE_FILE, "wb") as file:59            pickle.dump(obj=LLM_CACHE, file=file)60 61    result = LLM_CACHE[prompt][LLM_CACHE_USED[prompt]]62    LLM_CACHE_USED[prompt] += 163 64    return result65