Q Stateful PromptTarget

[grant-conine]

Is your feature request related to a problem? Please describe.

I'm working on a newly started enterprise Red Team and have been exploring PyRIT. I've found there's a class of prompt targets that fit somewhere between a PromptTarget and PromptChatTarget - LLM applications where it is possible to maintain a conversation history, but not possible to edit the chat history. Think a PlaywrightTarget where the target application maintains the conversation history, managing context internally, or an HTTPTarget or custom Target where the first response contains a conversation_id which allows for continued conversation in a shared context, or a GPT Realtime API where a websocket connection maintains conversation context but doesn't allow for editing the conversation history.

It's been difficult to red team custom LLM applications which don't align to a pre-set openAI / azure ML endpoints even if they're constructed on those tools, because the most successful attacks (crescendo) require a PromptChatTarget

Describe the solution you'd like

Could we create a core class between PromotTarget and PromptChatTarget (maybe called StatefulPromotTarget?) that outlines abstract methods for setting and resetting chat contexts, and modify attack strategies to check if these sets and resets exist and call them in their setup and teardown steps?

This would allow a PlayerightTarget to establish a context when the Page is loaded, and reset the context by refreshing a page, or an Custom Target establishing the context with an initial prompt (to get and set a conversation_id) and reset the context by flushing this conversation_id. This would be a step between "send multiple one-off attacks" and "have full control over conversation history" and enable a simplified crescendo and skeleton key style attacks even in cases where maintaining conversational context is a bit more complex.

Describe alternatives you've considered, if relevant

I've sort of worked around this by monkey-patching AttackStrategy._teardown_async method to refresh the PlaywrightTarget's Page during teardown but doing the same with an HTTPTarget led me to think about a broader solution.

Additional context

I am trying to get enterprise permission to contribute to this project, but for now this is all I can do.

[romanlutz]

@grant-conine thanks for reaching out! That is actually doable even today. Check out https://azure.github.io/PyRIT/code/targets/playwright_target_copilot.html for example. It keeps the state but sends only the latest message. PromptChatTarget does not necessarily mean that it will send the entire conversation history. That's really up to the implementation. There are lots of cases where we can't actually edit the history.

Now to the solution part: You decide what you're passing to send_prompt_async. If that's a message with an existing conversation ID then you're continuing that conversation. If it's a new conversation ID then it's a new conversation. Unless I'm missing something (which is entirely possible) this is all doable today. You may want to check out the page I linked above.

Would love to have you contribute in this direction if your employer permits it 🙂 PlawrightTarget is definitely one of those areas that don't work quite as well as I would hope and there's so much potential.

[grant-conine]

@grant-conine thanks for reaching out! That is actually doable even today. Check out https://azure.github.io/PyRIT/code/targets/playwright_target_copilot.html for example. It keeps the state but sends only the latest message. PromptChatTarget does not necessarily mean that it will send the entire conversation history. That's really up to the implementation. There are lots of cases where we can't actually edit the history.

Now to the solution part: You decide what you're passing to send_prompt_async. If that's a message with an existing conversation ID then you're continuing that conversation. If it's a new conversation ID then it's a new conversation. Unless I'm missing something (which is entirely possible) this is all doable today. You may want to check out the page I linked above.

Would love to have you contribute in this direction if your employer permits it 🙂 PlawrightTarget is definitely one of those areas that don't work quite as well as I would hope and there's so much potential.

Ah I hadn't seen that as I was looking at the docs for 0.10.0rc0.

I think my question still stands - if I want to isolate different attacks when I have no control over how the conversation history is maintained, but a way to "reset" the conversation (a easy example is a PlaywrightTarget, but other examples could be an agent built in Google ADK, Langchain, DSPy, etc), workflows would get a lot easier if that reset could be made to be part of an AttackStrategy's _teardown_async method.

The first app I readteamed I just used a PlaywrightTarget and refreshed the page between attacks by patching the refresh operation onto the AttackStrategy's teardown step, but it felt hacky. It worked incredibly well though. I was able to successfully demonstrate a vulnerability to a crescendo attack by a false customer persona.

Now I'm working on another app that has a really simple API - the first response has a conversation_id which I can use to persist the conversation across multiple calls (and is really easy to implement) but again, I want to change the conversation_id between attacks so everything is coming from a similar starting point, and it felt like I was about to repeat the same hack. That makes me think either (1) I'm not understanding the approach I should be using for PromptTargets or (2) There's potentially an improvement to be made in the package.

