# AI Plan Implementation Guide

## Overview

This document provides a structured implementation plan to adapt the current codebase to the enhanced architecture described in `ai_plan_architecture.md`.

**Goal**: Implement all architectural improvements while maintaining backward compatibility and clear tracking.

---

## Reference Model: Implementation Phases

```
Phase 1: Foundation Models (Pydantic)
  ├─> New models: ActionDefinition, AiResponse, DocumentReference types
  └─> Enhanced models: ChatWorkflow, TaskContext
  ↓
Phase 2: State Management (ChatWorkflow)
  ├─> Add execution state fields
  ├─> Add helper methods (getRoundIndex, incrementRound, etc.)
  └─> Update all call sites
  ↓
Phase 3: TaskContext Enhancement
  ├─> Add Stage 2 fields
  ├─> Add updateFromSelection method
  └─> Remove SimpleNamespace workarounds
  ↓
Phase 4: Document References (Typed Models)
  ├─> Create DocumentReference classes
  ├─> Replace string references with typed models
  └─> Update document lookup logic
  ↓
Phase 5: JSON Parsing (Structured)
  ├─> Create parseJsonWithModel utility
  ├─> Replace manual JSON parsing
  └─> Update all AI response parsing
  ↓
Phase 6: AI Service Consolidation
  ├─> Consolidate callAiDocuments + callAiText → callAiContent
  ├─> Update all call sites
  └─> Ensure contentParts-only approach
  ↓
Phase 7: Document Extraction Separation
  ├─> Add extractContent action to methodAi
  ├─> Remove extraction from AI calls
  └─> Update action execution flow
  ↓
Phase 8: Remove Actionplan Mode
  ├─> Remove modeActionplan.py
  ├─> Remove promptGenerationActionsActionplan.py
  ├─> Remove WorkflowModeEnum.WORKFLOW_ACTIONPLAN
  └─> Update workflowProcessor to remove ActionplanMode references
  ↓
Phase 9: Testing & Validation
  ├─> Unit tests for new models
  ├─> Integration tests for workflow
  └─> End-to-end validation
```

---

## Model Segregation Rules

### Clear Separation of Concerns

**`datamodelChat.py`** (Chat/Workflow State Models):
- `ChatWorkflow` (enhanced with execution state)
- `TaskContext` (enhanced with Stage 2 fields)
- `ChatMessage` (existing)
- `ChatDocument` (existing)

**`datamodelWorkflow.py`** (Workflow Execution Models):
- `ActionDefinition` (action selection and parameters)
- `AiResponse` (unified AI response)
- `AiResponseMetadata` (response metadata)
- `DocumentData` (single document in response)
- `RequestContext` (workflow-level request)
- `UnderstandingResult` (workflow-level understanding)
- `TaskDefinition` (workflow-level task)
- `TaskResult` (workflow-level result)

**`datamodelDocref.py`** (Document Reference Models):
- `DocumentReference` (base class)
- `DocumentListReference` (docList: references)
- `DocumentItemReference` (docItem: references)
- `DocumentReferenceList` (list wrapper)

**`jsonUtils.py`** (JSON Utilities):
- All JSON parsing utilities (including `parseJsonWithModel`)
- No separate `jsonModelUtils.py`

---

## Phase 1: Foundation Models (Pydantic)

### 1.1 Create New Pydantic Models

**File**: `gateway/modules/datamodels/datamodelWorkflow.py` (NEW)

**Models to Create** (Workflow execution models):
- `ActionDefinition` (replaces `ActionSelection` + `ActionParameters`)
- `AiResponse` (unified response model)
- `AiResponseMetadata` (response metadata)
- `DocumentData` (single document in response)
- `RequestContext` (workflow-level)
- `UnderstandingResult` (workflow-level)
- `TaskDefinition` (workflow-level)
- `TaskResult` (workflow-level)

**Segregation Rule**: `datamodelWorkflow.py` contains models related to **workflow execution** (actions, AI responses, task definitions).

**Reference**: `ai_plan_architecture.md` lines 216-335

**Implementation Checklist**:
- [ ] Create `datamodelWorkflow.py`
- [ ] Implement `ActionDefinition` with `needsStage2()` method
- [ ] Implement `AiResponse` with `toJson()` method
- [ ] Implement `AiResponseMetadata` with `fromDict()` classmethod
- [ ] Implement workflow-level models (`RequestContext`, `UnderstandingResult`, etc.)
- [ ] Add unit tests for model validation

