Function get_secrets_LLM

deeppass2.py:130–167

Run LLM check and confirmation on the retrieved xlm-RoBERTa-base model

(bert_output)

    return prompt

def get_secrets_LLM(bert_output):
    """Run LLM check and confirmation on the retrieved xlm-RoBERTa-base model"""
    secrets_list = []
    litellm_model = os.environ["AUGMENT_MODEL"]

    start = time.time()

    for secrets in bert_output:
        passw = secrets["secret"]
        text = secrets["chunk"]
        prompt = get_prompt(text, passw)

        ## Change the code here to accomodate non-litellm infra
        response = litellm_client.chat.completions.create(
            model=litellm_model,
            messages = [
                {
                "role": "system",
                "content": "You are a helpful assistant. Your expertise is to identify the presence of APIs, database links, or leaked passwords. Do not explain."
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
        )
        execution_time = (time.time()-start)
        message = response.choices[0].message.content

        credential_pattern = r"<credential>(.*?)</credential>"
        credential_matches = list(re.finditer(credential_pattern, message))

        
        credentials = [match.group(1) for match in credential_matches]

        secrets_list.append({'Extracted_Credentials':credentials, 'LLM_response':message,'chunk':text, 'generation_time': execution_time})

    return secrets_list


def get_secrets_BERT(chunks, model=model, tokenizer=tokenizer):