Grid Optimization - NAT v1.2.1 Integration Guide

Overview

This project is now fully integrated with NVIDIA NeMo Agent Toolkit (NAT) v1.2.1, using the latest function registration patterns and workflow configurations.

What Changed

1. Updated Dependencies

Changed from: aiq imports
Changed to: nat imports (NAT v1.2.1)
Added nvidia-nat[langchain]~=1.2.1 as a core dependency

2. Module Imports

All NAT-related imports now use the nat namespace:

from nat.builder.builder import Builder
from nat.builder.function_info import FunctionInfo
from nat.cli.register_workflow import register_function
from nat.data_models.function import FunctionBaseConfig

3. Entry Points

Updated pyproject.toml to use NAT-compliant entry points:

[project.entry-points."nat.components"]
nat_grid_optimization = "grid_optimization.integrations.nat.register"

Installation

Prerequisites

Python 3.11-3.13
NVIDIA NeMo Agent Toolkit v1.2.1+

Setup

# 1. Clone and navigate to repository
cd Grid_Optm

# 2. Create virtual environment
python -m venv .venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate

# 3. Install with NAT dependencies
pip install -e .

# 4. Initialize database
python -c "from grid_optimization.core.initialization import init_test_data; init_test_data()"

# 5. Set up environment variables
echo "OPENAI_API_KEY=your-api-key" > .env
# OR for NVIDIA NIM:
echo "NVIDIA_API_KEY=your-nvidia-api-key" > .env

Usage with NAT CLI

Run Grid Optimization Workflow

# Using the NAT CLI
nat run configs/grid_config.yml

# Interactive mode
nat run configs/grid_config.yml --interactive

# With specific query
nat run configs/grid_config.yml --query "Optimize grid for us-west region"

Available Tools/Functions

grid_optimize - Optimize power grid for a region
```
Example: "Optimize the grid for us-west"
```
grid_status - Get current grid status and optimization history
```
Example: "Show me the status of us-east grid"
```

grid_analyze - Analyze specific grid metrics

Example: "Analyze efficiency metrics for us-central"

current_datetime - Get current timestamp
```
Example: "What time is it?"
```

Legacy Compatibility

The system maintains backward compatibility with old function names:

optimize_grid (legacy) → grid_optimize (modern)
show_last_optimization (legacy) → grid_status (modern)

Configuration

LLM Configuration

The default configuration (configs/grid_config.yml) supports both:

NVIDIA NIM (default):

llms:
  nim_llm:
    _type: nim
    model_name: meta/llama-3.1-70b-instruct
    temperature: 0.0
    max_tokens: 2048

OpenAI (alternative):

llms:
  openai_llm:
    _type: openai
    model_name: gpt-4
    temperature: 0.0
    max_tokens: 2048

workflow:
  llm_name: openai_llm  # Change this line

Architecture

Grid_Optm/
├── grid_optimization/
│   ├── core/                           # Core optimization engine
│   │   ├── operations.py              # Grid optimization algorithms
│   │   ├── database.py                # SQLAlchemy models
│   │   └── models.py                  # Pydantic models
│   └── integrations/
│       └── nat/
│           └── register.py            # NAT v1.2.1 function registry
├── configs/
│   └── grid_config.yml                # NAT workflow configuration
└── pyproject.toml                     # Updated with NAT v1.2.1 deps

Registered Functions

All functions are registered using the NAT v1.2.1 pattern:

@register_function(config_type=GridOptimizeToolConfig)
async def grid_optimize(tool_config: GridOptimizeToolConfig, builder: Builder):
    async def _grid_optimize(region: str = "default") -> str:
        # Implementation
        ...

    yield FunctionInfo.from_fn(
        _grid_optimize,
        description="Optimizes power grid configuration..."
    )

Testing

# Run unit tests
pytest tests/unit/

# Run integration tests
pytest tests/integration/

# Test NAT integration specifically
pytest tests/test_nat_integration.py -v

Example Queries

Once running with NAT:

User: "Optimize the power grid for the us-west region"
→ Calls grid_optimize(region="us-west")

User: "What's the current status of the us-east grid?"
→ Calls grid_status(region="us-east")

User: "Analyze efficiency metrics for us-central"
→ Calls grid_analyze(region="us-central", metric="efficiency")

Available Regions

us-west - Western United States
us-east - Eastern United States
us-central - Central United States

Troubleshooting

ImportError: No module named 'nat'

# Ensure NAT is installed
pip install nvidia-nat[langchain]~=1.2.1

Database not initialized

python -c "from grid_optimization.core.initialization import init_test_data; init_test_data()"

API Key issues

# Check environment variables
env | grep API_KEY

# Set API key
export NVIDIA_API_KEY="your-key"
# OR
export OPENAI_API_KEY="your-key"

Migration from AIQ to NAT

If you have old code using aiq imports:

Replace aiq with nat in imports
Update pyproject.toml entry points from aiq.components to nat.components
Add nvidia-nat[langchain]~=1.2.1 to dependencies
Update Python version requirement to >=3.11,<3.14

Additional Resources

Support

For issues related to:

Grid Optimization: Check the main README.md
NAT Integration: Refer to NeMo Agent Toolkit docs
General Issues: Open an issue in the repository

Last Updated: October 2025 NAT Version: 1.2.1+ Status: ✅ Fully Integrated

FilesExpand file tree

NAT_INTEGRATION.md

Latest commit

History