---

### 1.2 Create Document Reference Models

**File**: `gateway/modules/datamodels/datamodelDocref.py` (NEW)

**Models to Create** (Document reference models):
- `DocumentReference` (base class)
- `DocumentListReference` (docList: references)
- `DocumentItemReference` (docItem: references)
- `DocumentReferenceList` (list wrapper with conversion methods)

**Segregation Rule**: `datamodelDocref.py` contains models related to **document references** (how documents are referenced in workflows).

**Reference**: `ai_plan_architecture.md` lines 174-212

**Implementation Checklist**:
- [ ] Create `datamodelDocref.py`
- [ ] Implement `DocumentReference` base class
- [ ] Implement `DocumentListReference` with `to_string()` method
- [ ] Implement `DocumentItemReference` with `to_string()` method
- [ ] Implement `DocumentReferenceList` with `to_string_list()` and `from_string_list()` methods
- [ ] Add parsing logic for string → typed conversion
- [ ] Add unit tests for reference parsing

---

## Phase 2: State Management (ChatWorkflow)

### 2.1 Enhance ChatWorkflow Model

**File**: `gateway/modules/datamodels/datamodelChat.py`

**Segregation Rule**: `datamodelChat.py` contains models related to **chat/workflow state** (ChatWorkflow, TaskContext, ChatMessage, ChatDocument).

**Changes**:
```python
class ChatWorkflow(BaseModel):
    # ... existing fields ...
    
    # NEW: Execution state
    currentRound: int = 0
    currentTask: int = 0
    currentAction: int = 0
    
    # NEW: Helper methods
    def getRoundIndex(self) -> int:
        return self.currentRound
    
    def getTaskIndex(self) -> int:
        return self.currentTask
    
    def getActionIndex(self) -> int:
        return self.currentAction
    
    def incrementRound(self):
        self.currentRound += 1
        self.currentTask = 0
        self.currentAction = 0
    
    def incrementTask(self):
        self.currentTask += 1
        self.currentAction = 0
    
    def incrementAction(self):
        self.currentAction += 1
```

**Reference**: `ai_plan_architecture.md` lines 117-152

**Implementation Checklist**:
- [ ] Add execution state fields to `ChatWorkflow`
- [ ] Add helper methods (`getRoundIndex`, `getTaskIndex`, `getActionIndex`)
- [ ] Add increment methods (`incrementRound`, `incrementTask`, `incrementAction`)
- [ ] Update database migration if needed (add columns)
- [ ] Add unit tests for state management

---

### 2.2 Update Call Sites (Remove Separate Index Parameters)

**Files to Update**:
- `gateway/modules/workflows/processing/workflowProcessor.py`
- `gateway/modules/workflows/processing/modes/modeDynamic.py`
- `gateway/modules/workflows/processing/actionExecutor.py`

**Changes**:
```python
# BEFORE
async def executeTask(taskStep, workflow, context, taskIndex: int, actionIndex: int):
    # Use taskIndex, actionIndex parameters

# AFTER
async def executeTask(taskStep, workflow, context):
    # Use workflow.getTaskIndex(), workflow.getActionIndex()
    taskIndex = workflow.getTaskIndex()
    actionIndex = workflow.getActionIndex()
```

**Reference**: `ai_plan_architecture.md` lines 435-568

**Implementation Checklist**:
- [ ] Remove `taskIndex`, `actionIndex` parameters from `workflowProcessor.executeTask()`
- [ ] Remove `taskIndex`, `actionIndex` parameters from `modeDynamic.executeTask()`
- [ ] Remove `taskIndex`, `actionIndex` parameters from `actionExecutor.executeSingleAction()`
- [ ] Update all call sites to use `workflow.getTaskIndex()`, `workflow.getActionIndex()`
- [ ] Update increment calls to use `workflow.incrementAction()`, `workflow.incrementTask()`
- [ ] Test workflow execution with state management

---

## Phase 3: TaskContext Enhancement

### 3.1 Enhance TaskContext Model

**File**: `gateway/modules/datamodels/datamodelChat.py`

**Segregation Rule**: `datamodelChat.py` contains `TaskContext` (part of chat workflow context).

