# TealFlow MCP: teal.transform Documentation Improvement Plan

**Date:** 2026-02-27
**Status:** Ready for Implementation
**Recommended Approach:** Option 2 (Moderate - Enhanced Agent Guidance + Tool Improvements)

---

## Executive Summary

### Problem
LLMs using TealFlowMCP frequently confuse `teal.transform::variable_choices()` and `teal.transform::value_choices()` functions, leading to incorrect R code generation. The codebase currently has minimal documentation explaining when to use each function.

### Root Cause
- Only 1 usage example exists in agent guidance (line 562 of agent.md)
- No formal documentation explaining function purposes or differences
- Module parameter metadata shows usage patterns but lacks explanatory guidance
- Code generation includes commented examples but no active documentation

### Critical Distinction

**`variable_choices()`** - "Which COLUMN do you want?"
- Extracts variable/column names from a dataset
- Used when users need to SELECT which COLUMNS to analyze
- Example: `variable_choices(ADSL, c("AGE", "SEX", "RACE"))`

**`value_choices()`** - "Which VALUE in this column?"
- Extracts unique VALUES from specific columns with labels
- Used when users need to FILTER data by specific values within a column
- Example: `value_choices(ADLB, "PARAMCD", "PARAM")` → shows "GLUC - Glucose"

### Recommended Solution: Option 2

**Why Option 2?**
- Best balance of effort (1 day) vs impact (High)
- Provides context-aware guidance through enhanced tools
- Low risk, manageable maintenance
- Can evolve to more comprehensive solution later if needed

**What's Included:**
1. Comprehensive documentation section in agent.md (~150-200 lines)
2. Enhanced module_details tool with parameter role classification
3. Improved code generation comments with inline guidance

---

## Implementation Plan

### Phase 1: Enhanced Agent Guidance Documentation (Priority: HIGH)

**File:** `/home/kuba/projects/TealFlowMCP/knowledge_base/agent.md`
**Location:** After "R Code Style Guidelines" section (around line 465)
**Estimated Time:** 2-3 hours

#### Content Structure

```markdown
## Understanding teal.transform Helper Functions

### Overview
[Brief intro to teal.transform's role in parameter construction]

### The Three Core Functions

#### 1. variable_choices()
- Function signature: variable_choices(data, subset = NULL, fill = FALSE, key = NULL)
- Purpose: Extract column names from datasets
- When to use: Column selection scenarios
- Parameters explained
- 3-4 examples with different patterns

#### 2. value_choices()
- Function signature: value_choices(data, var_choices, var_label = NULL, subset = NULL, sep = " - ")
- Purpose: Extract values from columns with labels
- When to use: Filtering scenarios
- Parameters explained
- 3-4 examples with ADLB/paramcd patterns

#### 3. choices_selected()
- Function signature: choices_selected(choices, selected, keep_order, fixed)
- Purpose: Wrapper that pairs choices with preselection
- How it works with the other two functions
- Examples of common combinations

### Decision Guide: Which Function to Use?

[Decision table showing scenarios and correct function choice]

### Common Patterns with ADaM Datasets

#### Pattern 1: Variable Selection (arm_var, strata_var)
[Example with explanation]

#### Pattern 2: Parameter Filtering (paramcd with ADLB)
[Example with explanation]

#### Pattern 3: Fixed Variable Selection (aval_var, cnsr_var)
[Example with explanation]

#### Pattern 4: Dynamic Column Selection
[Example with explanation]

#### Pattern 5: Multi-variable Combinations
[Example with explanation]

### Common Mistakes and How to Avoid Them

[4-5 common error patterns with corrections]

### References
[Links to official teal.transform documentation]
```

#### Key Examples to Include

**Example 1: Arm Variable Selection**
```r
# Use variable_choices() to let user select which column contains arm info
arm_var = teal.transform::choices_selected(
  teal.transform::variable_choices(ADSL, c("ARM", "ARMCD", "ACTARM")),
  selected = "ARM",
  fixed = FALSE
)
```

**Example 2: Parameter Code Filtering**
```r
# Use value_choices() to filter by specific parameter values
paramcd = teal.transform::choices_selected(
  teal.transform::value_choices(ADLB, "PARAMCD", "PARAM"),
  selected = "GLUC",
  fixed = FALSE
)
# User sees: "GLUC - Glucose (Fasting)" in dropdown
```

**Example 3: Fixed Analysis Variable**
```r
# Use variable_choices() with fixed=TRUE for required columns
aval_var = teal.transform::choices_selected(
  teal.transform::variable_choices(dataname, "AVAL"),
  selected = "AVAL",
  fixed = TRUE  # User cannot change this
)
```

**Example 4: Visit Filtering**
```r
# Use value_choices() for visit filtering
avisit = teal.transform::choices_selected(
  teal.transform::value_choices(ADLB, "AVISIT"),
  selected = c("BASELINE", "WEEK 8", "WEEK 16")
)
```

