Updated beam agent.

agents/agent_with_custom_beam_design_tools.py

@@ -1,14 +1,14 @@
 import os
+import math
 
 import smolagents
 import smolagents.tools
 
+
 # Rectangular section properties tool
 class RectSectionPropsTool(smolagents.tools.Tool):
     name = "rect_section_props"
-    description = (
-        "Rectangular section properties."
-    )
+    description = "Rectangular section properties."
     inputs = {
         "b": {"type": "number", "description": "Width in meters"},
         "h": {"type": "number", "description": "Height in meters"},
@@ -20,15 +20,14 @@ class RectSectionPropsTool(smolagents.tools.Tool):
             raise ValueError("b and h must be positive.")
         A = b * h
         I = b * (h**3) / 12.0
-        S = I / (h / 2.0)       # = b*h^2/6
+        S = I / (h / 2.0)  # = b*h^2/6
         return {"A": A, "I": I, "S": S}
 
+
 # Circle section properties tool
 class CircleSectionPropsTool(smolagents.tools.Tool):
     name = "circle_section_props"
-    description = (
-        "Circular section properties."
-    )
+    description = "Circular section properties."
     inputs = {
         "d": {"type": "number", "description": "Diameter in meters"},
     }
@@ -39,20 +38,21 @@ class CircleSectionPropsTool(smolagents.tools.Tool):
             raise ValueError("d must be positive.")
         A = math.pi * (d**2) / 4.0
         I = math.pi * (d**4) / 64.0
-        S = I / (d / 2.0)       # = π d^3 / 32
+        S = I / (d / 2.0)  # = π d^3 / 32
         return {"A": A, "I": I, "S": S}
 
 
 # Mass tool
 class BeamMassTool(smolagents.tools.Tool):
     name = "beam_mass"
-    description = (
-        "Beam mass [kg]."
-    )
+    description = "Beam mass [kg]."
     inputs = {
         "L": {"type": "number", "description": "Length in meters"},
         "rho": {"type": "number", "description": "Material density in kg/m^3"},
-        "A": {"type": "number", "description": "Cross-sectional area in m^2 (optional)"},
+        "A": {
+            "type": "number",
+            "description": "Cross-sectional area in m^2 (optional)",
+        },
     }
     output_type = "number"
 
@@ -61,40 +61,29 @@ class BeamMassTool(smolagents.tools.Tool):
             raise ValueError("L, rho, and A must be positive.")
         return rho * A * L
 
+
 # Deflection tool
 class BeamDeflectionTool(smolagents.tools.Tool):
     name = "beam_deflection"
-    description = (
-        "Max deflection δ [m] for simple cases."
-    )
+    description = "Max deflection δ [m] for simple cases."
     inputs = {
         "P": {"type": "number", "description": "Point load in Newtons"},
         "L": {"type": "number", "description": "Span length in meters"},
         "E": {"type": "number", "description": "Young's modulus in Pascals"},
         "I": {"type": "number", "description": "Second moment in m^4 (optional)"},
-        "bc": {"type": "string", "description": "Boundary case: 'cantilever_end' or 'ss_center'"},
     }
     output_type = "number"
 
-    def forward(self, P: float, L: float, E: float, bc: str, I: float) -> float:
+    def forward(self, P: float, L: float, E: float, I: float) -> float:
         if any(x <= 0 for x in [P, L, E, I]):
             raise ValueError("P, L, E, I must be positive.")
-        key = str(bc).lower().strip()
-        if key in ["cantilever", "cantilever_end", "cantilver_end", "cantilever-tip", "cantilever_tip"]:
-            return P * (L**3) / (3.0 * E * I)
-        elif key in ["ss", "ss_center", "simply_supported_center", "simply supported center", "simply_supported"]:
-            return P * (L**3) / (48.0 * E * I)
-        else:
-            raise ValueError("Unsupported bc. Use 'cantilever_end' or 'ss_center'.")
-
+        return P * (L**3) / (3.0 * E * I)
 
 
 # Bending stress tool
 class BeamBendingStressTool(smolagents.tools.Tool):
     name = "beam_bending_stress"
-    description = (
-        "Max bending stress σ_max [Pa]."
-    )
+    description = "Max bending stress σ_max [Pa]."
     inputs = {
         "M": {"type": "number", "description": "Bending moment in N·m"},
         "I": {"type": "number", "description": "Second moment in m^4 (optional)"},
@@ -105,7 +94,8 @@ class BeamBendingStressTool(smolagents.tools.Tool):
     def forward(self, M: float, I: float, c: float) -> float:
         if M <= 0:
             raise ValueError("M must be positive.")
-        return M *c/I
+        return M * c / I
+
 
 def generate_beam_agent() -> smolagents.CodeAgent:
     # Define the agent with all of these tools.
@@ -122,7 +112,8 @@ def generate_beam_agent() -> smolagents.CodeAgent:
             api_base=os.getenv("UPSTREAM_OPENAI_BASE", "").rstrip("/"),
             api_key=os.getenv("OPENAI_API_KEY"),
         ),
+        instructions="You are a beam design agent with custom tools to support that task. Use those tools whenever feasible.",
         name="BeamAgent",
         add_base_tools=False,
         max_steps=12,
-)
+    )