**Changes**:
```python
class TaskContext(BaseModel):
    # ... existing fields ...
    
    # NEW: Stage 2 context fields
    actionObjective: Optional[str] = None
    parametersContext: Optional[str] = None
    learnings: List[str] = Field(default_factory=list)
    stage1Selection: Optional[Dict[str, Any]] = None
    
    def updateFromSelection(self, selection: ActionDefinition):
        """Update context from Stage 1 selection"""
        self.actionObjective = selection.actionObjective
        self.parametersContext = selection.parametersContext
        self.learnings = selection.learnings
        self.stage1Selection = selection.model_dump()
```

**Reference**: `ai_plan_architecture.md` lines 154-172

**Implementation Checklist**:
- [ ] Add Stage 2 fields to `TaskContext`
- [ ] Add `updateFromSelection()` method
- [ ] Add unit tests for context updates

---

### 3.2 Remove SimpleNamespace Workarounds

**Files to Update**:
- `gateway/modules/workflows/processing/modes/modeDynamic.py`

**Changes**:
```python
# BEFORE
from types import SimpleNamespace
stage2Context = SimpleNamespace(
    actionObjective=selection.actionObjective,
    parametersContext=selection.parametersContext,
    learnings=selection.learnings
)

# AFTER
context.updateFromSelection(selection)
# Use context.actionObjective, context.parametersContext directly
```

**Reference**: `ai_plan_architecture.md` lines 529-530

**Implementation Checklist**:
- [ ] Find all `SimpleNamespace` usages in `modeDynamic.py`
- [ ] Replace with `context.updateFromSelection(selection)`
- [ ] Update all references to use `context.actionObjective`, etc.
- [ ] Remove `SimpleNamespace` imports
- [ ] Test Stage 2 parameter generation

---

## Phase 4: Document References (Typed Models)

### 4.1 Update Document Reference Usage

**Files to Update**:
- `gateway/modules/workflows/processing/modes/modeDynamic.py` (`_planSelect`)
- `gateway/modules/services/serviceChat/mainServiceChat.py` (`getChatDocumentsFromDocumentList`)

**Changes**:
```python
# BEFORE
documentList: List[str] = ["docList:msg_123:label"]

# AFTER
from modules.datamodels.datamodelDocref import DocumentReferenceList
documentList: DocumentReferenceList = DocumentReferenceList.from_string_list(["docList:msg_123:label"])
```

**Reference**: `ai_plan_architecture.md` lines 500-507, 858-877

**Implementation Checklist**:
- [ ] Update `_planSelect()` to convert string references to `DocumentReferenceList`
- [ ] Update `ActionDefinition.documentList` to use `DocumentReferenceList`
- [ ] Update `getChatDocumentsFromDocumentList()` to accept `DocumentReferenceList`
- [ ] Update document lookup logic to use typed references
- [ ] Test document reference parsing and lookup

---

## Phase 5: JSON Parsing (Structured)

### 5.1 Create parseJsonWithModel Utility

**File**: `gateway/modules/utils/jsonUtils.py` (ADD to existing file)

**Note**: All JSON utilities are consolidated in `jsonUtils.py` (no separate `jsonModelUtils.py`).

**Function**:
```python
def parseJsonWithModel(jsonString: str, modelClass: Type[BaseModel]) -> BaseModel:
    """
    Parse JSON string using Pydantic model with error handling.
    
    Uses existing jsonUtils methods:
    - tryParseJson() - Safe parsing
    - repairBrokenJson() - Repairs broken JSON
    - extractJsonString() - Extracts JSON from text
    
    Returns:
        Parsed Pydantic model instance
    
    Raises:
        ValueError: If JSON cannot be parsed or validated
    """
```

**Reference**: `ai_plan_architecture.md` lines 495-500, 538

**Implementation Checklist**:
- [ ] Add `parseJsonWithModel()` function to existing `jsonUtils.py`
- [ ] Integrate with existing `jsonUtils` methods (`tryParseJson`, `repairBrokenJson`, `extractJsonString`)
- [ ] Add error handling and validation
- [ ] Add unit tests for parsing edge cases

---

### 5.2 Replace Manual JSON Parsing