**Example 5: Dynamic Numeric Variable Selection**
```r
# Use variable_choices() without subset for flexible selection
x_var = teal.transform::variable_choices(ADSL)
# User can select from all columns in ADSL
```

#### Decision Table

| Scenario | Function | Example Use Case |
|----------|----------|------------------|
| User selects which COLUMN to analyze | `variable_choices()` | arm_var, strata_var, x_var, y_var |
| User filters by VALUES in a column | `value_choices()` | paramcd="GLUC", avisit="BASELINE" |
| Fixed to specific column name | `variable_choices(data, "COL")` + fixed=TRUE | aval_var="AVAL" |
| Dynamic column selection | `variable_choices(data)` | User picks any column |
| Filter with value labels | `value_choices(data, "VAR", "LABEL")` | Show "GLUC - Glucose" |
| Multiple column options | `variable_choices(data, c("A", "B"))` | arm_var with ARM/ARMCD choices |

---

### Phase 2: Enhanced Module Details Tool (Priority: MEDIUM)

**File:** `/home/kuba/projects/TealFlowMCP/tealflow_mcp/tools/module_details.py`
**Lines to Modify:** 89-145 (parameter documentation formatting)
**Estimated Time:** 3-4 hours

#### Implementation Steps

1. **Add Parameter Classification Logic**
   - Create helper function to classify parameter types
   - Categories: "variable_selection", "value_filtering", "configuration", "statistical"

2. **Enhance Markdown Output**
   - Add "Recommended Function" hint for choices_selected parameters
   - Add inline code examples for key parameters
   - Cross-reference agent.md section

3. **Enhance JSON Output**
   - Add `"parameter_role"` field to parameter objects
   - Add `"recommended_function"` field where applicable
   - Add `"example_code"` snippet for common patterns

#### Parameter Classification Rules

```python
def classify_parameter_role(param_name: str, param_type: str, description: str) -> str:
    """Classify parameter to determine which teal.transform function to use."""

    # Variable selection patterns (use variable_choices)
    variable_selection_params = {
        "arm_var", "strata_var", "facet_var", "id_var",
        "x_var", "y_var", "color_var", "fill_var",
        "aval_var", "cnsr_var", "time_unit_var", "baseline_var"
    }

    # Value filtering patterns (use value_choices)
    value_filtering_params = {
        "paramcd", "avisit", "visit_var"
    }

    # Check for filtering keywords in description
    if "filter" in description.lower() and param_type == "choices_selected":
        return "value_filtering"

    # Check for variable/column selection keywords
    if any(keyword in description.lower() for keyword in ["variable names", "column", "select variable"]):
        return "variable_selection"

    # Direct pattern matching
    if param_name in variable_selection_params:
        return "variable_selection"
    if param_name in value_filtering_params:
        return "value_filtering"

    return "configuration"
```

#### Example Enhanced Output

**Current markdown output:**
```
arm_var (choices_selected): Variable names that can be used as arm_var. Object with all available choices and preselected option for arm variable
```

**Enhanced markdown output:**
```
arm_var (choices_selected): Variable names that can be used as arm_var. Object with all available choices and preselected option for arm variable

  Recommended: Use variable_choices() to provide column selection
  Example:
    arm_var = choices_selected(
      variable_choices(ADSL, c("ARM", "ARMCD")),
      selected = "ARM"
    )

  See agent.md: "Understanding teal.transform Helper Functions" for details
```

---

### Phase 3: Improved Code Generation Comments (Priority: MEDIUM)

**File:** `/home/kuba/projects/TealFlowMCP/tealflow_mcp/tools/code_generation.py`
**Lines to Modify:** 55-75 (clinical modules), 96-125 (general modules)
**Estimated Time:** 2 hours

#### Changes for Clinical Modules

**Current (line 55):**
```python
'      choices = variable_choices("ADSL"),  # TODO: Specify available columns'
```

**Enhanced:**
```python
'      # Use variable_choices() to let users select which COLUMN to analyze'
'      # Example: choices = variable_choices(ADSL, c("ARM", "ARMCD"))'
'      choices = variable_choices("ADSL"),  # TODO: Specify column names'
```

**Add context-aware comments based on parameter names:**
```python
def generate_parameter_comment(param_name: str) -> str:
    """Generate helpful comment based on parameter type."""

    if param_name == "paramcd":
        return "# Use value_choices() to FILTER by specific parameter values (e.g., 'GLUC')"
    elif param_name in ["arm_var", "strata_var", "facet_var"]:
        return "# Use variable_choices() to let user SELECT which COLUMN contains this info"
    elif param_name in ["aval_var", "cnsr_var"]:
        return "# Use variable_choices() with fixed=TRUE for required analysis columns"
    else:
        return "# See tealflow_get_module_details for parameter guidance"
```

