58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
import json
|
|
import re
|
|
|
|
from json_minify import json_minify
|
|
from json_repair import repair_json
|
|
|
|
|
|
class ValidationException(Exception):
|
|
def __init__(self, message: str):
|
|
self.message = message
|
|
|
|
|
|
def parse_message_json(message: str) -> dict:
|
|
message = message.strip()
|
|
code_block_pattern = r"```json\s*([\s\S]+?)```"
|
|
code_block_match = re.search(code_block_pattern, message, re.DOTALL)
|
|
|
|
if code_block_match:
|
|
json_str = code_block_match.group(1).strip()
|
|
else:
|
|
bracket_pattern = r"\{.*\}"
|
|
bracket_match = re.search(bracket_pattern, message, re.DOTALL)
|
|
if not bracket_match:
|
|
raise ValidationException("Response does not have correct json format")
|
|
json_str = bracket_match.group(0).strip()
|
|
|
|
try:
|
|
json_str = json_minify(json_str)
|
|
data = json.loads(json_str)
|
|
except json.JSONDecodeError:
|
|
try:
|
|
json_str = repair_json(json_str)
|
|
data = json.loads(json_str)
|
|
except json.JSONDecodeError:
|
|
raise ValidationException("Response does not have correct json format")
|
|
return data
|
|
|
|
|
|
class GroundingOutput:
|
|
def __init__(
|
|
self,
|
|
description: str,
|
|
position: tuple[int, int],
|
|
end_position: tuple[int, int] = None,
|
|
):
|
|
self.description = description
|
|
self.position = position
|
|
self.end_position = end_position
|
|
|
|
|
|
class GroundingRequest:
|
|
def __init__(
|
|
self, description: str, image_base64: str, action_type: str | None = None
|
|
):
|
|
self.description = description
|
|
self.image_base64 = image_base64
|
|
self.action_type = action_type
|