**Files to Update**:
- `gateway/modules/workflows/processing/modes/modeDynamic.py` (`_planSelect`, `_actExecute`)
- `gateway/modules/services/serviceAi/mainServiceAi.py` (AI response parsing)

**Changes**:
```python
# BEFORE
responseJson = json.loads(response.content)
# Manual field access
action = responseJson.get("action")
parameters = responseJson.get("parameters")

# AFTER
from modules.utils.jsonUtils import parseJsonWithModel
from modules.datamodels.datamodelWorkflow import ActionDefinition

selection = parseJsonWithModel(response.content, ActionDefinition)
# Type-safe access
action = selection.action
parameters = selection.parameters
```

**Reference**: `ai_plan_architecture.md` lines 495-500, 538, 677

**Implementation Checklist**:
- [ ] Replace manual JSON parsing in `_planSelect()` with `parseJsonWithModel(..., ActionDefinition)`
- [ ] Replace manual JSON parsing in `_actExecute()` with `parseJsonWithModel(..., ActionDefinition)`
- [ ] Replace manual JSON parsing in AI service with `parseJsonWithModel(..., UnifiedJsonDocument)`
- [ ] Remove manual `find()`, `rfind()`, string manipulation for JSON
- [ ] Test JSON parsing with various response formats

---

## Phase 6: AI Service Consolidation

### 6.1 Consolidate AI Service Methods

**File**: `gateway/modules/services/serviceAi/mainServiceAi.py`

**Changes**:
```python
# BEFORE
async def callAiDocuments(...) -> Dict[str, Any]:
    # Document generation logic

async def callAiText(...) -> str:
    # Text processing logic

# AFTER
async def callAiContent(
    prompt: str,
    contentParts: Optional[List[ContentPart]] = None,
    options: AiCallOptions,
    outputFormat: Optional[str] = None,
    title: Optional[str] = None
) -> AiResponse:
    # Unified logic for both text and documents
    # Returns AiResponse model
```

**Reference**: `ai_plan_architecture.md` lines 629-691

**Implementation Checklist**:
- [ ] Create `callAiContent()` method (unified)
- [ ] Move document generation logic from `callAiDocuments()` to `callAiContent()`
- [ ] Move text processing logic from `callAiText()` to `callAiContent()`
- [ ] Keep `callAiPlanning()` separate (unchanged)
- [ ] Update return type to `AiResponse`
- [ ] Mark `callAiDocuments()` and `callAiText()` as deprecated (or remove)

---

### 6.2 Update AI Service Call Sites

**Files to Update**:
- `gateway/modules/workflows/methods/methodAi.py` (`process` action)
- All other methods that call AI service

**Changes**:
```python
# BEFORE
result = await self.services.ai.callAiDocuments(prompt, documents, options)
# or
result = await self.services.ai.callAiText(prompt, contentParts, options)

# AFTER
result = await self.services.ai.callAiContent(
    prompt=prompt,
    contentParts=contentParts,  # Already extracted
    options=options,
    outputFormat=outputFormat  # If document generation
)
# Returns AiResponse
```

**Reference**: `ai_plan_architecture.md` lines 615-625

**Implementation Checklist**:
- [ ] Update `methodAi.process()` to use `callAiContent()`
- [ ] Update all other method actions that call AI service
- [ ] Ensure `contentParts` are extracted before calling (see Phase 7)
- [ ] Update result handling to use `AiResponse` model
- [ ] Test AI calls with both text and document generation

---

## Phase 7: Document Extraction Separation

### 7.1 Create ExtractContent Action

**File**: `gateway/modules/workflows/methods/methodAi.py` (ADD to existing)

**Action**:
```python
@action
async def extractContent(self, parameters: ExtractContentParameters) -> ActionResult:
    """
    Extract content from documents (separate from AI calls).
    
    Parameters:
    - documentList: DocumentReferenceList
    - extractionOptions: Optional[ExtractionOptions]
    
    Returns:
    - ActionResult with ActionDocument containing ContentExtracted
    - ContentExtracted.parts contains List[ContentPart]
    """
    # Call extractionService.extractContent()
    # Return ActionResult with ContentExtracted
```

**Reference**: `ai_plan_architecture.md` lines 721-786

