|
4 | 4 | import os |
5 | 5 | import inspect |
6 | 6 | import importlib |
7 | | -from pyrit.common import default_values, initialize_pyrit, IN_MEMORY |
| 7 | +from pyrit.common import default_values |
| 8 | +from pyrit.setup import initialize_pyrit_async, IN_MEMORY |
8 | 9 | from pyrit.prompt_converter import PromptConverter |
9 | 10 | from pyrit.prompt_target import OpenAIChatTarget |
10 | | -from pyrit.orchestrator import PromptSendingOrchestrator |
| 11 | +from pyrit.executor.attack import PromptSendingAttack, MarkdownAttackResultPrinter |
11 | 12 | from pyrit.score import SelfAskTrueFalseScorer |
12 | 13 | from dotenv import load_dotenv |
13 | 14 |
|
@@ -52,22 +53,23 @@ def convert(converter_name:str): |
52 | 53 |
|
53 | 54 | @app.route('/prompt/generate', methods=['POST']) |
54 | 55 | def generate_prompt(): |
55 | | - # Initialize AOAI GPT-4o target |
| 56 | + # Initialize chat target |
56 | 57 | global chat_target |
57 | 58 | if (chat_target is None): |
58 | 59 | chat_target = initialize_chat_target() |
59 | 60 |
|
60 | | - promptSendingOrchestrator = PromptSendingOrchestrator(objective_target=chat_target) |
| 61 | + prompt_sending_attack = PromptSendingAttack(objective_target=chat_target) |
61 | 62 | # Extract input data from json payload |
62 | 63 | data = request.get_json() |
63 | 64 | prompt_goal = data['prompt_goal'] |
64 | 65 |
|
65 | | - generated_prompt = asyncio.run(promptSendingOrchestrator.send_prompts_async(prompt_list=[prompt_goal]))[0].request_pieces[0].converted_value |
66 | | - return jsonify({"prompt": generated_prompt}) |
| 66 | + result = asyncio.run(prompt_sending_attack.execute_async(objective=prompt_goal)).last_response.converted_value |
| 67 | + |
| 68 | + return jsonify({"prompt": result}) |
67 | 69 |
|
68 | 70 | @app.route('/prompt/score/SelfAskTrueFalseScorer', methods=['POST']) |
69 | 71 | def score(): |
70 | | - # Initialize AOAI GPT-4o target |
| 72 | + # Initialize chat target |
71 | 73 | global chat_target |
72 | 74 | if (chat_target is None): |
73 | 75 | chat_target = initialize_chat_target() |
@@ -97,7 +99,7 @@ def score(): |
97 | 99 | ) |
98 | 100 |
|
99 | 101 | def initialize_chat_target(): |
100 | | - initialize_pyrit(memory_db_type=IN_MEMORY) |
| 102 | + asyncio.run(initialize_pyrit_async(memory_db_type=IN_MEMORY)) |
101 | 103 |
|
102 | 104 | chat_target = OpenAIChatTarget( |
103 | 105 | model_name=os.environ.get("OPENAI_CHAT_MODEL_NAME"), |
|
0 commit comments