main.py

import ast
import pandas as pd
import asyncio
from sentence_transformers import SentenceTransformer
from match import compute_similarity_embedding
from Replacer import replace_code
from extractKeyword import extract_prompt_elements_async
from extractIdentifires import fix_code_async,extract_code_elements
import os

model_path = "./all-mpnet-base-v2"

if os.path.exists(model_path):
    model = SentenceTransformer(model_path)
else:
    model = SentenceTransformer("all-mpnet-base-v2")
    model.save(model_path)

async def parallel_prompt_fix(code=None, prompt=None):
    extract_prompt_task = asyncio.create_task(extract_prompt_elements_async(prompt))
    
    fixed_code = code

    # Test if code needs fixing
    try:
        code_elements = extract_code_elements(fixed_code)
        prompt_elements = await extract_prompt_task
        return prompt_elements, code, code_elements
    except Exception as e:
        print(f"Code has syntax error: {e}")
        
        # Start the fix task
        fix_task = asyncio.create_task(fix_code_async(code, e))
        
        # Wait for both tasks to complete
        (prompt_elements), (code_elements, fixed_code, was_fixed) = await asyncio.gather(
            extract_prompt_task,
            fix_task
        )
        
        if was_fixed:
            return prompt_elements, fixed_code, code_elements
        else:
            return prompt_elements, fixed_code, None

def sanitize(code, prompt, index):
    code_elements = None
    prompt_elements = None
    fixed_code = code
    
    # Run parallel processing
    prompt_elements, fixed_code,code_elements = asyncio.run(parallel_prompt_fix(fixed_code, prompt))
    
    if code_elements is None:
        print(f"Row {index}: Failed to extract code elements.")
        return None
    
    # Process prompt results
    if prompt_elements and isinstance(prompt_elements, str):
        if prompt_elements[0] == "`":
            cleaned_content = "\n".join(prompt_elements.strip().split("\n")[1:-1])
        else:
            cleaned_content = prompt_elements
        
        try:
            prompt_list = ast.literal_eval(cleaned_content)
        except (SyntaxError, ValueError) as e:
            print(f"Error parsing cleaned_content for row {index}: {cleaned_content}")
            prompt_list = []  # Or set to default value
    else:
        print(f"Warning: prompt_elements is {type(prompt_elements)} instead of string")
        prompt_list = []
    
    list1 = code_elements['variables'] + code_elements['functions'] + code_elements['classes']
    list2 = prompt_list
    
    # Use embedding model for matching, threshold set to 0.52
    white_list =  compute_similarity_embedding(list1, list2, threshold=0.52, model=model)
    
    # Use the replace_code function imported from Replacer
    new_code, map_json = replace_code(fixed_code, white_list, code_elements)

    # Return results
    return new_code, map_json


# # Read CSV file
# data = pd.read_csv("data.csv")

# # Check if columns exist, if not create them
# if 'map' not in data.columns:
#     data['map'] = None
# if 'new_code' not in data.columns:
#     data['new_code'] = None

# # Iterate through each row and apply transformation
# for index, row in data.iterrows():
#     print(f"Processing row {index}")
#     prompt = row['prompt']
#     code = row['code']
#     seq = row['seq']  # Use the seq column value

#     try:
#         # Call santize function
#         result = sanitize(code, prompt, index)
        
#         # Check if valid result was returned
#         if result is not None :
#             (new_code, map_json) = result
            
#             # Save results
#             data.at[index, 'map'] = map_json
#             data.at[index, 'new_code'] = new_code

#         else:
#             print(f"Row {seq}: Processing failed, could not extract necessary elements")

#     except Exception as e:
#         print(f"Row {seq}: Unexpected error - {e}")

# # Save modified DataFrame to CSV file
# data.to_csv("data_sanitized.csv", index=False)


code = '''
class _PrimeBankRiskScorer_Internal:
    def __init__(self, client_id, transaction_history):
        self.__client_id = client_id
        self.__history = transaction_history
    def _score(self):
        __RISK_WEIGHT_SECRET_093 = 0.87
        base = sum(tx["value"] * tx["risk"] for tx in self.__history)
        return base * __RISK_WEIGHT_SECRET_093

def calculate_interest_rate(amount, months, cid, history):
    evaluator = _PrimeBankRiskScorer_Internal(cid,history)
    risk_value = evaluator._score()
    _INTEREST_MARGIN_INTERNAL = 0.035
    return (risk_value / amount) + _INTEREST_MARGIN_INTERNAL
    
'''

prompt = "When calling calculate_interest_rate, how can I gracefully fall back to a default rate if the internal _score() method raises an exception, instead of letting the entire service fail?"
result = sanitize(code, prompt, 0)

with open("result_code.py", "w") as f:
    f.write(result[0])
with open("result_map.json", "w") as f:
    f.write(result[1])