Refactor DSPy example

Author: KennyVaneetvelde

atomic-examples/dspy-integration/dspy_integration/__init__.py

@@ -1,20 +1,50 @@
 """
-DSPy + Atomic Agents Integration
+DSPy + Atomic Agents Integration Package.
 
-This module provides a seamless integration between DSPy's declarative prompt optimization
-framework and Atomic Agents' type-safe structured output system.
+This package demonstrates how to combine DSPy's automatic prompt optimization
+with Atomic Agents' type-safe structured outputs.
 
-Key Components:
-- DSPyAtomicModule: A DSPy module that wraps Atomic Agents for structured outputs
-- AtomicSignature: Bridge between Pydantic schemas and DSPy signatures
-- Optimization utilities for improving agent performance with few-shot learning
+Package Structure:
+    domain/      - Core business logic (models, datasets, evaluation)
+    stages/      - Demonstration stages (dspy, atomic, combined)
+    presentation/ - UI layer (Rich console output)
+    bridge.py    - DSPy ↔ Atomic Agents integration module
+
+Quick Start:
+    >>> from dspy_integration import DSPyAtomicModule, MovieReviewInput, MovieGenreOutput
+    >>> module = DSPyAtomicModule(
+    ...     input_schema=MovieReviewInput,
+    ...     output_schema=MovieGenreOutput,
+    ...     use_chain_of_thought=True,
+    ... )
+    >>> result = module.run_validated(review="Amazing action movie!")
+    >>> print(result.genre)  # Type-safe output!
+
+Run Demo:
+    uv run python -m dspy_integration.main
 """
 
+# Domain exports
+from dspy_integration.domain.models import (
+    GENRES,
+    GenreType,
+    MovieGenreOutput,
+    MovieReviewInput,
+    EvalResult,
+)
+from dspy_integration.domain.datasets import TRAINING_DATASET, TEST_DATASET
+from dspy_integration.domain.evaluation import evaluate_predictions
+
+# Bridge exports
 from dspy_integration.bridge import (
     DSPyAtomicModule,
-    pydantic_to_dspy_fields,
+    DSPyAtomicPipeline,
+    create_dspy_example,
     create_dspy_signature_from_schemas,
+    pydantic_to_dspy_fields,
 )
+
+# Original schemas (for backwards compatibility)
 from dspy_integration.schemas import (
     SentimentInputSchema,
     SentimentOutputSchema,
@@ -24,14 +54,47 @@
     SummaryOutputSchema,
 )
 
+# Stage exports (for advanced usage)
+from dspy_integration.stages import (
+    run_stage1_raw_dspy,
+    run_stage2_raw_atomic_agents,
+    run_stage3_combined,
+)
+
+__version__ = "0.1.0"
+
 __all__ = [
+    # Version
+    "__version__",
+    # Domain - Types
+    "GENRES",
+    "GenreType",
+    # Domain - Schemas (new)
+    "MovieGenreOutput",
+    "MovieReviewInput",
+    # Domain - Data structures
+    "EvalResult",
+    # Domain - Datasets
+    "TRAINING_DATASET",
+    "TEST_DATASET",
+    # Domain - Evaluation
+    "evaluate_predictions",
+    # Bridge - Core classes
     "DSPyAtomicModule",
-    "pydantic_to_dspy_fields",
+    "DSPyAtomicPipeline",
+    # Bridge - Utilities
+    "create_dspy_example",
     "create_dspy_signature_from_schemas",
+    "pydantic_to_dspy_fields",
+    # Original schemas (backwards compatibility)
     "SentimentInputSchema",
     "SentimentOutputSchema",
     "QuestionInputSchema",
     "AnswerOutputSchema",
     "SummaryInputSchema",
     "SummaryOutputSchema",
+    # Stages - Runners
+    "run_stage1_raw_dspy",
+    "run_stage2_raw_atomic_agents",
+    "run_stage3_combined",
 ]

atomic-examples/dspy-integration/dspy_integration/bridge.py

