pyt

from pyrit.common import IN_MEMORY, initialize_pyrit
from pyrit.orchestrator import PromptSendingOrchestrator
from pyrit.prompt_target.http_target.http_target import (
  HTTPTarget, 
)
from pyrit.prompt_target import (
    get_http_target_json_response_callback_function,
    get_http_target_regex_matching_callback_function,
)

import requests
import sys
import asyncio

import logging
logging.basicConfig(level=logging.DEBUG)

    
initialize_pyrit(memory_db_type=IN_MEMORY)


url = "http://ollama:11434/api/chat"
api_key = "test"

raw_http_request = f"""
    POST {url} HTTP/1.1
    Content-Type: application/json

    {{
        "model": "llama3",
        "messages": [
            {{"role": "user", "content": "{{PROMPT}}" }}
        ],
        "max_tokens": 50,
        "temperature": 0.2,
        "stream": false
    }}
"""
parsing_fn = get_http_target_json_response_callback_function(key="message.content")

print(raw_http_request)

http_target = HTTPTarget(
    http_request=raw_http_request,
    prompt_regex_string="{PROMPT}",   # literal substring to replace
    callback_function=parsing_fn,
    use_tls = False,
    timeout=3600.0
)


orchestrator = PromptSendingOrchestrator(objective_target=http_target, verbose=True)

responses = await orchestrator.send_prompts_async(
  prompt_list=[
    "How do I bake a cake?",
    "Tell me about the French Revolution",
  ]
)
    
await orchestrator.print_conversations_async()


print(f"{"-"*100}")