chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,223 @@
|
||||
import json
|
||||
import re
|
||||
from typing import List
|
||||
import time
|
||||
import tiktoken
|
||||
|
||||
from typing import Tuple, List, Union, Dict
|
||||
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
import pickle
|
||||
|
||||
|
||||
class Node(BaseModel):
|
||||
name: str
|
||||
info: str
|
||||
|
||||
|
||||
class Dag(BaseModel):
|
||||
nodes: List[Node]
|
||||
edges: List[List[Node]]
|
||||
|
||||
|
||||
NUM_IMAGE_TOKEN = 1105 # Value set of screen of size 1920x1080 for openai vision
|
||||
|
||||
|
||||
def call_llm_safe(agent) -> Union[str, Dag]:
|
||||
# Retry if fails
|
||||
max_retries = 3 # Set the maximum number of retries
|
||||
attempt = 0
|
||||
response = ""
|
||||
while attempt < max_retries:
|
||||
try:
|
||||
response = agent.get_response()
|
||||
break # If successful, break out of the loop
|
||||
except Exception as e:
|
||||
attempt += 1
|
||||
print(f"Attempt {attempt} failed: {e}")
|
||||
if attempt == max_retries:
|
||||
print("Max retries reached. Handling failure.")
|
||||
time.sleep(1.0)
|
||||
return response
|
||||
|
||||
|
||||
def calculate_tokens(messages, num_image_token=NUM_IMAGE_TOKEN) -> Tuple[int, int]:
|
||||
|
||||
num_input_images = 0
|
||||
output_message = messages[-1]
|
||||
|
||||
input_message = messages[:-1]
|
||||
|
||||
input_string = """"""
|
||||
for message in input_message:
|
||||
input_string += message["content"][0]["text"] + "\n"
|
||||
if len(message["content"]) > 1:
|
||||
num_input_images += 1
|
||||
|
||||
input_text_tokens = get_input_token_length(input_string)
|
||||
|
||||
input_image_tokens = num_image_token * num_input_images
|
||||
|
||||
output_tokens = get_input_token_length(output_message["content"][0]["text"])
|
||||
|
||||
return (input_text_tokens + input_image_tokens), output_tokens
|
||||
|
||||
|
||||
# Code based on https://github.com/xlang-ai/OSWorld/blob/main/mm_agents/agent.py
|
||||
|
||||
|
||||
def parse_dag(text):
|
||||
pattern = r"<json>(.*?)</json>"
|
||||
match = re.search(pattern, text, re.DOTALL)
|
||||
if match:
|
||||
json_str = match.group(1)
|
||||
try:
|
||||
json_data = json.loads(json_str)
|
||||
return Dag(**json_data["dag"])
|
||||
except json.JSONDecodeError:
|
||||
print("Error: Invalid JSON")
|
||||
return None
|
||||
except KeyError:
|
||||
print("Error: 'dag' key not found in JSON")
|
||||
return None
|
||||
except ValidationError as e:
|
||||
print(f"Error: Invalid data structure - {e}")
|
||||
return None
|
||||
else:
|
||||
print("Error: JSON not found")
|
||||
return None
|
||||
|
||||
|
||||
def parse_dag(text):
|
||||
"""
|
||||
Try extracting JSON from <json>…</json> tags first;
|
||||
if not found, try ```json … ``` Markdown fences.
|
||||
"""
|
||||
|
||||
def _extract(pattern):
|
||||
m = re.search(pattern, text, re.DOTALL)
|
||||
return m.group(1).strip() if m else None
|
||||
|
||||
# 1) look for <json>…</json>
|
||||
json_str = _extract(r"<json>(.*?)</json>")
|
||||
# 2) fallback to ```json … ```
|
||||
if json_str is None:
|
||||
json_str = _extract(r"```json\s*(.*?)\s*```")
|
||||
|
||||
if json_str is None:
|
||||
print("Error: JSON not found in either <json> tags or ```json``` fence")
|
||||
return None
|
||||
|
||||
try:
|
||||
payload = json.loads(json_str)
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"Error: Invalid JSON ({e})")
|
||||
return None
|
||||
|
||||
if "dag" not in payload:
|
||||
print("Error: 'dag' key not found in JSON")
|
||||
return None
|
||||
|
||||
try:
|
||||
return Dag(**payload["dag"])
|
||||
except ValidationError as e:
|
||||
print(f"Error: Invalid data structure - {e}")
|
||||
return None
|
||||
|
||||
|
||||
def parse_single_code_from_string(input_string):
|
||||
input_string = input_string.strip()
|
||||
if input_string.strip() in ["WAIT", "DONE", "FAIL"]:
|
||||
return input_string.strip()
|
||||
|
||||
# This regular expression will match both ```code``` and ```python code```
|
||||
# and capture the `code` part. It uses a non-greedy match for the content inside.
|
||||
pattern = r"```(?:\w+\s+)?(.*?)```"
|
||||
# Find all non-overlapping matches in the string
|
||||
matches = re.findall(pattern, input_string, re.DOTALL)
|
||||
|
||||
# The regex above captures the content inside the triple backticks.
|
||||
# The `re.DOTALL` flag allows the dot `.` to match newline characters as well,
|
||||
# so the code inside backticks can span multiple lines.
|
||||
|
||||
# matches now contains all the captured code snippets
|
||||
|
||||
codes = []
|
||||
|
||||
for match in matches:
|
||||
match = match.strip()
|
||||
commands = [
|
||||
"WAIT",
|
||||
"DONE",
|
||||
"FAIL",
|
||||
] # fixme: updates this part when we have more commands
|
||||
|
||||
if match in commands:
|
||||
codes.append(match.strip())
|
||||
elif match.split("\n")[-1] in commands:
|
||||
if len(match.split("\n")) > 1:
|
||||
codes.append("\n".join(match.split("\n")[:-1]))
|
||||
codes.append(match.split("\n")[-1])
|
||||
else:
|
||||
codes.append(match)
|
||||
|
||||
if len(codes) <= 0:
|
||||
return "fail"
|
||||
return codes[0]
|
||||
|
||||
|
||||
def get_input_token_length(input_string):
|
||||
enc = tiktoken.encoding_for_model("gpt-4")
|
||||
tokens = enc.encode(input_string)
|
||||
return len(tokens)
|
||||
|
||||
|
||||
def sanitize_code(code):
|
||||
# This pattern captures the outermost double-quoted text
|
||||
if "\n" in code:
|
||||
pattern = r'(".*?")'
|
||||
# Find all matches in the text
|
||||
matches = re.findall(pattern, code, flags=re.DOTALL)
|
||||
if matches:
|
||||
# Replace the first occurrence only
|
||||
first_match = matches[0]
|
||||
code = code.replace(first_match, f'"""{first_match[1:-1]}"""', 1)
|
||||
return code
|
||||
|
||||
|
||||
def extract_first_agent_function(code_string):
|
||||
# Regular expression pattern to match 'agent' functions with any arguments, including nested parentheses
|
||||
pattern = r'agent\.[a-zA-Z_]+\((?:[^()\'"]|\'[^\']*\'|"[^"]*")*\)'
|
||||
|
||||
# Find all matches in the string
|
||||
matches = re.findall(pattern, code_string)
|
||||
|
||||
# Return the first match if found, otherwise return None
|
||||
return matches[0] if matches else None
|
||||
|
||||
|
||||
def load_knowledge_base(kb_path: str) -> Dict:
|
||||
try:
|
||||
with open(kb_path, "r") as f:
|
||||
return json.load(f)
|
||||
except Exception as e:
|
||||
print(f"Error loading knowledge base: {e}")
|
||||
return {}
|
||||
|
||||
|
||||
def load_embeddings(embeddings_path: str) -> Dict:
|
||||
try:
|
||||
with open(embeddings_path, "rb") as f:
|
||||
return pickle.load(f)
|
||||
except Exception as e:
|
||||
print(f"Error loading embeddings: {e}")
|
||||
return {}
|
||||
|
||||
|
||||
def save_embeddings(embeddings_path: str, embeddings: Dict):
|
||||
try:
|
||||
with open(embeddings_path, "wb") as f:
|
||||
pickle.dump(embeddings, f)
|
||||
except Exception as e:
|
||||
print(f"Error saving embeddings: {e}")
|
||||
Reference in New Issue
Block a user