#### Changes for General Modules

**Enhance existing commented examples (lines 96-122):**

**Current:**
```python
lines.append("# Example with actual configuration for ADSL dataset:")
lines.append("# x = data_extract_spec(")
lines.append('#   choices = variable_choices(ADSL, c("AGE", "BMRKR1", "BMRKR2")),')
```

**Enhanced:**
```python
lines.append("# Example 1: Column selection with variable_choices()")
lines.append("#   Use variable_choices() when user needs to SELECT which COLUMN")
lines.append("# x = data_extract_spec(")
lines.append('#   dataname = "ADSL",')
lines.append("#   select = select_spec(")
lines.append('#     choices = variable_choices(ADSL, c("AGE", "BMRKR1", "BMRKR2")),')
lines.append('#     selected = "AGE"')
lines.append("#   )")
lines.append("# )")
lines.append("")
lines.append("# Example 2: Value filtering with value_choices()")
lines.append("#   Use value_choices() when user needs to FILTER by VALUES in a column")
lines.append("# y = data_extract_spec(")
lines.append('#   dataname = "ADLB",')
lines.append("#   filter = filter_spec(")
lines.append('#     vars = "PARAMCD",  # Column to filter on')
lines.append('#     choices = value_choices(ADLB, "PARAMCD", "PARAM"),  # Values with labels')
lines.append('#     selected = "GLUC"  # Pre-selected value')
lines.append("#   )")
lines.append("# )")
```

---

## Testing and Validation Plan

### Phase 1 Testing: Agent Guidance

1. **Content Verification**
   ```bash
   # Verify section was added correctly
   grep -n "Understanding teal.transform Helper Functions" knowledge_base/agent.md

   # Check examples are present
   grep -c "variable_choices" knowledge_base/agent.md
   grep -c "value_choices" knowledge_base/agent.md
   ```

2. **MCP Tool Test**
   ```bash
   # Test that agent guidance returns new content
   # Use MCP inspector or test script to call tealflow_agent_guidance
   # Verify the new section appears in output
   ```

### Phase 2 Testing: Module Details Tool

1. **Test with Known Modules**
   ```bash
   # Test modules that use both function types
   uv run python -c "
   from tealflow_mcp.tools.module_details import tealflow_get_module_details
   from tealflow_mcp.models import ModuleDetailsInput

   # Test tm_g_km (has paramcd - value filtering)
   result = await tealflow_get_module_details(
       ModuleDetailsInput(module_name='tm_g_km', response_format='markdown')
   )
   print(result)
   assert 'value_choices' in result

   # Test tm_g_lineplot (has arm_var - variable selection)
   result = await tealflow_get_module_details(
       ModuleDetailsInput(module_name='tm_g_lineplot', response_format='markdown')
   )
   print(result)
   assert 'variable_choices' in result
   "
   ```

2. **Verify JSON Output**
   - Check that `parameter_role` field is present
   - Verify `recommended_function` is correct for known parameters
   - Confirm examples are well-formatted

### Phase 3 Testing: Code Generation

1. **Generate Code with Comments**
   ```bash
   uv run python -c "
   from tealflow_mcp.tools.code_generation import tealflow_generate_module_code
   from tealflow_mcp.models import GenerateModuleCodeInput

   # Test clinical module generation
   result = await tealflow_generate_module_code(
       GenerateModuleCodeInput(
           module_name='tm_g_km',
           label='Test KM',
           dataname='ADTTE',
           include_comments=True
       )
   )
   print(result)
   assert 'value_choices' in result
   assert 'FILTER' in result or 'filter' in result
   "
   ```

2. **Validate R Syntax**
   - Generate code for 3-4 modules
   - Save to temporary R file
   - Run basic R syntax check (if R is available)

### Integration Testing

1. **End-to-End LLM Test Scenarios**

   **Test 1: Explain the difference**
   - Prompt: "What's the difference between variable_choices and value_choices?"
   - Expected: Clear explanation with decision criteria

   **Test 2: Generate paramcd filtering code**
   - Prompt: "Generate code to filter ADLB by paramcd for glucose"
   - Expected: Uses `value_choices(ADLB, "PARAMCD", "PARAM")`

   **Test 3: Generate arm_var selection code**
   - Prompt: "Generate code to let user select arm variable"
   - Expected: Uses `variable_choices(ADSL, c("ARM", "ARMCD"))`

   **Test 4: Create complete module**
   - Prompt: "Create tm_g_km module for time-to-event analysis"
   - Expected: Correct function usage for all parameters
   - Validation: Run `tealflow_check_shiny_startup` on generated app