My colleague is checking out promptfoo parallel to me and they have a "afterEach" hook which easily allows a page to be refreshed or a conversation_id set back to "" - I was hoping something similar could be found in PyRIT to reset the conversation history when I don't actually have a lot of control over it (managed by targeted app/ai agent)

PS: Yes, I could reset the entire attack for each prompt, but it seems potentially wasteful and a silly way to handle this. Just wondering what solutions others have found or if the team is open to an extension of PromptTarget to allow inter-attack cleanup.

[romanlutz]

I see your point for reusing a PlaywrightTarget. The page would need to be reset between each attack execution, and we wouldn't be able to run this in parallel (at least not as is, because it would require multiple browsers). I need to check where the best spot would be for such a teardown. @bashirpartovi or @rlundeen2 may already have thoughts. Attacks have a _teardown_async method which could be utilized to call a target teardown method (which does not exist yet) for tearing down a specific conversation. We'd just need to add that sort of thing and for most targets it wouldn't do anything.

Regarding your HTTPTarget example, the attack object should be using a different conversation ID per conversation. If you have a minimal example I'm happy to investigate. HTTPTarget only works for individual prompts right now.

For most endpoints, there is no real cleanup. The moment you switch the conversation ID the target won't find any records of the conversation in the database and start a new conversation.

[grant-conine]

I see your point or reusing a PlaywrightTarget. The page would need to be reset between each attack execution, and we wouldn't be able to run this in parallel (at least not as is, because it would require multiple browsers). I need to check where the best spot would be for such a teardown. @bashirpartovi or @rlundeen2 may already have thoughts. Attacks have a _teardown_async method which could be utilized to call a target teardown method (which does not exist yet) for tearing down a specific conversation. We'd just need to add that sort of thing and for most targets it wouldn't do anything.

Regarding your HTTPTarget example, the attack object should be using a different conversation ID per conversation. If you have a minimal example I'm happy to investigate. HTTPTarget only works for individual prompts right now.

For most endpoints, there is no real cleanup. The moment you switch the conversation ID the target won't find any records of the conversation in the database and start a new conversation.

Here's some minimal examples.

PlaywrightTarget modification

This is more-or-less how I set up a PlayWright target class extension.

class ResetFunction(Protocol):
    async def __call__(self, page: "Page") -> None: ...

Class ResettingPlaywrightTarget(PlaywrightTarget):
    
    def __init__(
        self, 
        *,
        interaction_func: InteractionFunction,
        reset_func: ResetFunc,
        page: "Page",
        max_requests_per_minute: Optional[int],
    ) -> None:
        super().__init__(interaction_func=interaction_func, page=page, max_requests_per_minute=max_requests_per_minute)
        self._reset_func = reset_func

    async def reset_target_async(self) -> None:
        self._reset_func(self.page)

Custom Target

This is more or less how I set up a custom target. An HTTPTarget could work with logic similar to the {PROMPT} regex to insert the external app's conversation ID, but I don't think that would necessarily be that consistent from target to target so I probably wouldn't include the logic in a base class.

class ResettingTarget(PromptTarget):
    """Imagine this target takes a request like
        {
            "message": "{PROMPT}",
            "conversation_id": ""
        }

        On the first message, no conversation id is sent, and the response comes back
        {
            "conversation_id": "123-abc-456-def",
            "messages": ["content": "{RESPONSE}"]
        }

        If I put the conversation_id from the response into subsequent messages, the target application treats it
        as part of the same conversation. If I send a new message without the conversation_id, it treats it as a new conversation

        Yes, it's confusing that the application also has a conversation_id, and it is distinct from PyRIT's, but it's a common concept.
    """

    def __init__(
        self,
        *,
        endpoint: str,
        max_requests_per_minute: Optional[int],
    ) -> None:
        super().__init__(endpoint=endpoint, max_requests_per_minute=max_requests_per_minute)
        self._external_conversation_id = ""
    
    def _validate_request(self, *, message: Message) -> None:
        # validation logic here
 
    @limit_requests_per_minute
    async def send_prompt_async(self, *, message: Message) -> list[Message}:
        self._validate_request(message)
        response = await self._complete_text_async(message)
        
        response_json = response.json()

        if self._external_conversation_id == "":
            self._external_conversation_id = response.get("conversation_id")
        response_text = response.get("messages")[0].get("content")
        
        response_entry = construct_response_from_request(request=request, response_text_pieces=[response_text])
        return [response_entry]

    async def reset_target_async(self) -> None:
        self._external_conversation_id = ""
    
    # More logic to actually make the requests with net_utility, using the externaal conversation id and do other things could follow.