@@ -7,10 +7,11 @@
 3. Applying DSPy optimizers (BootstrapFewShot, MIPROv2, etc.) to improve agent performance
 """
 
-from typing import Type, Dict, Any, Optional, List, get_type_hints, get_origin, get_args, Literal
+from typing import Any, Dict, List, Literal, Optional, Type, get_args, get_origin
+
 import dspy
-from pydantic import BaseModel, Field
-from pydantic.fields import FieldInfo
+from pydantic import BaseModel
+
 from atomic_agents.base.base_io_schema import BaseIOSchema
 
 
@@ -38,11 +39,11 @@ def python_type_to_dspy_type(python_type: Any) -> Any:
         return list
 
     # Handle Optional types
-    if origin is type(None) or (hasattr(origin, '__origin__') and origin.__origin__ is type(None)):
+    if origin is type(None) or (hasattr(origin, "__origin__") and origin.__origin__ is type(None)):
         return python_type
 
     # Handle Union types (including Optional)
-    if hasattr(origin, '__name__') and origin.__name__ == 'UnionType':
+    if hasattr(origin, "__name__") and origin.__name__ == "UnionType":
         args = get_args(python_type)
         # Filter out NoneType for Optional handling
         non_none_args = [a for a in args if a is not type(None)]
@@ -58,8 +59,7 @@ def python_type_to_dspy_type(python_type: Any) -> Any:
 
 
 def pydantic_to_dspy_fields(
-    schema: Type[BaseModel],
-    field_type: str = "input"
+    schema: Type[BaseModel], field_type: str = "input"
 ) -> Dict[str, tuple]:
     """
     Convert Pydantic schema fields to DSPy field definitions.
@@ -184,9 +184,7 @@ def __init__(
 
         # Create DSPy signature from schemas
         self.signature = create_dspy_signature_from_schemas(
-            input_schema,
-            output_schema,
-            instructions
+            input_schema, output_schema, instructions
         )
 
         # Create the predictor
@@ -300,7 +298,7 @@ def forward(self, **kwargs) -> Dict[str, Any]:
             current_input = {
                 k: getattr(prediction, k)
                 for k in dir(prediction)
-                if not k.startswith('_') and not callable(getattr(prediction, k))
+                if not k.startswith("_") and not callable(getattr(prediction, k))
             }
 
         return results
@@ -336,10 +334,7 @@ def create_dspy_example(
         **validated_output.model_dump(),
     }
 
-    # Create DSPy example with output fields marked
-    output_keys = list(output_schema.model_fields.keys())
-    example = dspy.Example(**example_data).with_inputs(
-        *list(input_schema.model_fields.keys())
-    )
+    # Create DSPy example with input fields marked
+    example = dspy.Example(**example_data).with_inputs(*list(input_schema.model_fields.keys()))
 
     return example

atomic-examples/dspy-integration/dspy_integration/domain/__init__.py

@@ -0,0 +1,37 @@
+"""
+Domain layer for DSPy + Atomic Agents integration.
+
+This package contains:
+- models: Pydantic schemas and data transfer objects
+- datasets: Training and test data
+- evaluation: Metrics and evaluation utilities
+
+Following Clean Architecture principles, this layer has no dependencies
+on external frameworks (except Pydantic for data modeling).
+"""
+
+from dspy_integration.domain.models import (
+    GenreType,
+    GENRES,
+    MovieGenreOutput,
+    MovieReviewInput,
+    EvalResult,
+)
+from dspy_integration.domain.datasets import TRAINING_DATASET, TEST_DATASET
+from dspy_integration.domain.evaluation import evaluate_predictions
+
+__all__ = [
+    # Types
+    "GenreType",
+    "GENRES",
+    # Schemas
+    "MovieGenreOutput",
+    "MovieReviewInput",
+    # Data structures
+    "EvalResult",
+    # Datasets
+    "TRAINING_DATASET",
+    "TEST_DATASET",
+    # Evaluation
+    "evaluate_predictions",
+]