Skip to content

test: add rigging Model integration and tests for core data models#28

Merged
l50 merged 1 commit intomainfrom
jayson/cap-832-integrate-rigging-framework-across-ares-codebase
Jan 11, 2026
Merged

test: add rigging Model integration and tests for core data models#28
l50 merged 1 commit intomainfrom
jayson/cap-832-integrate-rigging-framework-across-ares-codebase

Conversation

@l50
Copy link
Copy Markdown
Contributor

@l50 l50 commented Jan 11, 2026

Key Changes:

  • Refactored core SOC and red team data models to inherit from rigging's Model
  • Enabled automatic serialization, LLM output parsing, and XML export for models
  • Added comprehensive tests covering all models and parsing utilities
  • Re-exported rigging parsing utilities directly from the models module

Added:

  • Test suite for ares.core.models - Introduced tests/test_models.py to validate
    model creation, serialization (dict/XML), computed fields, parsing utilities,
    and pydantic validation across all SOC and red team models

Changed:

  • Core data models (Evidence, TimelineEvent, InvestigativeQuestion,
    Target, Host, User, Credential, Hash, Share) now inherit from
    rigging's Model for enhanced serialization and XML/LLM parsing
  • Switched from dataclasses to Model-based classes, enabling model validation,
    .model_dump(), .model_validate(), and .to_xml() methods
  • Updated field definitions to use rigging's wrapped and element utilities
    for robust XML/LLM parsing, especially for lists of strings
  • to_dict methods updated to use Model's model_dump(mode="json") for
    backward compatibility and consistent output
  • Re-exported rigging parsing utilities (parse, parse_many, parse_set,
    try_parse, etc.) in ares.core.models for convenient import and usage
  • Expanded module docstring with usage examples and improved documentation

Removed:

  • Legacy dataclass usage and manual default factories for fields in all models
  • Custom to_dict implementations replaced with Model's serialization methods
  • Redundant or obsolete imports and docstring content to streamline the module

…l tests

**Added:**

- Introduced a comprehensive `tests/test_models.py` suite for SOC and red team models,
  covering instantiation, serialization, validation, parsing utilities, and XML output
- Added detailed docstring and example usage for parsing LLM output using rigging

**Changed:**

- Migrated core SOC and red team data models (`Evidence`, `TimelineEvent`,
  `InvestigativeQuestion`, `Target`, `Host`, `User`, `Credential`, `Hash`, `Share`)
  from `@dataclass` to inherit from `rigging.Model` for automatic serialization and
  LLM parsing support
- Updated field definitions to use rigging's `wrapped` and `element` helpers for
  XML compatibility and defaults, and `Field` for default factories
- Replaced custom `to_dict` methods with calls to `model_dump(mode="json")` for
  standardized serialization
- Re-exported rigging parsing utilities (`parse`, `parse_set`, etc.) directly from
  the module for easier access
- Updated `__all__` to include new symbols and parsing helpers
- Improved and expanded module-level documentation to clarify usage and design

**Removed:**

- Removed legacy `@dataclass` decorators and direct `dataclasses.field` usage in
  favor of rigging and pydantic-based definitions
- Eliminated manual serialization logic now handled by rigging/pydantic methods
@linear
Copy link
Copy Markdown

linear Bot commented Jan 11, 2026

CAP-832 Integrate Rigging Framework Across Ares Codebase

Description:
Enhance the integration of the rigging framework throughout the Ares codebase. Transition away from manual data modeling, parsing, and prompt handling by leveraging rigging’s structured models, prompt functions, and parsing utilities. This will improve maintainability, enable type-safe LLM interactions, and facilitate future enhancements.


Objective:

Fully adopt the rigging framework for data modeling, prompt management, and LLM response parsing in Ares, replacing legacy approaches and reducing technical debt. Ensure all key components use rigging’s features for serialization, validation, and output schema generation.


Scope of Work:

  • Convert manual dataclasses (e.g., Evidence, InvestigativeQuestion) to rg.Model classes in src/ares/core/models.py and update relevant usages.
  • Integrate PyramidLevel enum and update InvestigationState to utilize rigging models.
  • Implement schema generation for agent prompts using rigging models.
  • Replace Jinja2-based prompt templates with @rg.prompt functions; migrate or deprecate src/ares/core/templates.py.
  • Add return type models to prompt functions for structured LLM outputs.
  • Evaluate and ensure compatibility with the dreadnode Agent system.
  • Identify all manual LLM response parsing locations (regex/JSON) and replace with rg.parse() and corresponding models.
  • Add fallback handling for parsing failures.
  • Document migration steps and update relevant references (e.g., TODO-RIGGING.md).

Dependencies:

  • Familiarity with the rigging documentation and API.
  • Access to and understanding of existing Ares codebase, especially files listed in the overview.
  • Coordination with maintainers of the dreadnode Agent system for compatibility review.
  • Review of TODO-RIGGING.md for additional context.
  • No external blockers identified.

Acceptance Criteria:

  1. All major dataclasses and enums in src/ares/core/models.py and related files are refactored to use rg.Model.
  2. Jinja2 prompt templates are replaced with @rg.prompt functions, and legacy template files are migrated or deprecated.
  3. LLM response parsing is handled via rg.parse() with appropriate models; no manual regex/JSON parsing remains.
  4. Structured schema generation for agent prompts is present and validated.
  5. Compatibility with the dreadnode Agent system is maintained or improved.
  6. All changes are documented, and references in TODO-RIGGING.md are updated.
  7. All affected unit and integration tests are updated and pass successfully.

Additional Notes:

  • For additional context and rationale, see the Rigging GitHub and Rigging Documentation.
  • Use migration phases outlined in the input as guidance for implementation order.
  • Consider future extensibility for prompt caching and pipeline refactoring, but focus on high-priority items in this task.
  • For questions or architectural considerations, consult with the Ares lead engineer.

@l50 l50 merged commit 470d975 into main Jan 11, 2026
8 checks passed
@l50 l50 deleted the jayson/cap-832-integrate-rigging-framework-across-ares-codebase branch January 11, 2026 02:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant