greydgl--pentestgpt
Grey_D
58 行
1.5 KiB
Python
58 行
1.5 KiB
Python
import json
|
|
import re
|
|
|
|
from loguru import logger
|
|
|
|
|
|
class Message:
|
|
def __init__(self, content: str):
|
|
self.content = content
|
|
self.role = None
|
|
|
|
def get_content(self):
|
|
raise NotImplementedError("get_content() not implemented")
|
|
|
|
|
|
class SystemMessage(Message):
|
|
def __init__(self, content: str):
|
|
super().__init__(content)
|
|
self.role = "system"
|
|
|
|
def get_content(self):
|
|
return self.content
|
|
|
|
|
|
class UserMessage(Message):
|
|
def __init__(self, content: str):
|
|
super().__init__(content)
|
|
self.role = "user"
|
|
|
|
def get_content(self):
|
|
return self.content
|
|
|
|
|
|
class AssistantMessage(Message):
|
|
def __init__(self, content: str, is_json: bool = False):
|
|
super().__init__(content)
|
|
self.is_json = is_json
|
|
self.role = "assistant"
|
|
|
|
def get_content(self):
|
|
json_patterns = [r"```json\s*(\{.*?\})\s*```", r"```json\s*(\[.*?\])\s*```"]
|
|
for pattern in json_patterns:
|
|
match = re.search(pattern, self.content, re.DOTALL)
|
|
if match:
|
|
formatted_json = match.group(1)
|
|
try:
|
|
return json.loads(formatted_json)
|
|
except Exception as e:
|
|
logger.exception(f"Failed to parse JSON content: {e}")
|
|
|
|
if self.is_json:
|
|
try:
|
|
return json.loads(self.content)
|
|
except Exception as e:
|
|
logger.exception(f"Failed to parse JSON content: {e}")
|
|
return {}
|
|
return self.content
|