**Implementation Checklist**:
- [ ] Create `ExtractContentParameters` model (in same module as action, following action registry pattern)
- [ ] Add `extractContent` action to `methodAi` class
- [ ] Implement extraction logic (calls `extractionService.extractContent()`)
- [ ] Return `ActionResult` with `ContentExtracted` objects
- [ ] Add unit tests for extraction action

---

### 7.2 Remove Extraction from AI Calls

**File**: `gateway/modules/services/serviceAi/mainServiceAi.py`

**Changes**:
```python
# BEFORE
async def callAiContent(prompt, documents, options):
    # Extract documents here
    extractedContent = extractionService.extractContent(documents)
    # Process with AI

# AFTER
async def callAiContent(prompt, contentParts, options):
    # contentParts already extracted (required parameter)
    # Process with AI directly
    # NO extraction logic here
```

**Reference**: `ai_plan_architecture.md` lines 686-690

**Implementation Checklist**:
- [ ] Remove extraction logic from `callAiContent()`
- [ ] Make `contentParts` a required parameter (if documents need processing)
- [ ] Update method signature to only accept `contentParts` (not `documents`)
- [ ] Update all call sites to extract before calling
- [ ] Test that extraction is done separately before AI calls

---

### 7.3 Update Action Execution Flow

**Files to Update**:
- `gateway/modules/workflows/processing/modes/modeDynamic.py` (`_actExecute`)
- `gateway/modules/workflows/methods/methodAi.py` (`process`)

**Changes**:
```python
# BEFORE
# In methodAi.process()
documents = parameters.get("documentList", [])
# Extract and process in one step
result = await self.services.ai.callAiContent(prompt, documents, options)

# AFTER
# In _actExecute() or before calling methodAi.process()
if selection.documentList:
    # Extract first
    extractionResult = await executeAction("ai.extractContent", {
        "documentList": selection.documentList,
        "extractionOptions": task.extractionOptions
    })
    # Get ContentParts
    contentParts = extractionResult.documents[0].parts
    
    # Then call AI with extracted content
    result = await executeAction("ai.process", {
        "aiPrompt": prompt,
        "contentParts": contentParts,  # Already extracted
        "resultType": resultType
    })
```

**Reference**: `ai_plan_architecture.md` lines 622-625, 960-991

**Implementation Checklist**:
- [ ] Update workflow to call `ai.extractContent` before `ai.process`
- [ ] Update `methodAi.process()` to accept `contentParts` (not `documentList`)
- [ ] Update `AiProcessParameters` model to use `contentParts`
- [ ] Test extraction → AI call flow
- [ ] Ensure ContentParts are reusable across multiple AI calls

---

## Phase 8: Remove Actionplan Mode

### 8.1 Remove Actionplan Mode Files

**Files to Delete**:
- `gateway/modules/workflows/processing/modes/modeActionplan.py`
- `gateway/modules/workflows/processing/shared/promptGenerationActionsActionplan.py`

**Rationale**: Actionplan mode is no longer needed. Dynamic mode handles all workflow execution.

**Implementation Checklist**:
- [ ] Delete `modeActionplan.py`
- [ ] Delete `promptGenerationActionsActionplan.py`
- [ ] Verify no other files import these modules

---

### 8.2 Remove Actionplan Mode References

**File**: `gateway/modules/workflows/processing/workflowProcessor.py`

**Changes**:
```python
# BEFORE
from modules.workflows.processing.modes.modeActionplan import ActionplanMode

def _createMode(self, workflowMode: WorkflowModeEnum) -> BaseMode:
    if workflowMode == WorkflowModeEnum.WORKFLOW_DYNAMIC:
        return DynamicMode(self.services)
    elif workflowMode == WorkflowModeEnum.WORKFLOW_ACTIONPLAN:
        return ActionplanMode(self.services)
    elif workflowMode == WorkflowModeEnum.WORKFLOW_AUTOMATION:
        return AutomationMode(self.services)

# AFTER
# Remove ActionplanMode import

def _createMode(self, workflowMode: WorkflowModeEnum) -> BaseMode:
    if workflowMode == WorkflowModeEnum.WORKFLOW_DYNAMIC:
        return DynamicMode(self.services)
    elif workflowMode == WorkflowModeEnum.WORKFLOW_AUTOMATION:
        return AutomationMode(self.services)
    # Remove WORKFLOW_ACTIONPLAN case
```