2. **Success Metrics**
   - [ ] LLM correctly uses `variable_choices` for arm_var 90%+ of time
   - [ ] LLM correctly uses `value_choices` for paramcd filtering 90%+ of time
   - [ ] Generated code passes R syntax validation
   - [ ] Generated apps pass `tealflow_check_shiny_startup`
   - [ ] Agent can explain the difference when asked
   - [ ] Agent references the new documentation section

---

## File Modifications Summary

### Files to Modify

1. **`/home/kuba/projects/TealFlowMCP/knowledge_base/agent.md`**
   - **Line:** ~465 (after "R Code Style Guidelines")
   - **Changes:** Add 150-200 line section
   - **Content:** Full documentation with examples and decision guide

2. **`/home/kuba/projects/TealFlowMCP/tealflow_mcp/tools/module_details.py`**
   - **Lines:** 89-145
   - **Changes:**
     - Add `classify_parameter_role()` function
     - Enhance markdown formatting (lines 94-118)
     - Enhance JSON output (lines 131-145)
     - Add example code snippets

3. **`/home/kuba/projects/TealFlowMCP/tealflow_mcp/tools/code_generation.py`**
   - **Lines:** 55-75, 96-125
   - **Changes:**
     - Add `generate_parameter_comment()` helper
     - Enhance clinical module comments
     - Improve and annotate general module examples

### No Changes Required

- JSON metadata files (save for potential Phase 4)
- Test files (unless adding new tests)
- MCP server registration (tealflow_mcp.py)
- Other tool files

---

## Implementation Timeline

### Recommended Schedule

**Day 1 - Morning (3 hours)**
- Phase 1: Write and add agent.md documentation section
- Test: Verify content with `tealflow_agent_guidance`

**Day 1 - Afternoon (4 hours)**
- Phase 2: Implement parameter classification in module_details.py
- Phase 2: Enhance markdown and JSON output
- Test: Verify with tm_g_km, tm_g_lineplot, tm_t_ancova

**Day 2 - Morning (2 hours)**
- Phase 3: Improve code generation comments
- Test: Generate code for multiple modules

**Day 2 - Afternoon (2 hours)**
- Integration testing with LLM scenarios
- Documentation cleanup
- Final validation

**Total Time:** ~1.5 days (11 hours)

---

## Future Enhancements (Optional)

### Option 4 Evolution Path

If metrics show specific parameters cause persistent confusion:

1. **Targeted JSON Metadata Enhancement**
   - Add `"recommended_function"` field to 15-20 high-confusion parameters
   - Focus on: arm_var, paramcd, strata_var, avisit, facet_var
   - Estimated effort: 3-4 hours

2. **Decision Tree Visualization**
   - Add ASCII decision tree to agent.md
   - Interactive flowchart for when to use each function
   - Estimated effort: 1 hour

### Optional: Standalone Reference Tool

Could add new MCP tool `tealflow_get_transform_reference`:
- Returns quick function reference
- Accepts optional `function_name` parameter
- Lightweight alternative to reading full agent.md
- Estimated effort: 2-3 hours

---

## Success Criteria

### Immediate Goals (After Implementation)

- [x] Documentation section exists in agent.md
- [x] Section includes 8+ annotated examples
- [x] Decision criteria clearly documented
- [x] Module details tool shows parameter guidance
- [x] Code generation includes helpful comments
- [x] All tests pass

### Long-term Goals (After 2-4 weeks)

- [ ] 90%+ correct function usage by LLMs
- [ ] Reduced confusion-related questions
- [ ] Generated code passes validation consistently
- [ ] No reported issues with incorrect function choice
- [ ] Positive feedback from users/developers

---

## External References

**Official teal.transform Documentation:**
- [variable_choices()](https://insightsengineering.github.io/teal.transform/v0.5.0-rc1/reference/variable_choices.html)
- [value_choices()](https://insightsengineering.github.io/teal.transform/latest-tag/reference/value_choices.html)
- [choices_selected()](https://insightsengineering.github.io/teal.transform/latest-tag/reference/choices_selected.html)
- [Package Index](https://insightsengineering.github.io/teal.transform/latest-tag/reference/index.html)

**Related TealFlow MCP Files:**
- [CLAUDE.md](/home/kuba/projects/TealFlowMCP/CLAUDE.md) - Project overview
- [knowledge_base/agent.md](/home/kuba/projects/TealFlowMCP/knowledge_base/agent.md) - Current agent guidance
- Research plan: [/home/kuba/.claude/plans/zesty-drifting-finch.md](/home/kuba/.claude/plans/zesty-drifting-finch.md)

---

## Notes

- This plan focuses on Option 2 (Moderate approach) as recommended
- Can evolve to Option 4 later with targeted metadata enhancements
- Low risk: Documentation-focused changes with tool enhancements
- High impact: Addresses root cause of LLM confusion
- Manageable maintenance: Clear ownership of documentation sections

**Last Updated:** 2026-02-27
**Next Review:** After implementation and 2 weeks of usage
