# Predictive Maintenance Agent Development Guide

This document provides comprehensive guidance for developing with the Predictive Maintenance Agent using Cursor+Claude. It covers NAT (NeMo Agent Toolkit) concepts, architecture, and development patterns.

## Table of Contents
- [What is NAT (NeMo Agent Toolkit)](#what-is-nat-nemo-agent-toolkit)
- [Architecture Overview](#architecture-overview)
- [Development Environment](#development-environment)
- [Adding New Tools/Functions](#adding-new-toolsfunctions)
- [Workspace Utilities System](#workspace-utilities-system)
- [Configuration Management](#configuration-management)
- [Database Integration](#database-integration)
- [Code Generation Assistant](#code-generation-assistant)
- [Evaluation and Testing](#evaluation-and-testing)
- [Development Best Practices](#development-best-practices)
- [Debugging and Troubleshooting](#debugging-and-troubleshooting)
- [File Structure Guide](#file-structure-guide)

## What is NAT (NeMo Agent Toolkit)

**NeMo Agent Toolkit (NAT)** is NVIDIA's framework for building agentic AI workflows that can reason, plan, and execute complex tasks using LLMs and tools.

### Core Concepts

**Functions/Tools**: Reusable components that agents can invoke to perform specific tasks
- SQL retrieval, RUL prediction, plotting, anomaly detection, code generation

**Agents**: AI entities that can reason and make decisions
- `ReAct Agent`: Uses reasoning and acting patterns
- `Reasoning Agent`: High-level planning and orchestration

**Workflows**: Complete systems that combine agents and functions
- Our predictive maintenance workflow uses a reasoning agent that orchestrates a ReAct agent

**Builders**: NAT's dependency injection system
- Automatically manages LLMs, functions, and their dependencies
- Handles configuration and initialization

**Configuration-Driven**: Everything is defined in YAML configuration files
- LLMs, functions, agents, and workflows are all configured declaratively
- No hard-coding, everything is composable and reusable

### Key NAT Patterns

```python
# 1. Function Registration
@register_function(config_type=YourToolConfig)
async def your_tool(config: YourToolConfig, builder: Builder):
    # Tool implementation
    yield FunctionInfo.from_fn(fn=_inner_function, description="Tool description")

# 2. Configuration Classes
class YourToolConfig(FunctionBaseConfig, name="your_tool"):
    parameter: str = Field(description="Parameter description")
    llm_name: LLMRef = Field(description="LLM reference")

# 3. Builder Pattern for Dependencies
llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
other_function = builder.get_function(config.other_function_ref)
```

## Architecture Overview

### Multi-Agent Architecture

```
User Query → Reasoning Agent → Plan → ReAct Agent → Tools → Results
```

**Reasoning Agent**: 
- High-level planning and orchestration
- Generates step-by-step execution plans
- Handles complex task decomposition

**ReAct Agent (Data Analysis Assistant)**:
- Executes the plan using available tools
- Follows Reasoning-Acting pattern
- Handles tool selection and coordination

**Tools/Functions**:
- `sql_retriever`: Database queries using NIM LLM + Vanna
- `predict_rul`: RUL prediction using MOMENT foundation model
- `plot_*`: Various visualization tools
- `anomaly_detection`: Time series anomaly detection
- `code_generation_assistant`: Custom Python code generation and execution

### Workspace Utilities Innovation

**Problem Solved**: Multi-agent instruction passing loses context and creates unreliable results

**Solution**: Pre-built utilities in `/workspace/utils/` that can be invoked with simple instructions

```python
# Instead of complex pseudo-code through agent chains:
"Apply piecewise RUL transformation using workspace utility"

# Results in reliable, pre-tested code:
import utils
utils.apply_piecewise_rul_transformation(file_path, maxlife=125)
```

## Development Environment

### Required Setup

```bash
# 1. Conda Environment
conda create -n pdm-nat python=3.11
conda activate pdm-nat

# 2. NAT Installation
cd /path/to/NeMo-Agent-Toolkit
uv sync --all-groups --all-extras

# 3. PDM Agent Installation
cd /path/to/predictive_maintenance_agent
uv pip install -e .

# 4. Environment Variables
export NVIDIA_API_KEY="your-key"
```

### Code Execution Sandbox

**Critical Component**: Docker-based Python execution environment

```bash
# Start sandbox (maps output_data to /workspace)
cd /path/to/NeMo-Agent-Toolkit/src/nat/tool/code_execution/
./local_sandbox/start_local_sandbox.sh local-sandbox /path/to/output_data/
```

**Key Insight**: The sandbox provides isolated Python execution with volume mounting for file persistence.

## Adding New Tools/Functions

### Step 1: Create Tool Configuration

```python
# src/predictive_maintenance_agent/your_category/your_tool.py
from pydantic import Field
from nat.data_models.function import FunctionBaseConfig
from nat.data_models.component_ref import LLMRef

class YourToolConfig(FunctionBaseConfig, name="your_tool_name"):
    """
    Configuration for your custom tool.
    """
    llm_name: LLMRef = Field(description="LLM to use for reasoning")
    output_folder: str = Field(description="Output directory", default="/workspace")
    custom_param: int = Field(description="Custom parameter", default=100)
```

### Step 2: Implement Tool Function

```python
from nat.builder.builder import Builder
from nat.builder.function_info import FunctionInfo
from nat.cli.register_workflow import register_function

@register_function(config_type=YourToolConfig)
async def your_tool(config: YourToolConfig, builder: Builder):
    # Get dependencies from builder
    llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
    
    async def _inner_function(input_param: str) -> str:
        """
        Your tool's main logic.
        
        Args:
            input_param: Description of parameter
            
        Returns:
            String result with details
        """
        # Implementation here
        result = f"Processed {input_param} with param {config.custom_param}"
        return result
    
    yield FunctionInfo.from_fn(
        fn=_inner_function,
        description="Clear description of what this tool does and when to use it"
    )
```

### Step 3: Register in Configuration

```yaml
# configs/config-reasoning.yml
functions:
  your_tool_name:
    _type: your_tool_name
    llm_name: coding_llm
    output_folder: "/workspace"
    custom_param: 150

  data_analysis_assistant:
    # Add your tool to the ReAct agent's available tools
    tool_names: [sql_retriever, your_tool_name, predict_rul, ...]
```

### Step 4: Update System Prompts (if needed)

```yaml
# Add tool usage guidelines to system prompts
system_prompt: |
  ### TOOL GUIDELINES ###
  **Your New Tool**:
  - Use your_tool_name for [specific use case]
  - Input format: [description]
  - Output: [description]
```

### Tool Development Best Practices

1. **Error Handling**: Always include comprehensive error handling
2. **Type Hints**: Use proper type annotations for all parameters
3. **Documentation**: Clear docstrings and descriptions
4. **Configuration**: Make parameters configurable, not hard-coded
5. **Testing**: Create test cases for your tool
6. **Logging**: Add appropriate logging for debugging

## Workspace Utilities System

### Philosophy

**Traditional Approach** (Problematic):
```
Reasoning Agent → Complex Algorithm → ReAct Agent → Degraded Instructions → Code Gen → Error-Prone Code
```

**Utilities Approach** (Robust):
```
Reasoning Agent → "Use utility X" → ReAct Agent → Code Gen → Pre-built Utility → Reliable Result
```

### Creating New Utilities

```python
# output_data/utils/your_utility.py
def your_custom_utility(file_path: str, param: int = 100) -> str:
    """
    Your utility function with comprehensive error handling.
    
    Args:
        file_path: Path to input file
        param: Your parameter
        
    Returns:
        Success message with transformation details
        
    Raises:
        FileNotFoundError: If file doesn't exist
        ValueError: If data is invalid
    """
    try:
        # Load and validate data
        with open(file_path, 'r') as f:
            data = json.load(f)
        
        if not data:
            raise ValueError("File is empty")
            
        # Your transformation logic
        result_count = len(data)
        
        # Save back (in-place modification preferred)
        with open(file_path, 'w') as f:
            json.dump(data, f, indent=2)
            
        return (f"✅ Custom utility executed successfully!\n"
                f"📁 File: {file_path}\n"
                f"📊 Records processed: {result_count}\n"
                f"⚙️ Parameter used: {param}")
        
    except FileNotFoundError:
        return f"❌ Error: File '{file_path}' not found"
    except Exception as e:
        return f"❌ Unexpected error: {str(e)}"
```

```python
# output_data/utils/__init__.py - Export your utility
from .your_utility import your_custom_utility

__all__ = ['your_custom_utility', ...]
```

### Updating Code Generation Assistant

```python
# Update system prompt to include your utility
**WORKSPACE UTILITIES AVAILABLE:**
- utils.your_custom_utility(file_path, param=100): Description of what it does
```

## Configuration Management

### Configuration Structure

```yaml
# config-reasoning.yml structure
general:          # Global settings (logging, telemetry)
llms:            # LLM definitions with models and parameters
embedders:       # Embedding models for vector operations
functions:       # All available tools/functions
workflow:        # Main workflow definition
eval:           # Evaluation configuration
```

### Key Configuration Patterns

**LLM References**:
```yaml
llms:
  my_llm:
    _type: nim
    model_name: "nvidia/llama-3.3-nemotron-super-49b-v1"
    
functions:
  my_tool:
    llm_name: my_llm  # Reference to LLM above
```

**Function References**:
```yaml
functions:
  code_execution:
    _type: code_execution
    uri: http://127.0.0.1:6000/execute
    
  code_generation_assistant:
    code_execution_tool: code_execution  # Reference to function above
```

**Path Management**:
```yaml
# Use absolute paths for critical directories
db_path: "/full/path/to/database/nasa_turbo.db"

# Use relative paths within workspace
output_folder: "/workspace"  # Maps to mounted volume
```

## Database Integration

### Database Setup

The system uses **SQLite** with NASA turbofan datasets:

```python
# Database structure
- FD001_train, FD001_test, FD001_RUL (and FD002-FD004)
- Columns: unit_number, time_in_cycles, setting_1, setting_2, setting_3, sensor_1...sensor_21
- RUL tables: unit_number, RUL
```

### SQL Tool (Vanna Integration)

**Key Component**: `sql_retriever` uses Vanna AI + NIM LLM for natural language to SQL:

```python
# Vanna training data includes:
- Table schemas (DDL)
- Example queries  
- Domain documentation
- Question-SQL pairs
```

**Usage Pattern**:
```python
# In agent prompts:
"Retrieve time in cycles and sensor data for engine unit 24"
# Automatically generates and executes SQL
```

### Database Development Tips

1. **Schema Understanding**: Learn the NASA turbofan data structure
2. **Vanna Training**: Add new question-SQL pairs to improve performance
3. **Query Patterns**: Study common patterns in `vanna_training_data.yaml`
4. **Data Validation**: Always validate query results before processing

## Code Generation Assistant

### Architecture

**Purpose**: Generate and execute Python code based on natural language instructions

**Key Features**:
- Automatic workspace utilities discovery
- In-place file modification preference
- Comprehensive error handling and completion
- Integration with code execution sandbox

### Customizing Code Generation

**System Prompt Structure**:
```python
system_prompt = """
**WORKSPACE UTILITIES AVAILABLE:**
[List of available utilities with descriptions]

**CODE REQUIREMENTS:**
[Coding standards and requirements]

**FILE MODIFICATION PREFERENCE:**
[Guidelines for file handling]
"""
```

**User Prompt Structure**:
```python
user_prompt = """
**INSTRUCTIONS:**
{instructions}

**UTILITIES REMINDER:**
[Remind about available utilities]
"""
```

### Code Generation Best Practices

1. **Utilities First**: Always check if workspace utilities can handle the task
2. **Error Handling**: Include comprehensive error handling in generated code
3. **File Operations**: Prefer in-place modifications over creating new files
4. **Path Handling**: Use consistent path patterns (`/workspace/filename`)
5. **Output Messages**: Generate clear success/failure messages

## Evaluation and Testing

### Evaluation System

**Configuration**:
```yaml
eval:
  general:
    output:
      dir: "eval_output"
    dataset:
      _type: json
      file_path: "eval_data/eval_set_master.json"
  evaluators:
    multimodal_eval:
      _type: multimodal_llm_judge_evaluator
```

**Running Evaluations**:
```bash
nat eval --config_file configs/config-reasoning.yml
```

### Testing Patterns

**Unit Testing Tools**:
```python
# Test individual functions
pytest test_pdm_workflow.py -k "test_sql_retriever" -v

# End-to-end testing
pytest test_pdm_workflow.py -m e2e -v
```

**Manual Testing**:
```python
# Test workflow programmatically
from nat.cli.build_workflow import build_workflow
workflow = await build_workflow("configs/config-reasoning.yml")
result = await workflow("Your test query")
```

### Creating Test Cases

```python
# tests/test_your_tool.py
import pytest
from your_tool import your_tool_function

@pytest.mark.asyncio
async def test_your_tool():
    # Arrange
    config = YourToolConfig(param=100)
    
    # Act
    result = await your_tool_function("test input")
    
    # Assert
    assert "expected_output" in result
    assert result.startswith("✅")
```

## Development Best Practices

### Code Organization

1. **Separation of Concerns**: Keep tools, utilities, and configuration separate
2. **Function Naming**: Use descriptive names that indicate purpose
3. **Configuration**: Make everything configurable, avoid hard-coding
4. **Error Handling**: Always include comprehensive error handling
5. **Documentation**: Document all functions, parameters, and use cases

### Agent Design Patterns

**Tool Selection Guidelines**:
```python
# In system prompts
"""
### TOOL SELECTION GUIDELINES ###
1. SQL Retrieval: For database queries and data extraction
2. Prediction Tools: For RUL and anomaly detection
3. Plotting Tools: For visualizations
4. Code Generation: For custom processing and complex transformations
5. Workspace Utilities: For reliable, pre-tested operations
"""
```

**Workflow Patterns**:
1. **Data Extraction**: Use SQL retrieval first
2. **Data Processing**: Use utilities or plotting tools
3. **Custom Logic**: Use code generation assistant
4. **Results**: Return file paths and summaries

### Configuration Management

```yaml
# Use environment-specific configs
configs/
  ├── config-reasoning.yml      # Main configuration
  ├── config-development.yml    # Development overrides
  └── config-evaluation.yml     # Evaluation-specific settings
```

### Version Control Best Practices

```bash
# Git patterns for this project
.gitignore:
  output_data/          # Generated files
  eval_output/          # Evaluation results
  *.log                 # Log files
  .env                  # Environment variables
```

## NAT Version Compatibility

### NAT 1.2.1 vs 1.3.0

**Current Version**: NAT 1.2.1 (with pydantic 2.10.x)

**Key Compatibility Rules**:

1. **Optional String Fields**:
```python
# ❌ WRONG - Will fail validation
elasticsearch_url: str = Field(default=None)

# ✅ CORRECT - Use Optional for nullable strings
from typing import Optional
elasticsearch_url: Optional[str] = Field(default=None)
```

2. **Reference Field Types (NAT 1.2.1)**:
```python
# NAT 1.2.1 uses plain strings for references
llm_name: str = Field(description="LLM reference")
embedding_name: str = Field(description="Embedder reference")
```

3. **Reference Field Types (NAT 1.3.0 - Future)**:
```python
# NAT 1.3.0 requires typed references
from nat.data_models.component_ref import LLMRef, EmbedderRef, FunctionRef

llm_name: LLMRef = Field(description="LLM reference")
embedding_name: EmbedderRef = Field(description="Embedder reference")
code_execution_tool: FunctionRef = Field(description="Function reference")
```

4. **YAML Configuration Quoting**:
```yaml
# Always quote string references in YAML configs for pydantic 2.10+
functions:
  sql_retriever:
    llm_name: "sql_llm"              # Quoted
    embedding_name: "vanna_embedder"  # Quoted
    vector_store_type: "chromadb"     # Quoted
    db_type: "sqlite"                 # Quoted
    
  data_analysis_assistant:
    tool_names: [
      "sql_retriever",   # All tool names quoted
      "predict_rul",
      "plot_distribution"
    ]
```

### Pydantic 2.10+ Best Practices

**Type Annotations**:
```python
from typing import Optional

class ToolConfig(FunctionBaseConfig):
    # Required fields
    required_param: str = Field(description="Must be provided")
    
    # Optional fields with None default
    optional_param: Optional[str] = Field(default=None, description="Can be None")
    
    # Optional fields with non-None default
    param_with_default: str = Field(default="default_value", description="Has default")
    
    # Numeric fields (can use None without Optional if you want)
    max_retries: int = Field(default=3, description="Number of retries")
```

**Common Validation Errors**:
```
ValidationError: Input should be a valid string [input_value=None, input_type=NoneType]
→ Solution: Use Optional[str] instead of str for fields with default=None

ValidationError: functions: Input should be a valid string (4 times)
→ Solution: Quote all string values in YAML config, especially references
```

### Upgrading to NAT 1.3.0 (Future)

When upgrading, you'll need to:

1. Update pyproject.toml:
```toml
dependencies = [
  "nvidia-nat[profiling,langchain,telemetry]==1.3.0",
  "pydantic>=2.11.0,<3.0.0",
]
```

2. Update all tool configs:
```python
# Before (NAT 1.2.1)
llm_name: str = Field(...)

# After (NAT 1.3.0)
from nat.data_models.component_ref import LLMRef
llm_name: LLMRef = Field(...)
```

3. Update evaluator configs:
```python
# multimodal_llm_judge_evaluator_register.py
# llm_judge_evaluator_register.py
from nat.data_models.component_ref import LLMRef
llm_name: LLMRef = Field(...)
```

4. Keep Optional[str] for nullable fields (both versions need this)

## Debugging and Troubleshooting

### Common Issues and Solutions

**Issue**: "Function not found in builder"
```python
# Solution: Check function registration and config
@register_function(config_type=YourConfig)  # Must have this decorator
# And in config:
functions:
  your_function_name:    # Must match the name in config class
    _type: your_function_name
```

**Issue**: "Code execution failed"
```bash
# Check sandbox status
curl http://localhost:6000/health

# Check sandbox logs  
docker logs local-sandbox

# Restart sandbox
docker stop local-sandbox
./local_sandbox/start_local_sandbox.sh ...
```

**Issue**: "Workspace utilities not found"
```python
# Ensure path setup in generated code
import sys
sys.path.append('/workspace')
import utils
```

### Debugging Tools

**Logging Configuration**:
```yaml
general:
  telemetry:
    logging:
      console:
        level: DEBUG  # Increase verbosity
      file:
        path: "debug.log"
        level: DEBUG
```

**Phoenix Tracing** (Optional):
```yaml
# Uncomment in config for detailed tracing
tracing:
  phoenix:
    endpoint: http://localhost:6006/v1/traces
    project: pdm-debug
```

**Manual Testing**:
```python
# Test specific components
from nat.cli.build_workflow import build_workflow

workflow = await build_workflow("configs/config-reasoning.yml")
result = await workflow("Debug query")
print(result)
```

## File Structure Guide

```
predictive_maintenance_agent/
├── configs/
│   └── config-reasoning.yml          # Main configuration
├── src/predictive_maintenance_agent/
│   ├── register.py                   # Function registration
│   ├── plotting/                     # Visualization tools
│   │   ├── code_generation_assistant.py
│   │   ├── plot_*.py                # Plotting functions
│   └── predictors/                   # ML prediction tools
│       ├── predict_rul_tool.py
│       └── moment_*.py               # MOMENT model integration
├── data/                             # Database files
│   └── nasa_turbo.db
├── output_data/                      # Generated outputs (maps to /workspace)
│   ├── utils/                        # Workspace utilities
│   │   ├── __init__.py
│   │   └── rul_transforms.py
│   └── *.json, *.html, *.png        # Generated files
├── eval_data/                        # Evaluation datasets
│   ├── eval_set_master.json
│   └── eval_set_test.json
├── eval_output/                      # Evaluation results
├── tests/                            # Test files
└── vanna_training_data.yaml          # SQL training data
```

### Key Files to Understand

1. **`register.py`**: Central function registration
2. **`config-reasoning.yml`**: Complete system configuration  
3. **`code_generation_assistant.py`**: Core code generation logic
4. **`utils/`**: Workspace utilities system
5. **`vanna_training_data.yaml`**: SQL generation training data

## Development Workflow

### Typical Development Cycle

1. **Understand Requirements**: What new capability is needed?
2. **Choose Approach**: New tool, utility, or configuration change?
3. **Implement**: Following NAT patterns and best practices
4. **Test**: Unit tests and integration testing
5. **Document**: Update prompts, configurations, and documentation
6. **Evaluate**: Run evaluation suite to ensure no regressions

### Environment Setup Checklist

- [ ] Conda environment activated (`pdm-nat`)
- [ ] NAT installed with all dependencies
- [ ] Environment variables set (`NVIDIA_API_KEY`)
- [ ] Database setup complete
- [ ] Code execution sandbox running
- [ ] Configuration paths updated for your system

### Before Committing

- [ ] All tests pass
- [ ] Code follows project patterns
- [ ] Configuration is environment-agnostic
- [ ] Documentation is updated
- [ ] No sensitive information in code
- [ ] Evaluation shows no regressions

---

This guide should provide comprehensive coverage for developing with this agentic workflow. Remember that NAT is designed for composability and configuration-driven development, so leverage these patterns for maintainable and robust implementations.