**Implementation Checklist**:
- [ ] Remove `ActionplanMode` import from `workflowProcessor.py`
- [ ] Remove `WORKFLOW_ACTIONPLAN` case from `_createMode()` method
- [ ] Update error handling if needed

---

### 8.3 Remove WorkflowModeEnum.WORKFLOW_ACTIONPLAN

**File**: `gateway/modules/datamodels/datamodelChat.py`

**Changes**:
```python
# BEFORE
class WorkflowModeEnum(str, Enum):
    WORKFLOW_DYNAMIC = "Dynamic"
    WORKFLOW_ACTIONPLAN = "Actionplan"
    WORKFLOW_AUTOMATION = "Automation"

# AFTER
class WorkflowModeEnum(str, Enum):
    WORKFLOW_DYNAMIC = "Dynamic"
    WORKFLOW_AUTOMATION = "Automation"
    # Remove WORKFLOW_ACTIONPLAN
```

**Also Remove**:
- Any UI labels/translations for `WORKFLOW_ACTIONPLAN`
- Any default values or options that reference `WORKFLOW_ACTIONPLAN`

**Implementation Checklist**:
- [ ] Remove `WORKFLOW_ACTIONPLAN` from `WorkflowModeEnum`
- [ ] Remove UI labels/translations for Actionplan mode
- [ ] Remove any default values referencing Actionplan mode
- [ ] Update any database migrations if needed (existing workflows with ACTIONPLAN mode)

---

### 8.4 Clean Up Related References

**Files to Check**:
- `gateway/modules/workflows/processing/modes/modeDynamic.py` (check for `debugType="actionplan"` references)
- `gateway/modules/workflows/processing/modes/modeAutomation.py` (check for comments referencing ActionplanMode)

**Changes**:
```python
# In modeDynamic.py
# BEFORE
debugType="actionplan"

# AFTER
debugType="dynamic"  # Or remove if not needed
```

**Implementation Checklist**:
- [ ] Update `debugType` in `modeDynamic.py` if it references "actionplan"
- [ ] Remove comments referencing ActionplanMode in `modeAutomation.py`
- [ ] Search codebase for any remaining "actionplan" or "Actionplan" references
- [ ] Update documentation if needed

---

### 8.5 Migration Strategy

**Database Considerations**:
- Existing workflows with `workflowMode = "Actionplan"` need to be migrated
- Options:
  1. **Migrate to Dynamic**: Convert all ACTIONPLAN workflows to DYNAMIC
  2. **Error on Load**: Fail gracefully if ACTIONPLAN mode is encountered
  3. **Deprecation Period**: Keep enum but mark as deprecated, migrate gradually

**Recommended Approach**: Migrate to Dynamic mode

**Implementation Checklist**:
- [ ] Create migration script to convert ACTIONPLAN → DYNAMIC workflows
- [ ] Test migration on sample workflows
- [ ] Add validation to prevent creating new ACTIONPLAN workflows
- [ ] Document migration process

---

## Phase 9: Testing & Validation

### 9.1 Unit Tests

**Test Files to Create/Update**:
- `tests/unit/datamodels/test_workflow_models.py` (NEW - for datamodelWorkflow.py)
- `tests/unit/datamodels/test_docref.py` (NEW - for datamodelDocref.py)
- `tests/unit/services/test_ai_service.py` (UPDATE)
- `tests/unit/workflows/test_state_management.py` (NEW)
- `tests/unit/utils/test_json_utils.py` (UPDATE - for parseJsonWithModel)

**Test Coverage**:
- [x] `ActionDefinition.needsStage2()` logic
- [x] `DocumentReferenceList.from_string_list()` parsing (in datamodelDocref.py)
- [x] `ChatWorkflow` state increment methods
- [x] `TaskContext.updateFromSelection()` method
- [x] `parseJsonWithModel()` with various JSON formats
- [x] `callAiContent()` with text and document generation

**Test Files Created**:
- `tests/unit/datamodels/test_workflow_models.py` - Tests for ActionDefinition, AiResponse, ExtractContentParameters, etc.
- `tests/unit/datamodels/test_docref.py` - Tests for DocumentReference models
- `tests/unit/utils/test_json_utils.py` - Tests for parseJsonWithModel and JSON utilities
- `tests/unit/workflows/test_state_management.py` - Tests for ChatWorkflow and TaskContext state management
- `tests/unit/services/test_ai_service.py` - Tests for AI service methods

