Spaces:
Running
Running
File size: 6,664 Bytes
7eb1167 | 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 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | from typing import TypedDict, List, Optional
from config import Config
from utils import LLMService, GraphState
from langgraph.graph import StateGraph, START, END
from langgraph.types import Command
def llm_requires_custom_mesh(state: GraphState) -> int:
"""
Use LLM to determine if user requires custom mesh based on their requirement.
Args:
state: Current graph state containing user requirement and LLM service
Returns:
int: 1 if custom mesh is required, 2 if gmsh mesh is required, 0 otherwise
"""
user_requirement = state["user_requirement"]
system_prompt = (
"You are an expert in OpenFOAM workflow analysis. "
"Analyze the user requirement to determine if they want to use a custom mesh file. "
"Look for keywords like: custom mesh, mesh file, .msh, .stl, .obj, gmsh, snappyHexMesh, "
"or any mention of importing/using external mesh files. "
"If the user explicitly mentions or implies they want to use a custom mesh file, return 'custom_mesh'. "
"If they want to use standard OpenFOAM mesh generation (blockMesh, snappyHexMesh with STL, etc.), return 'standard_mesh'. "
"Look for keywords like gmsh and determine if they want to create mesh using gmsh. If they want to create mesh using gmsh, return 'gmsh_mesh'. "
"Be conservative - if unsure, assume 'standard_mesh' unless clearly specified otherwise."
"Only return 'custom_mesh' or 'standard_mesh' or 'gmsh_mesh'. Don't return anything else."
)
user_prompt = (
f"User requirement: {user_requirement}\n\n"
"Determine if the user wants to use a custom mesh file. "
"Return exactly 'custom_mesh' if they want to use a custom mesh file, "
"'standard_mesh' if they want standard OpenFOAM mesh generation or 'gmsh_mesh' if they want to create mesh using gmsh."
)
response = state["llm_service"].invoke(user_prompt, system_prompt)
if "custom_mesh" in response.lower():
return 1
elif "gmsh_mesh" in response.lower():
return 2
else:
return 0
def llm_requires_hpc(state: GraphState) -> bool:
"""
Use LLM to determine if user requires HPC/cluster execution based on their requirement.
Args:
state: Current graph state containing user requirement and LLM service
Returns:
bool: True if HPC execution is required, False otherwise
"""
user_requirement = state["user_requirement"]
system_prompt = (
"You are an expert in OpenFOAM workflow analysis. "
"Analyze the user requirement to determine if they want to run the simulation on HPC (High Performance Computing) or locally. "
"Look for keywords like: HPC, cluster, supercomputer, SLURM, PBS, job queue, "
"parallel computing, distributed computing, or any mention of running on remote systems. "
"If the user explicitly mentions or implies they want to run on HPC/cluster, return 'hpc_run'. "
"If they want to run locally or don't specify, return 'local_run'. "
"Be conservative - if unsure, assume local run unless clearly specified otherwise."
"Only return 'hpc_run' or 'local_run'. Don't return anything else."
)
user_prompt = (
f"User requirement: {user_requirement}\n\n"
"return 'hpc_run' or 'local_run'"
)
response = state["llm_service"].invoke(user_prompt, system_prompt)
return "hpc_run" in response.lower()
def llm_requires_visualization(state: GraphState) -> bool:
"""
Use LLM to determine if user requires visualization based on their requirement.
Args:
state: Current graph state containing user requirement and LLM service
Returns:
bool: True if visualization is required, False otherwise
"""
user_requirement = state["user_requirement"]
system_prompt = (
"You are an expert in OpenFOAM workflow analysis. "
"Analyze the user requirement to determine if they want visualization of results. "
"Look for keywords like: plot, visualize, graph, chart, contour, streamlines, paraview, post-processing."
"Only if the user explicitly mentions they want visualization, return 'yes_visualization'. "
"If they don't mention visualization or only want to run the simulation, return 'no_visualization'. "
"Be conservative - if unsure, assume visualization is wanted unless clearly specified otherwise."
"Only return 'yes_visualization' or 'no_visualization'. Don't return anything else."
)
user_prompt = (
f"User requirement: {user_requirement}\n\n"
"return 'yes_visualization' or 'no_visualization'"
)
response = state["llm_service"].invoke(user_prompt, system_prompt)
return "yes_visualization" in response.lower()
def route_after_architect(state: GraphState):
"""
Route after architect node based on whether user wants custom mesh.
For current version, if user wants custom mesh, user should be able to provide a path to the mesh file.
"""
mesh_type = state.get("mesh_type", "standard_mesh")
if mesh_type == "custom_mesh":
print("Router: Custom mesh requested. Routing to meshing node.")
return "meshing"
elif mesh_type == "gmsh_mesh":
print("Router: GMSH mesh requested. Routing to meshing node.")
return "meshing"
else:
print("Router: Standard mesh generation. Routing to input_writer node.")
return "input_writer"
def route_after_input_writer(state: GraphState):
"""
Route after input_writer node based on whether user wants to run on HPC.
"""
if llm_requires_hpc(state):
print("LLM determined: HPC run requested. Routing to hpc_runner node.")
return "hpc_runner"
else:
print("LLM determined: Local run requested. Routing to local_runner node.")
return "local_runner"
def route_after_runner(state: GraphState):
if state.get("error_logs") and len(state["error_logs"]) > 0:
return "reviewer"
elif llm_requires_visualization(state):
return "visualization"
else:
return END
def route_after_reviewer(state: GraphState):
loop_count = state.get("loop_count", 0)
max_loop = state["config"].max_loop
if loop_count >= max_loop:
print(f"Maximum loop count ({max_loop}) reached. Ending workflow.")
if llm_requires_visualization(state):
return "visualization"
else:
return END
print(f"Loop {loop_count}: Continuing to fix errors.")
return "input_writer"
|