Adding resetting logic

What I did, that felt really hacky was...

def patch_resetting_logic(attack_strategy: AttackStrategy) -> AttackStrategy:
    old_teardownattack_strategy._teardown_async
    
    def new_teardown(**kwargs) -> None:
        await old_teardown(**kwargs)
        if hasattr(attack_strategy.objective_target, "reset_target_async"):
            await attack_strategy.objective_target.reset_target_async
    attack_strategy._teardown_async = new_teardown

I'd love to see a broader change that negates the need for the weird patching that I did. Maybe a mixin?

[grant-conine]

I think the main use case here is more applications built on an LLM but which don't adhere to any strict openai apis. I see a lot of these being made right now, and limited ways to successfully test them for compliance,

[rlundeen2]

Recap/Background

So PromptChatTargets you have full control of the history, and with PromptTarget you have none, but can always start a new conversation. We should probably have another designation to distinguish things like StorageAccountTarget where you don't get a response. But let's table this for now and focus on the targets where we get responses that you bring up here @grant-conine; e.g. differences between Realtime/HTTP/Playwright (have responses and can't modify history) and OpenAI/AML (where you can modify history).

I don't think there is fundamentally an "in-between" type of target between PromptTarget and PromptChatTarget exactly. You can modify history with them or not. But there are still opportunities to both figure out how to modify history in some cases, or modify the attack strategies in cases where you can't.

Resetting Targets

From the perspective of a target, every time you set a new conversation id should be a brand new conversation with reset context. I know with RealTimeTarget, Bolor/Roman gave that some thought as to how to manage that and when to create new connections. So really, every time conversation id is new, it should be reset.

@grant-conine brings up an excellent point with HTTPTarget and Playwright, where it doesn't do this; the target is initialized with auth and a request; so most likely a new conversation ID would just continue the existing conversation. Probably not possible to generically reset, and definitely not ideal.

Opportunities - these are all categories of work I would absolutely love

1. Adding a designator for non AI targets: We should be able to better distinguish between something like a blob store target and RealTimeTarget. Sure, both can't modify history. But with a Blob target there isn't any history.
2. Any target improvement that adds history support, even if flexibly: Sometimes, we may be able to modify a PromptTarget so we CAN modify history. E.g. Even though HTTPTarget will never be able to support history modifications in a completely generic way, it could do so sometimes (e.g. if the HTTP call supplies the history somehow). And it would be good to be flexible in these cases so you can use it as a PromptChatTarget. If you have a way to expand the conversation history in a request and pass that in, all PromptChatTarget attack strategies will work. As an aside, how I would do this is a new class like HTTPChatTarget(HTTPTarget, PromptChatTarget) that takes a required callback on how to expand the history.
3. Explicit target logic to reset conversation when there is a new conversation id: It's the responsibility of a target to do this - and really EVERY target should be a resetting target. Despite not being able to modify history, RealTimeTarget keeps a connection dictionary for new sockets when there is a new conversation id.
   • HTTPTarget and Playwright target have unique challenges, but we could tackle it in a similar way. "If the conversation id is new, use this callback function which changes my auth tokens"
   • It might also be nice to have some configurable error here. If you can't reset the history, but get a new conversation id, that might be nice to know so you don't get unexpected behavior. Because I do think the expectation that every conversation id is a new conversation.
4. Additional/modified non PromptChatTarget strategies: Many attacks strategies need you to be able to modify the history for them to work how they are written... but sometimes they could potentially support variations if you can't modify history. For example, Crescendo with no backtracks would only need a PromptTarget. On the other hand, I think these are already pretty close... Because Crescendo without backtracks could essentially just be done with RedTeamAttack using a different system prompt.

Thanks for bringing this up @grant-conine! Great thoughts! Do these align with your thinking? I think #3 is most relevant to you.

For the PyRIT team, we have competing priorities the next few months so I don't see us tackling this in the near term, but I could see it in the next six months. But if any of these are things you're interested in tackling, let us know and happy to dive into guidance/brainstorming on implementation details.