Adding AgentLogger

.github/workflows/test.yml

@@ -2,7 +2,7 @@ name: Tests
 
 on:
   push:
-    branches: ["main"]
+    branches: [ "main" ]
   pull_request:
   workflow_dispatch:
 

.gitignore

@@ -207,3 +207,9 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+
+
+#
+autobolt/
+*.db-*
+*.db

README.md

@@ -2,17 +2,21 @@
 
 # AutoBoltAgent
 
-AutoBoltAgent is a small Python package that provides AI-powered assistance for bolt design tasks. It aims to help engineers prototype bolt specifications and explore design options programmatically.
+AutoBoltAgent is a small Python package that provides AI-powered assistance for bolt design tasks. It aims to help
+engineers prototype bolt specifications and explore design options programmatically.
 
 Key points:
+
 - Lightweight library packaged under `autoboltagent`.
 - Includes example tools for low- and high-fidelity bolt generation and helper inputs.
 
 ## Requirements
+
 - Python 3.10+
 - (Optional) Conda for easy environment setup — an `environment.yml` is included in the repository.
 
 ## Quick start
+
 1. Create and activate a Python environment (optional, using conda):
 
 ```bash

environment.yml

@@ -17,4 +17,5 @@ dependencies:
       - pandas
       - "autobolt @ git+https://github.com/sriyanc2001/AutoBolt.git"
       - black
-      - pytest
+      - pytest
+      - sqlalchemy

examples/simple.py

@@ -1,6 +1,7 @@
-import autoboltagent
 import smolagents
 
+import autoboltagent
+
 # Use a local model on macOS for faster testing
 model = smolagents.MLXModel(
     model_id="mlx-community/Qwen3-4B-Instruct-2507-4bit",

pyproject.toml

@@ -8,8 +8,8 @@ version = "0.1.0"
 description = "AI agent for automatically designing bolts"
 readme = "README.md"
 requires-python = ">=3.10"
-license = {text = "MIT"}
-authors = [{name="Sriya"}, {name="Chris McComb", email="ccmcc2012@gmail.com"}]
+license = { text = "MIT" }
+authors = [{ name = "Sriya" }, { name = "Chris McComb", email = "ccmcc2012@gmail.com" }]
 dependencies = [
     "autobolt @ git+https://github.com/sriyanc2001/AutoBolt.git",
     "smolagents[transformers]",
@@ -18,7 +18,7 @@ dependencies = [
 ]
 [project.optional-dependencies]
 test = [
-  "pytest",
+    "pytest",
 ]
 
 [tool.pytest.ini_options]

src/autoboltagent/agents.py

@@ -6,17 +6,18 @@
     DUAL_FIDELITY_COORDINATION,
 )
 from .tools import AnalyticalTool, FiniteElementTool
+from .tools.logger import AgentLogger
 
 
-class GuessingAgent(smolagents.ToolCallingAgent):
+class GuessingAgent(smolagents.agents.ToolCallingAgent):
     """
     An agent that makes guesses without using any tools.
 
     This agent operates solely based on its internal model and does not utilize any external tools for analysis or calculations.
     It is designed to provide initial estimates or solutions based on its knowledge and reasoning capabilities.
     """
 
-    def __init__(self, model: smolagents.Model) -> None:
+    def __init__(self, model: smolagents.models.Model) -> None:
         """
         Initializes a GuessingAgent that does not use any tools.
 
@@ -33,40 +34,67 @@ def __init__(self, model: smolagents.Model) -> None:
         )
 
 
-class LowFidelityAgent(smolagents.ToolCallingAgent):
+class LowFidelityAgent(smolagents.agents.ToolCallingAgent):
     """
     An agent that utilizes a low-fidelity analytical tool for bolted connection design.
 
     This agent leverages an analytical tool to perform calculations and analyses related to bolted connections.
     It is designed to provide solutions based on simplified models and assumptions, making it suitable for quick estimates and preliminary designs.
     """
 
-    def __init__(self, model: smolagents.Model) -> None:
+    def __init__(
+        self,
+        model: smolagents.models.Model,
+        agent_id: str,
+        run_id: str,
+        target_fos: float,
+        agent_logger: AgentLogger | None = None,
+        max_steps=20,
+    ) -> None:
         """
         Initializes a LowFidelityAgent that uses an analytical tool.
 
         Args:
             model: An instance of smolagents.Model to be used by the agent.
         """
+
+        self.agent_logger = agent_logger
+        self.agent_id = agent_id
+        self.run_id = run_id
+        self.target_fos = target_fos
+
+        callbacks = [self.log] if self.agent_logger else []
+
         super().__init__(
             name="LowFidelityAgent",
             tools=[AnalyticalTool()],
             add_base_tools=False,
             model=model,
             instructions=BASE_INSTRUCTIONS + TOOL_USING_INSTRUCTION,
+            step_callbacks=callbacks,
             verbosity_level=2,
+            max_steps=max_steps,
         )
 
+    def log(self, step, agent):
+        if self.agent_logger and step.__class__.__name__ == "ActionStep":
+            self.agent_logger.log(
+                agent_id=self.agent_id,
+                run_id=self.run_id,
+                target_fos=self.target_fos,
+                action_step=step,
+            )
+
 
-class HighFidelityAgent(smolagents.ToolCallingAgent):
+class HighFidelityAgent(smolagents.agents.ToolCallingAgent):
     """
     An agent that utilizes a high-fidelity finite element analysis tool for bolted connection design.
 
     This agent leverages a finite element tool to perform detailed calculations and analyses related to bolted connections.
     It is designed to provide accurate and reliable solutions based on comprehensive models, making it suitable for
     """
 
-    def __init__(self, model: smolagents.Model) -> None:
+    def __init__(self, model: smolagents.models.Model) -> None:
         """
         Initializes a HighFidelityAgent that uses a finite element tool.
 
@@ -83,15 +111,15 @@ def __init__(self, model: smolagents.Model) -> None:
         )
 
 
-class DualFidelityAgent(smolagents.ToolCallingAgent):
+class DualFidelityAgent(smolagents.agents.ToolCallingAgent):
     """
     An agent that utilizes both low-fidelity and high-fidelity tools for bolted connection design.
 
     This agent leverages both an analytical tool and a finite element tool to perform calculations and analyses related to bolted connections.
     It is designed to provide solutions that balance speed and accuracy by using the low-fidelity tool
     """
 
-    def __init__(self, model: smolagents.Model) -> None:
+    def __init__(self, model: smolagents.models.Model) -> None:
         """
         Initializes a DualFidelityAgent that uses both analytical and finite element tools.
 

src/autoboltagent/tools/high_fidelity_tool.py

@@ -1,10 +1,12 @@
+from typing import Union, cast
+
 import autobolt
 import smolagents
 
 from .inputs import INPUTS
 
 
-class FiniteElementTool(smolagents.Tool):
+class FiniteElementTool(smolagents.tools.Tool):
     """
     A tool that calculates the factor of safety for a bolted connection using finite element analysis.
 
@@ -15,7 +17,8 @@ class FiniteElementTool(smolagents.Tool):
     name = "fea_fos_calculation"
     description = "Calculates the factor of safety using finite element analysis."
 
-    inputs = INPUTS
+    input_type = dict[str, dict[str, Union[str, type, bool]]]
+    inputs: input_type = cast(input_type, INPUTS)
 
     output_type = "number"
 

src/autoboltagent/tools/logger.py

@@ -0,0 +1,127 @@
+from datetime import datetime, timezone
+from pathlib import Path
+
+from sqlalchemy import create_engine
+from sqlalchemy.orm import (
+    declarative_base,
+    sessionmaker,
+    Mapped,
+    mapped_column,
+)
+
+Base = declarative_base()
+
+
+class Iteration(Base):
+    __tablename__ = "iterations"
+
+    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
+    agent_id: Mapped[str] = mapped_column()
+    run_id: Mapped[str] = mapped_column()
+    iteration_no: Mapped[int] = mapped_column()
+
+    start_time: Mapped[datetime] = mapped_column(nullable=True)
+    end_time: Mapped[datetime] = mapped_column(nullable=True)
+
+    status: Mapped[str] = mapped_column(nullable=True)
+    tool_call: Mapped[str] = mapped_column(nullable=True)
+    observations: Mapped[str] = mapped_column(nullable=True)
+    target_fos: Mapped[float] = mapped_column(nullable=True)
+
+    failure_reason: Mapped[str] = mapped_column(nullable=True)
+    llm_output: Mapped[str] = mapped_column(nullable=True)
+    error_message: Mapped[str] = mapped_column(nullable=True)
+
+
+class AgentLogger:
+    _instance = None
+
+    def connect_to_db(self, db_url: str):
+        self.db_url = db_url
+        self.engine = create_engine(db_url, future=True, pool_pre_ping=True)
+
+        try:
+            with self.engine.connect() as conn:
+                conn.exec_driver_sql("PRAGMA journal_mode=WAL;")
+                conn.exec_driver_sql("PRAGMA synchronous=NORMAL;")
+
+            Base.metadata.create_all(self.engine)
+
+            self.db_session = sessionmaker(
+                bind=self.engine, expire_on_commit=False, future=True
+            )
+        except Exception as e:
+            raise IOError("Failed to connect to DB, check if file in use", repr(e))
+
+    def __new__(cls, db_url: str):
+        if not cls._instance:
+            cls._instance = super().__new__(cls)
+            cls._instance.db_url = None
+            cls._instance.engine = None
+            cls._instance.db_session = None
+            cls._instance.connect_to_db(db_url)
+        return cls._instance
+
+    @classmethod
+    def reset(cls):
+        if not cls._instance:
+            return
+
+        inst = cls._instance
+
+        if inst.engine:
+            inst.engine.dispose()
+
+        if inst.db_url:
+            for suffix in ("", "-wal", "-shm"):
+                file_path = Path(inst.db_url.replace("sqlite:///", "") + suffix)
+                file_path.unlink(missing_ok=True)
+
+        cls._instance = None
+
+    def log(self, run_id, agent_id, target_fos, action_step):
+
+        iteration_no = action_step.step_number
+        start_dt = datetime.fromtimestamp(
+            action_step.timing.start_time, tz=timezone.utc
+        )
+        end_dt = datetime.fromtimestamp(action_step.timing.end_time, tz=timezone.utc)
+
+        error = getattr(action_step, "error", None)
+        tool_calls = getattr(action_step, "tool_calls", None)
+        observations = getattr(action_step, "observations", None)
+        llm_message = getattr(action_step, "model_output_message", None)
+        llm_output = getattr(llm_message, "content", None)
+
+        print(tool_calls)
+        print(error)
+        print(observations)
+        print(llm_message)
+        print(action_step.token_usage)
+
+        try:
+
+            with self.db_session() as session:
+                session.add(
+                    Iteration(
+                        run_id=run_id,
+                        agent_id=agent_id,
+                        iteration_no=iteration_no,
+                        start_time=start_dt,
+                        end_time=end_dt,
+                        tool_call=str(tool_calls[0] if tool_calls else None),
+                        observations=observations,
+                        target_fos=target_fos,
+                        llm_output=llm_output,
+                        error_message=(
+                            error.message if (error and error.message) else None
+                        ),
+                    )
+                )
+                session.flush()
+                session.commit()
+
+        except Exception as e:
+            print("\n\n")
+            print(e)
+            print("\n\n")

tests/test_agents.py

@@ -10,15 +10,15 @@ def is_macos() -> bool:
     return platform.system() == "Darwin"
 
 
-def get_testing_model() -> smolagents.Model:
+def get_testing_model() -> smolagents.models.Model:
     if is_macos():
         # Use a local model on macOS for faster testing
         return smolagents.MLXModel(
             model_id="Qwen/Qwen3-1.7B-MLX-4bit",
         )
     else:
         # Use the smallest Instruct model available for fast CI feedback
-        return smolagents.TransformersModel(
+        return smolagents.models.TransformersModel(
             model_id="HuggingFaceTB/SmolLM-135M-Instruct",
             max_new_tokens=200,  # Keep generation short for speed
         )
@@ -38,10 +38,16 @@ def test_guessing_agent():
 def test_low_fidelity_agent():
 
     # Create the LowFidelityAgent and run it
-    response = autoboltagent.LowFidelityAgent(get_testing_model()).run(
-        autoboltagent.prompts.EXAMPLE_TASK_INSTRUCTIONS
+    agent = autoboltagent.LowFidelityAgent(
+        model=get_testing_model(),
+        agent_id="low fidelity agent",
+        run_id="test 1",
+        target_fos=3.0,
+        max_steps=5,
     )
 
+    response = agent.run(autoboltagent.prompts.EXAMPLE_TASK_INSTRUCTIONS)
+
     # Make sure the response exists
     assert response is not None