gpt_inf.py

import os
from openai import OpenAI
from count_tokens import num_tokens_from_string, convert_to_string


def inference(prompt):
    client = OpenAI(
        # This is the default and can be omitted
        api_key=os.environ.get("OPENAI_API_KEY"),
        )
    if prompt:
        prompt_str=convert_to_string(prompt)
        number=num_tokens_from_string(prompt_str, "gpt-3.5-turbo")
        if number <=4096:
            response = client.chat.completions.create(
                model="gpt-3.5-turbo",
                messages=prompt,
                temperature=0.3
                )
            return response.choices[0].message.content.strip()
        else :
            return 'null'


def sp_inference(query, codeA, codeB, context):
    prompt= [context]
    if codeA and codeB: 
        prompt.append(
            {
                "role": "user", 
                "content": f"""{query} \n Code A: ``` {codeA} ``` \n Code B: ```{codeB}``` \n"""
            },
        )
        response = inference(prompt=prompt)
        return response
        
def sct_inference(desc1, desc2, context):
    prompt= [context]
    if desc1 and desc2: 
        prompt.append(
            {
                "role": "user", 
                "content": f"""Analyze the following functions of two code snippets and determine if they are code clones, regardless of the programming language. The function of the first code snippet is ``` {desc1} ``` and the function of the second is ``` {desc2} ```. Please answer ‘yes’ if the code snippets are clones or ‘no’ if they are not."""
            },
        )
        response = inference(prompt=prompt)
        return response
        
def get_description(query, code, context):
    prompt= [context]
    if query and code: 
        prompt.append(
            {
                "role": "user", 
                "content": f"""{query} \n ``` {code} ```"""
            },
        )
        response = inference(prompt=prompt)
        return response