Spaces:
Sleeping
Sleeping
File size: 4,885 Bytes
a5bca5e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | """
==============================================================================
Universal-Node-Resolver — Local Test & Baseline
==============================================================================
Runs the NodeResolverAgent locally using a dummy heuristic function to
validate the reward shaping and environment logic. Generates a plot
for the project README.
"""
import json
import logging
import os
import random
import re
import sys
import matplotlib.pyplot as plt
# Ensure imports work from project root
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from client.agent import NodeResolverAgent
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("test_local")
def random_heuristic_agent(prompt: str) -> str:
"""
Dummy LLM inference function.
Parses the prompt to find a package involved in a conflict, and
randomly decides to 'update' or 'delete' it.
"""
packages = []
# Attempt to extract package names from the Active Conflicts log
try:
if "### Active Conflicts" in prompt:
conflicts_section = prompt.split("### Active Conflicts")[1].split("## Your Task")[0]
for line in conflicts_section.split("\n"):
if "ERESOLVE:" in line:
# Match packages in "pkg-name@version"
matches = re.findall(r"([a-zA-Z0-9-]+)@", line)
packages.extend(matches)
except Exception:
pass
# Fallback to installed packages if no errors parsed
if not packages:
try:
json_match = re.search(r"```json\s+(.*?)\s+```", prompt, re.DOTALL)
if json_match:
state = json.loads(json_match.group(1))
deps = state.get("dependencies", {})
packages = list(deps.keys())
except Exception:
pass
if not packages:
packages = ["pkg-001"] # Ultimate fallback
# Remove duplicates
packages = list(set(packages))
# Select random target package
target_pkg = random.choice(packages)
# 80% update, 20% delete (delete is highly likely to trigger nuke penalty)
if random.random() < 0.8:
action_type = "update"
# Guess a semantic version
version = f"{random.randint(0, 3)}.{random.randint(0, 5)}.0"
else:
action_type = "delete"
version = None
# Return formatted JSON string as requested
action = {
"action_type": action_type,
"package_name": target_pkg,
"version_target": version
}
return json.dumps(action)
def main():
logger.info("Initializing NodeResolverAgent for local sanity check...")
agent = NodeResolverAgent()
num_episodes = 50
rewards = []
steps_taken = []
logger.info(f"Running {num_episodes} baseline episodes with random heuristic agent...")
for i in range(num_episodes):
# We suppress verbose logging per step to keep output clean,
# but print episode summaries
reward, steps, solved = agent.run_episode(
llm_inference_function=random_heuristic_agent,
verbose=False
)
rewards.append(reward)
steps_taken.append(steps)
logger.info(f"Episode {i+1:02d}: Reward = {reward:6.1f} | Steps = {steps:2d} | Solved = {solved}")
# Plotting
logger.info("Generating baseline performance plots...")
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))
# Reward Plot
ax1.plot(rewards, color='#4CAF50', marker='o', linestyle='-', markersize=4, alpha=0.8)
ax1.axhline(0, color='gray', linestyle='--', alpha=0.5)
ax1.set_title("Dummy Agent: Episode vs Total Reward", fontsize=12, fontweight='bold')
ax1.set_xlabel("Episode")
ax1.set_ylabel("Total Reward")
ax1.grid(True, linestyle=':', alpha=0.7)
# Steps Plot
ax2.plot(steps_taken, color='#2196F3', marker='o', linestyle='-', markersize=4, alpha=0.8)
ax2.set_title("Dummy Agent: Episode vs Steps Taken", fontsize=12, fontweight='bold')
ax2.set_xlabel("Episode")
ax2.set_ylabel("Steps Taken")
ax2.grid(True, linestyle=':', alpha=0.7)
plt.tight_layout()
# Ensure assets directory exists
assets_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "assets"))
os.makedirs(assets_dir, exist_ok=True)
plot_path = os.path.join(assets_dir, "baseline_rewards.png")
plt.savefig(plot_path, dpi=150, bbox_inches='tight')
logger.info(f"Successfully saved plot to: {plot_path}")
logger.info("Local sanity check complete!")
if __name__ == "__main__":
main()
|