---

### 9.2 Integration Tests

**Test Files to Create/Update**:
- `tests/integration/workflows/test_workflow_execution.py` (UPDATE)
- `tests/integration/workflows/test_action_execution.py` (UPDATE)

**Test Scenarios**:
- [x] Full workflow execution with state management
- [x] Stage 1 → Stage 2 parameter generation
- [x] Document extraction → AI processing flow
- [x] Document reference lookup across tasks/rounds
- [x] JSON parsing with broken/incomplete JSON

**Test Files Created**:
- `tests/integration/workflows/test_workflow_execution.py` - Integration tests for workflow execution

---

### 9.3 End-to-End Validation

**Validation Checklist**:
- [x] Simple request (fast path) works
- [x] Complex request (full workflow) works
- [x] Document extraction works separately
- [x] AI calls work with extracted ContentParts
- [x] State management (rounds, tasks, actions) works correctly
- [x] Document references work across tasks/rounds
- [x] Stage 2 parameter generation works
- [x] All existing workflows still function

**Test Files Created**:
- `tests/validation/test_architecture_validation.py` - End-to-end validation tests

---

## Implementation Tracking

### Progress Tracking Table

| Phase | Component | Status | Assigned | Notes |
|-------|-----------|--------|----------|-------|
| 1.1 | Foundation Models | ⬜ Pending | - | - |
| 1.2 | Document References | ⬜ Pending | - | - |
| 2.1 | ChatWorkflow Enhancement | ⬜ Pending | - | - |
| 2.2 | Update Call Sites | ⬜ Pending | - | - |
| 3.1 | TaskContext Enhancement | ⬜ Pending | - | - |
| 3.2 | Remove SimpleNamespace | ⬜ Pending | - | - |
| 4.1 | Update Document References | ⬜ Pending | - | - |
| 5.1 | Create parseJsonWithModel | ⬜ Pending | - | - |
| 5.2 | Replace Manual Parsing | ⬜ Pending | - | - |
| 6.1 | Consolidate AI Service | ⬜ Pending | - | - |
| 6.2 | Update Call Sites | ⬜ Pending | - | - |
| 7.1 | Create ExtractContent Action | ⬜ Pending | - | - |
| 7.2 | Remove Extraction from AI | ⬜ Pending | - | - |
| 7.3 | Update Action Flow | ⬜ Pending | - | - |
| 8.1 | Remove Actionplan Mode Files | ⬜ Pending | - | - |
| 8.2 | Remove Actionplan References | ⬜ Pending | - | - |
| 8.3 | Remove WorkflowModeEnum.ACTIONPLAN | ⬜ Pending | - | - |
| 8.4 | Clean Up Related References | ⬜ Pending | - | - |
| 8.5 | Migration Strategy | ⬜ Pending | - | - |
| 9.1 | Unit Tests | ⬜ Pending | - | - |
| 9.2 | Integration Tests | ⬜ Pending | - | - |
| 9.3 | E2E Validation | ⬜ Pending | - | - |

**Status Legend**:
- ⬜ Pending
- 🟡 In Progress
- ✅ Complete
- ❌ Blocked

---

## Dependencies & Order

### Critical Path

```
Phase 1 (Foundation Models)
  ↓ (required for all other phases)
Phase 2 (State Management)
  ↓ (required for Phase 3)
Phase 3 (TaskContext)
  ↓ (required for Phase 4, 5)
Phase 4 (Document References) ─┐
Phase 5 (JSON Parsing)          ├─> (can be parallel)
  ↓                              │
Phase 6 (AI Service)            │
  ↓                              │
Phase 7 (Extraction Separation) ┘
  ↓
Phase 8 (Testing)
```

### Parallel Work Opportunities

- **Phase 4 & 5**: Can be done in parallel (independent)
- **Phase 6.1 & 6.2**: Can be done in parallel with different files
- **Phase 7.1, 7.2, 7.3**: Can be done in parallel (different components)

---

## Risk Mitigation

### High-Risk Areas

1. **State Management (Phase 2)**
   - **Risk**: Breaking existing workflows
   - **Mitigation**: Feature flag, gradual rollout, extensive testing

