GuardAgent-code/config.py at main · zlzGithub-0801/GuardAgent-code · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def model_config(model):
    if model == 'gpt-3.5-turbo':
        config = {
            "model": "gpt-3.5-turbo",
            "api_key": "<YOUR_API_KEY>",
        }
    else:
        config = {
            "model": "gpt-4",
            "api_key": "<YOUR_API_KEY>",
        }

    return config


def llm_config_list(seed, config_list):
    llm_config_list = {
        "functions": [
            {
                "name": "python",
                "description": "run the entire code and return the execution result. Only generate the code.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "cell": {
                            "type": "string",
                            "description": "Valid Python code to execute.",
                        }
                    },
                    "required": ["cell"],
                },
            },
        ],
        "config_list": config_list,
        "timeout": 120,
        "cache_seed": seed,
        "temperature": 0,
    }
    return llm_config_list