test_gguf_llm_works.py

import libcst as cst
from collections import OrderedDict
from llama_cpp import Llama

class RenameTransformer(cst.CSTTransformer):
    def __init__(self):
        super().__init__()
        self.var_map = OrderedDict()
        self.func_map = OrderedDict()
        self.cnt_var = 0
        self.cnt_func = 0

    def _map_var(self, name: str) -> str:
        if name not in self.var_map:
            self.cnt_var += 1
            self.var_map[name] = f"VAR{self.cnt_var}"
        return self.var_map[name]

    def _map_func(self, name: str) -> str:
        if name not in self.func_map:
            self.cnt_func += 1
            self.func_map[name] = f"FUNC{self.cnt_func}"
        return self.func_map[name]

    def leave_FunctionDef(self, original_node, updated_node):
        new_name = self._map_func(original_node.name.value)
        return updated_node.with_changes(name=original_node.name.with_changes(value=new_name))

    def leave_Param(self, original_node, updated_node):
        original_var = original_node.name.value
        new_var = self._map_var(original_var)
        return updated_node.with_changes(
            name=updated_node.name.with_changes(value=new_var)
        )

    def leave_Name(self, original_node, updated_node):
        val = original_node.value
        if val in self.func_map:
            return updated_node.with_changes(value=self.func_map[val])
        if val not in ("print",):
            return updated_node.with_changes(value=self._map_var(val))
        return updated_node

def anonymize_with_libcst(src: str):
    module = cst.parse_module(src)
    transformer = RenameTransformer()
    new_module = module.visit(transformer)
    return new_module.code, transformer.var_map, transformer.func_map

def call_privacy_llm(prompt_code: str):
    llm = Llama(
        model_path="Your local model path",
        n_gpu_layers=4,
        n_ctx=4096,
        chat_format="deepseek",
        verbose=False
    )
    system_prompt = (
        '''You are PrivacyLLM, a formatter agent.
        You will receive fully anonymized code using only VAR* and FUNC* labels.
        **You must not rename or replace any VAR*/FUNC* tag.**
        Instead, if code is incomplete (missing returns or documentation), you may **append content inside existing functions** but keep labels unchanged.
        Input anonymized code:
        '''
    )
    full_prompt = system_prompt + "\n\nInput:\n" + prompt_code
    resp = llm(full_prompt, max_tokens=512)
    return resp["choices"][0]["text"]

if __name__ == "__main__":
    src = """
def add_numbers(x, y):
    return x + y

print(add_numbers(x, y))
"""
    anon_code, var_map, func_map = anonymize_with_libcst(src)
    print("===Anonymized Code===")
    print(anon_code)
    print("VAR_MAP:", var_map)
    print("FUNC_MAP:", func_map)

    print("\n===The prompt submitted to PrivacyLLM===")
    print(anon_code)

    print("\n===PrivacyLLM Output===")
    out = call_privacy_llm(anon_code)
    print(out)