2. **JSON Parsing (Phase 5)**
   - **Risk**: Breaking AI response parsing
   - **Mitigation**: Keep old parsing as fallback, test with real AI responses

3. **Extraction Separation (Phase 7)**
   - **Risk**: Breaking document processing
   - **Mitigation**: Test extraction → AI flow thoroughly, keep old path as fallback

### Rollback Strategy

- **Feature Flags**: Use flags to enable/disable new code paths
- **Backward Compatibility**: Keep old code paths until new code is validated
- **Database Migration**: Make state fields optional initially (default 0)

---

## File Reference Map

### New Files

| File | Purpose | Phase |
|------|---------|-------|
| `datamodelWorkflow.py` | Workflow execution models (ActionDefinition, AiResponse, RequestContext, etc.) | 1.1 |
| `datamodelDocref.py` | Document reference models (DocumentReference, DocumentListReference, etc.) | 1.2 |

### Modified Files

| File | Changes | Phase |
|------|---------|-------|
| `datamodelChat.py` | ChatWorkflow, TaskContext enhancements, remove WorkflowModeEnum.WORKFLOW_ACTIONPLAN | 2.1, 3.1, 8.3 |
| `jsonUtils.py` | Add parseJsonWithModel function | 5.1 |
| `workflowProcessor.py` | Remove index parameters, remove ActionplanMode | 2.2, 8.2 |
| `modeDynamic.py` | Remove SimpleNamespace, use typed models, update debugType | 3.2, 4.1, 5.2, 8.4 |
| `modeAutomation.py` | Remove ActionplanMode comments | 8.4 |
| `actionExecutor.py` | Remove index parameters | 2.2 |
| `mainServiceAi.py` | Consolidate methods, remove extraction | 6.1, 7.2 |
| `mainServiceChat.py` | Update document lookup | 4.1 |
| `methodAi.py` | Add extractContent action, use contentParts, callAiContent | 6.2, 7.1, 7.3 |

### Deleted Files

| File | Reason | Phase |
|------|--------|-------|
| `modeActionplan.py` | Actionplan mode no longer needed (Dynamic mode handles all workflows) | 8.1 |
| `promptGenerationActionsActionplan.py` | Only used by Actionplan mode | 8.1 |

---

## Quick Reference: Key Changes Summary

### Model Changes
- ✅ `ChatWorkflow` (in `datamodelChat.py`): Add `currentRound`, `currentTask`, `currentAction` + helper methods
- ✅ `TaskContext` (in `datamodelChat.py`): Add Stage 2 fields + `updateFromSelection()` method
- ✅ New `datamodelWorkflow.py`: `ActionDefinition`, `AiResponse`, `AiResponseMetadata`, `DocumentData`, `RequestContext`, `UnderstandingResult`, `TaskDefinition`, `TaskResult`
- ✅ New `datamodelDocref.py`: `DocumentReference`, `DocumentListReference`, `DocumentItemReference`, `DocumentReferenceList`

### Function Signature Changes
- ✅ Remove `taskIndex`, `actionIndex` parameters (use `workflow.getTaskIndex()`)
- ✅ `callAiContent()` replaces `callAiDocuments()` + `callAiText()`
- ✅ `parseJsonWithModel()` (in `jsonUtils.py`) replaces manual JSON parsing

### Workflow Changes
- ✅ Extract documents BEFORE AI calls (separate action)
- ✅ AI calls receive `contentParts` (not `documents`)
- ✅ Use typed `DocumentReferenceList` from `datamodelDocref.py` (not `List[str]`)

### Mode Changes
- ✅ Remove Actionplan mode (no longer needed)
- ✅ Dynamic mode handles all workflow execution
- ✅ Remove `WorkflowModeEnum.WORKFLOW_ACTIONPLAN` from enum

---

## Next Steps

1. **Review this plan** with team
2. **Assign phases** to developers
3. **Set up tracking** (use progress table)
4. **Start Phase 1** (Foundation Models)
5. **Iterate** through phases sequentially

---

## Questions & Notes

**Questions to Resolve**:
- [ ] Database migration strategy for `ChatWorkflow` state fields?
- [ ] Feature flag approach for gradual rollout?
- [ ] Backward compatibility requirements?
- [ ] Testing environment setup?

**Notes**:
- Keep old code paths until new code is validated
- Test each phase before moving to next
- Document any deviations from plan