Skip to content

Commit b5b528b

Browse files
authored
Merge pull request #151 from Dooders/dev
Dev
2 parents 251bcb5 + e8b4db2 commit b5b528b

14 files changed

Lines changed: 15917 additions & 1817048 deletions

File tree

converter/config.py

Lines changed: 40 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -4,81 +4,86 @@
44

55
import logging
66
from dataclasses import dataclass, field
7-
from typing import Dict, Optional, List, Union
7+
from typing import Dict, List, Optional
88

9-
from .tiering import TieringStrategy, create_tiering_strategy
109
from .mapping import MemoryTypeMapper
10+
from .tiering import TieringStrategy, create_tiering_strategy
1111

1212
logger = logging.getLogger(__name__)
1313

14+
1415
@dataclass
1516
class ConverterConfig:
1617
"""Configuration for the AgentFarm DB to Memory System converter."""
17-
18+
1819
# Redis configuration
1920
use_mock_redis: bool = True
20-
21+
2122
# Validation settings
2223
validate: bool = True
2324
error_handling: str = "skip" # One of: "skip", "fail", "log"
24-
25+
2526
# Processing settings
2627
batch_size: int = 100
2728
show_progress: bool = True
28-
29+
2930
# Memory type mapping
30-
memory_type_mapping: Dict[str, str] = field(default_factory=lambda: {
31-
'AgentStateModel': 'state',
32-
'ActionModel': 'action',
33-
'SocialInteractionModel': 'interaction'
34-
})
31+
memory_type_mapping: Dict[str, str] = field(
32+
default_factory=lambda: {
33+
"AgentStateModel": "state",
34+
"ActionModel": "action",
35+
"SocialInteractionModel": "interaction",
36+
}
37+
)
3538
memory_type_mapper: Optional[MemoryTypeMapper] = None
36-
39+
3740
# Tiering strategy
38-
tiering_strategy_type: str = "simple" # One of: "simple", "step_based", "importance_aware"
41+
tiering_strategy_type: str = (
42+
"simple" # One of: "simple", "step_based", "importance_aware"
43+
)
3944
tiering_strategy: Optional[TieringStrategy] = None
40-
45+
4146
# Import settings
4247
import_mode: str = "full" # One of: "full", "incremental"
4348
selective_agents: Optional[List[int]] = None
4449
total_steps: Optional[int] = None # Total number of steps in the simulation
45-
50+
4651
def __post_init__(self):
4752
"""Validate configuration after initialization."""
4853
self._validate_error_handling()
4954
self._validate_import_mode()
5055
self._validate_batch_size()
5156
self._validate_tiering_strategy()
52-
57+
5358
# Initialize memory type mapper
5459
if self.memory_type_mapper is None:
5560
self.memory_type_mapper = MemoryTypeMapper(mapping=self.memory_type_mapping)
56-
61+
5762
# Initialize tiering strategy if not provided
5863
if self.tiering_strategy is None:
5964
self.tiering_strategy = create_tiering_strategy(self.tiering_strategy_type)
60-
65+
6166
def _validate_error_handling(self):
6267
"""Validate error handling mode."""
6368
if self.error_handling not in ["skip", "fail", "log"]:
6469
raise ValueError(
6570
f"Invalid error_handling mode: {self.error_handling}. "
6671
"Must be one of: skip, fail, log"
6772
)
68-
73+
6974
def _validate_import_mode(self):
7075
"""Validate import mode."""
7176
if self.import_mode not in ["full", "incremental"]:
7277
raise ValueError(
7378
f"Invalid import_mode: {self.import_mode}. "
7479
"Must be one of: full, incremental"
7580
)
76-
81+
7782
def _validate_batch_size(self):
7883
"""Validate batch size."""
7984
if self.batch_size < 1:
8085
raise ValueError("batch_size must be greater than 0")
81-
86+
8287
def _validate_tiering_strategy(self):
8388
"""Validate tiering strategy settings."""
8489
valid_types = ["simple", "step_based", "importance_aware"]
@@ -88,19 +93,20 @@ def _validate_tiering_strategy(self):
8893
f"Must be one of: {valid_types}"
8994
)
9095

96+
9197
# Default configuration
9298
DEFAULT_CONFIG = {
93-
'use_mock_redis': True,
94-
'validate': True,
95-
'error_handling': 'skip',
96-
'batch_size': 100,
97-
'show_progress': True,
98-
'memory_type_mapping': {
99-
'AgentStateModel': 'state',
100-
'ActionModel': 'action',
101-
'SocialInteractionModel': 'interaction'
99+
"use_mock_redis": True,
100+
"validate": True,
101+
"error_handling": "skip",
102+
"batch_size": 100,
103+
"show_progress": True,
104+
"memory_type_mapping": {
105+
"AgentStateModel": "state",
106+
"ActionModel": "action",
107+
"SocialInteractionModel": "interaction",
102108
},
103-
'tiering_strategy_type': 'simple',
104-
'import_mode': 'full',
105-
'selective_agents': None
106-
}
109+
"tiering_strategy_type": "simple",
110+
"import_mode": "full",
111+
"selective_agents": None,
112+
}

converter/db.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,8 +63,14 @@ def __init__(self, db_path: str, config: ConverterConfig):
6363
def initialize(self) -> None:
6464
"""Initialize the database connection and session factory."""
6565
try:
66+
# Handle in-memory database
67+
if self.db_path == 'sqlite:///:memory:':
68+
engine_url = 'sqlite:///:memory:'
69+
else:
70+
engine_url = f'sqlite:///{self.db_path}'
71+
6672
self._engine = create_engine(
67-
f'sqlite:///{self.db_path}',
73+
engine_url,
6874
poolclass=QueuePool,
6975
pool_size=5,
7076
max_overflow=10,

0 commit comments

Comments
 (0)