demeter

  1import json
  2import ast
  3from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage
  4from agent.sub_agents.water_and_atmospheric_dependencies.physics_engine import predict_outcome
  5from agent.sub_agents.water_and_atmospheric_dependencies.retrieval import ask_historian, ask_rag, diagnose_plant, ask_memory
  6from agent.sub_agents.water_and_atmospheric_dependencies.tools import calculate_vpd, web_search, check_ph_safety
  7from agent.guardrails.validation import validate_bounds, HARD_BOUNDS
  8
  9# 🟢 Add diagnose_plant and ask_memory to the map
 10TOOL_MAP = {
 11    "ask_historian": ask_historian,
 12    "ask_rag": ask_rag,
 13    "web_search": web_search,
 14    "calculate_vpd": calculate_vpd,
 15    "diagnose_plant": diagnose_plant,
 16    "ask_memory": ask_memory,
 17    "check_ph_safety": check_ph_safety
 18}
 19
 20def validate_plan_constraints(draft):
 21    """
 22    Check if draft plan violates hard constraints.
 23    Returns (is_valid, violation_message)
 24    """
 25    if not isinstance(draft, dict) or not draft:
 26        return False, "Plan must be a non-empty JSON object"
 27    
 28    violations = []
 29    for param, value in draft.items():
 30        if param in HARD_BOUNDS:
 31            # Skip None values (no-ops)
 32            if value is None or (isinstance(value, (int, float)) and value == 0):
 33                continue
 34            
 35            is_valid, message = validate_bounds(param, value)
 36            if not is_valid:
 37                violations.append(message)
 38    
 39    if violations:
 40        return False, "\n".join(violations)
 41    return True, ""
 42
 43def decide_node(state, model, system_prompt):
 44    """
 45    Node 1: Drafts a plan OR calls a tool.
 46    """
 47    retry_num = state['retry_count'] + 1
 48    print(f"   🤔 Thinking (Attempt {retry_num}/3)...")
 49    
 50    messages = state.get("messages", [])
 51    if not messages:
 52        messages = [SystemMessage(content=system_prompt)]
 53        user_msg = (
 54            f"Current Sensors: {state['sensors']}\n"
 55            f"Strategy: {state['strategy']}\n"
 56            f"Research: {state['research_context']}\n"
 57            f"History Context: {state.get('history', 'None provided')}\n"
 58        )
 59        if state.get("critique"):
 60            user_msg += f"\n\n❌ PREVIOUS SIMULATION FAILED: {state['critique']}"
 61            
 62        messages.append(HumanMessage(content=user_msg))
 63    else:
 64        # On retry, check if there's new critique to add
 65        if state.get("critique"):
 66            # Add critique as a new feedback message
 67            critique_msg = HumanMessage(content=f"❌ PREVIOUS ATTEMPT FAILED: {state['critique']}\n\nPlease try a different plan.")
 68            messages = messages + [critique_msg]
 69
 70    response = model.invoke(messages)
 71    new_messages = messages + [response]
 72    
 73    if response.tool_calls:
 74        print(f"   📞 Calling Tool: {response.tool_calls[0]['name']}")
 75        return {
 76            "messages": new_messages, 
 77            "next_step": "tools"
 78        }
 79
 80    content = response.content.replace("```json", "").replace("```", "").strip()
 81    
 82    try:
 83        # 1. Try standard JSON parsing first
 84        draft = json.loads(content)
 85    except json.JSONDecodeError:
 86        try:
 87            # 2. Fallback: Python literal eval (Handles single quotes)
 88            draft = ast.literal_eval(content)
 89        except Exception as e:
 90            draft = {}
 91        
 92    return {
 93        "draft_plan": draft, 
 94        "messages": new_messages,
 95        "next_step": "simulate",
 96        "retry_count": retry_num
 97    }
 98
 99def execute_tools_node(state):
100    """
101    Executes the tool call and returns the result to the LLM.
102    Handles 'Hidden State Injection' for heavy data like images.
103    """
104    print("   ⚙️ Executing Tools...")
105    
106    if "messages" not in state or not state["messages"]:
107        raise ValueError("No messages found in state.")
108
109    last_message = state["messages"][-1]
110    tool_results = []
111    
112    for tool_call in last_message.tool_calls:
113        tool_name = tool_call["name"]
114        tool_args = tool_call["args"].copy()
115        
116        # 🟢 INJECTION LOGIC: Pass image_b64 from state to the COPY
117        if tool_name == "diagnose_plant":
118            tool_args["image_b64"] = state.get("image_b64")
119
120        if tool_name in TOOL_MAP:
121            try:
122                output = TOOL_MAP[tool_name].invoke(tool_args)
123                result_content = str(output)
124            except Exception as e:
125                result_content = f"Error executing {tool_name}: {e}"
126        else:
127            result_content = f"Error: Tool {tool_name} is not available."
128            
129        print(f"      -> {tool_name}: {result_content[:100]}...") # Truncated log
130
131        tool_results.append(ToolMessage(
132            tool_call_id=tool_call["id"],
133            name=tool_name,
134            content=result_content
135        ))
136
137    return {"messages": state["messages"] + tool_results}
138
139def simulate_node(state):
140    print("   🧪 Simulating Outcome...")
141    draft = state.get('draft_plan')
142    
143    if not draft:
144        return {"simulation_result": {"passed": True, "reason": "No valid JSON plan generated."}}
145    
146    # 🛡️ CHECK CONSTRAINTS BEFORE SIMULATION (avoid wasting LLM calls)
147    is_valid, violation_msg = validate_plan_constraints(draft)
148    if not is_valid:
149        critique = f"Hard constraint violation:\n{violation_msg}"
150        return {
151            "simulation_result": {
152                "passed": False,
153                "reason": critique
154            },
155            "critique": critique  # Pass back to LLM for next retry
156        }
157
158    current = state['sensors']
159    
160    # Ensure physics engine is imported correctly at top
161    prediction = predict_outcome(current, draft)
162    
163    health = prediction.get('predicted_health', 0)
164    risk = prediction.get('risk_warning', "None")
165    
166    result = {"passed": True, "reason": ""}
167    
168    if health < 70.0:
169        reason = f"Predicted Health drops to {health}%. Warning: {risk}"
170        result["reason"] = reason
171        result["passed"] = False
172        return {
173            "simulation_result": result,
174            "critique": reason  # Pass back to LLM for next retry
175        }
176    else:
177        result["passed"] = True
178        
179    return {"simulation_result": result}
180
181def finalize_node(state):
182    plan = state['draft_plan']
183    print("   ✅ Plan Approved.")
184    reason = state.get("simulation_result", {}).get("reason", "")
185    if reason:
186        print(f"      Reason: {reason}")
187    return {"final_action": plan}