Refactor on Optimizers and Removing Proxy

CHANGELOG.md

@@ -42,7 +42,7 @@ and this project adheres to [Calendar Versioning](https://calver.org/) with form
 - Add `ForwardRecord` dataclass and `record` parameter to `run()` for backward pass support
 - Add `Module` base class with child/parameter registration and introspection methods
 - Add `LLMInference` atomic module for LLM API calls
-- Add tracing infrastructure with `Tracer`, `Proxy`, `GraphNode`, and `InferenceGraph`
+- Add tracing infrastructure with `Tracer`, `Value`, `GraphNode`, and `InferenceGraph`
 - Add `ExecutionState` for task management with dependency tracking and failure handling
 - Add `Scheduler` with concurrency control and `run()` function for end-to-end execution
 - Add data access operations (`__getitem__`, `__iter__`) and `NodeRef` for type-safe references

CLAUDE.md

@@ -99,7 +99,7 @@ plait/
 │       ├── graph.py         # InferenceGraph, GraphNode
 │       ├── types.py         # Core type definitions
 │       ├── errors.py        # Exception hierarchy
-│       ├── tracing/         # Tracer, Proxy, context
+│       ├── tracing/         # Tracer (Value-driven), context
 │       ├── execution/       # Scheduler, Executor, ExecutionState
 │       ├── resources/       # ResourceManager, config, rate limiting
 │       ├── optimization/    # Loss, Optimizer, backward, feedback
@@ -142,7 +142,7 @@ plait/
 The system has four main layers:
 
 1. **User Code**: `Module`, `LLMInference`, `Parameter`
-2. **Tracing**: `Tracer`, `Proxy`, `InferenceGraph` - captures DAG from forward()
+2. **Tracing**: `Tracer`, `Value`, `InferenceGraph` - captures DAG from forward()
 3. **Execution**: `Scheduler`, `ExecutionState` - async execution with priority queue
 4. **Infrastructure**: LLM clients, rate limiting, checkpointing
 
@@ -183,19 +183,19 @@ def record_call(
     module: Module,
     args: tuple[Any, ...],
     kwargs: dict[str, Any],
-) -> Proxy:
+) -> Value:
     """Record a module invocation and return a proxy for its output.
 
     Creates a new graph node representing this call and tracks dependencies
-    based on any Proxy objects in the arguments.
+    based on any Value objects in the arguments.
 
     Args:
         module: The module being called.
         args: Positional arguments passed to the module.
         kwargs: Keyword arguments passed to the module.
 
     Returns:
-        A Proxy representing the eventual output of this call.
+        A Value representing the eventual output of this call.
 
     Raises:
         TracingError: If called outside of an active trace context.

Makefile

@@ -3,20 +3,20 @@
 ci: lint types test
 
 test:
-	uv run pytest
+	uv run --group dev -m pytest
 
 test-unit:
-	uv run pytest tests/unit
+	uv run --group dev -m pytest tests/unit
 
 test-integration:
-	uv run pytest tests/integration
+	uv run --group dev -m pytest tests/integration
 
 lint:
-	uv run ruff format .
-	uv run ruff check --fix .
+	uv run --group dev ruff format .
+	uv run --group dev ruff check --fix .
 
 types:
-	uv run ty check
+	uv run --group dev ty check
 
 example:
 	@for f in examples/[0-9]*.py; do \
@@ -25,11 +25,11 @@ example:
 	done
 
 docs:
-	uv run mkdocs build
+	uv run --group dev mkdocs build
 	cp index.html styles.css public/
 
 docs-serve:
-	uv run mkdocs serve
+	uv run --group dev mkdocs serve
 
 doctest:
 	@echo "=== Running plait vs Pydantic AI comparison ==="

design_docs/DESIGN.md

@@ -321,7 +321,6 @@ plait/
 ├── src/
 │   ├── op.py           # Op base class
 │   ├── values.py       # Value, ValueKind, helpers
-│   ├── proxy.py        # (legacy) Proxy helpers, if needed
 │   ├── tracer.py       # Tracer implementation
 │   ├── dag.py          # DAG and Node definitions
 │   ├── scheduler.py    # Priority queue + rate limiter

design_docs/README.md

@@ -195,7 +195,7 @@ plait/
 │       ├── values.py           # Value, ValueKind, helpers
 │       ├── functional.py       # plait.functional
 │       ├── graph.py            # InferenceGraph, GraphNode
-│       ├── tracing/            # Tracer + trace context (+ legacy proxy)
+│       ├── tracing/            # Tracer + trace context (Value-driven)
 │       ├── execution/          # Scheduler, ExecutionState, checkpoints
 │       ├── resources/          # ResourceConfig/ResourceManager, rate limiting
 │       ├── optimization/       # Loss/Feedback/Optimizer, backward propagation

design_docs/architecture.md

@@ -478,7 +478,6 @@ plait/
 │       │   ├── __init__.py
 │       │   ├── tracer.py       # Tracer
 │       │   ├── values.py       # Value, ValueKind, helpers
-│       │   ├── proxy.py        # (legacy) Proxy helpers, if needed
 │       │   └── context.py      # TraceContext
 │       ├── execution/
 │       │   ├── __init__.py

design_docs/optimization.md

@@ -201,6 +201,7 @@ class ForwardRecord:
     node_inputs: dict[str, dict[str, Any]]   # node_id -> resolved input values
     node_outputs: dict[str, Any]              # node_id -> output value
     module_map: dict[str, Module]    # node_id -> module instance
+    node_parameters: dict[str, list[Parameter]]  # node_id -> direct parameters
 
     # Optional metadata
     execution_order: list[str] = field(default_factory=list)

design_docs/tracing.md

@@ -14,7 +14,7 @@ embedded in args/kwargs.
 
 ## Tracing Modes
 
-The Tracer supports two modes for backwards compatibility:
+The Tracer uses a single Value-driven mode:
 
 ### Value-Driven Mode (Recommended)
 
@@ -30,22 +30,8 @@ tracer = Tracer()
 graph = tracer.trace_values(module, "input text")
 ```
 
-### Proxy-Based Mode (Legacy)
-
-Use `Tracer.trace()` for backwards-compatible Proxy-based tracing:
-
-- Inputs are wrapped as Proxy objects with node IDs
-- Module calls return Proxy objects
-- Dependencies are tracked via `Proxy.node_id`
-- Arguments are stored with `NodeRef` placeholders
-
-```python
-tracer = Tracer()
-graph = tracer.trace(module, "input text")
-```
-
-Both modes produce equivalent `InferenceGraph` structures. Value-driven mode
-is preferred for new code as it aligns with the broader Value-based data model.
+Value-driven mode aligns with the broader Value-based data model and is the
+only supported tracing approach.
 
 ## Core Concepts
 

docs/comparison/compare_dspy.py

@@ -197,7 +197,7 @@ class _FactsCombiner(Module):
 
     This module formats the facts from two documents into a single
     prompt for comparison. Using a module ensures proper tracing
-    during the forward pass (Proxy objects are resolved correctly).
+    during the forward pass (Value objects are resolved correctly).
     """
 
     def forward(self, facts1: str, facts2: str) -> str:
@@ -236,7 +236,7 @@ def forward(self, doc1: str, doc2: str) -> str:
         facts1 = self.extractor(doc1)
         facts2 = self.extractor(doc2)
 
-        # Combine facts using the combiner module (resolves Proxy objects)
+        # Combine facts using the combiner module (resolves Value objects)
         combined = self.combiner(facts1, facts2)
 
         # This depends on both facts, so it waits for both to complete

docs/comparison/compare_langgraph.py

@@ -194,7 +194,7 @@ class _FactsCombiner(Module):
 
     This module formats the facts from two documents into a single
     prompt for comparison. Using a module ensures proper tracing
-    during the forward pass (Proxy objects are resolved correctly).
+    during the forward pass (Value objects are resolved correctly).
     """
 
     def forward(self, facts1: str, facts2: str) -> str:
@@ -235,7 +235,7 @@ def forward(self, doc1: str, doc2: str) -> str:
         facts1 = self.extractor(doc1)
         facts2 = self.extractor(doc2)
 
-        # Combine facts using the combiner module (resolves Proxy objects)
+        # Combine facts using the combiner module (resolves Value objects)
         combined = self.combiner(facts1, facts2)
 
         # This depends on both facts, so it waits for both to complete

docs/comparison/compare_pydantic_ai.py

@@ -149,7 +149,7 @@ class _FactsCombiner(Module):
 
     This module formats the facts from two documents into a single
     prompt for comparison. Using a module ensures proper tracing
-    during the forward pass (Proxy objects are resolved correctly).
+    during the forward pass (Value objects are resolved correctly).
     """
 
     def forward(self, facts1: str, facts2: str) -> str:
@@ -190,7 +190,7 @@ def forward(self, doc1: str, doc2: str) -> str:
         facts1 = self.extractor(doc1)
         facts2 = self.extractor(doc2)
 
-        # Combine facts using the combiner module (resolves Proxy objects)
+        # Combine facts using the combiner module (resolves Value objects)
         combined = self.combiner(facts1, facts2)
 
         # This depends on both facts, so it waits for both to complete

docs/comparison/dspy.md

@@ -57,7 +57,7 @@ class ExtractAndCompare(Module):
         facts1 = self.extractor(doc1)
         facts2 = self.extractor(doc2)
 
-        # Combine facts using the combiner module (resolves Proxy objects)
+        # Combine facts using the combiner module (resolves Value objects)
         combined = self.combiner(facts1, facts2)
 
         # This depends on both facts, waits for both to complete

docs/comparison/langgraph.md

@@ -58,7 +58,7 @@ class ExtractAndCompare(Module):
         facts1 = self.extractor(doc1)
         facts2 = self.extractor(doc2)
 
-        # Combine facts using the combiner module (resolves Proxy objects)
+        # Combine facts using the combiner module (resolves Value objects)
         combined = self.combiner(facts1, facts2)
 
         # This depends on both facts, waits for both to complete

docs/comparison/pydantic-ai.md

@@ -58,7 +58,7 @@ class ExtractAndCompare(Module):
         facts1 = self.extractor(doc1)
         facts2 = self.extractor(doc2)
 
-        # Combine facts using the combiner module (resolves Proxy objects)
+        # Combine facts using the combiner module (resolves Value objects)
         combined = self.combiner(facts1, facts2)
 
         # This depends on both facts, waits for both to complete

docs/design/architecture.md

@@ -478,7 +478,6 @@ plait/
 │       │   ├── __init__.py
 │       │   ├── tracer.py       # Tracer
 │       │   ├── values.py       # Value, ValueKind, helpers
-│       │   ├── proxy.py        # (legacy) Proxy helpers, if needed
 │       │   └── context.py      # TraceContext
 │       ├── execution/
 │       │   ├── __init__.py

docs/design/optimization.md

@@ -201,6 +201,7 @@ class ForwardRecord:
     node_inputs: dict[str, dict[str, Any]]   # node_id -> resolved input values
     node_outputs: dict[str, Any]              # node_id -> output value
     module_map: dict[str, Module]    # node_id -> module instance
+    node_parameters: dict[str, list[Parameter]]  # node_id -> direct parameters
 
     # Optional metadata
     execution_order: list[str] = field(default_factory=list)

docs/design/tracing.md

@@ -14,7 +14,7 @@ embedded in args/kwargs.
 
 ## Tracing Modes
 
-The Tracer supports two modes for backwards compatibility:
+The Tracer uses a single Value-driven mode:
 
 ### Value-Driven Mode (Recommended)
 
@@ -30,22 +30,8 @@ tracer = Tracer()
 graph = tracer.trace_values(module, "input text")
 ```
 
-### Proxy-Based Mode (Legacy)
-
-Use `Tracer.trace()` for backwards-compatible Proxy-based tracing:
-
-- Inputs are wrapped as Proxy objects with node IDs
-- Module calls return Proxy objects
-- Dependencies are tracked via `Proxy.node_id`
-- Arguments are stored with `NodeRef` placeholders
-
-```python
-tracer = Tracer()
-graph = tracer.trace(module, "input text")
-```
-
-Both modes produce equivalent `InferenceGraph` structures. Value-driven mode
-is preferred for new code as it aligns with the broader Value-based data model.
+Value-driven mode aligns with the broader Value-based data model and is the
+only supported tracing approach.
 
 ## Core Concepts
 

examples/02_llm_pipeline.py

@@ -108,7 +108,7 @@ def __init__(self) -> None:
     def forward(self, text: str) -> str:
         analyses = self.perspectives(text)
         # Use explicit key access instead of .items() iteration
-        # (Proxy objects don't support iteration during tracing)
+        # (Value objects don't support full iteration during tracing)
         combined = (
             f"## Technical\n{analyses['technical']}\n\n"
             f"## Business\n{analyses['business']}\n\n"
@@ -307,10 +307,15 @@ def forward(self, document: str) -> str:
     multi_dict = MultiPerspectiveDict()
     print(f"   Keys: {list(multi_dict.analyzers.keys())}")
     print("   Access by key:")
-    for key in multi_dict.analyzers:
-        print(f"      analyzers['{key}']: alias={multi_dict.analyzers[key].alias}")
+    for key, module in multi_dict.analyzers.items():
+        alias = module.alias if isinstance(module, LLMInference) else "<unknown>"
+        print(f"      analyzers['{key}']: alias={alias}")
     print("   Attribute access:")
-    print(f"      analyzers.technical: alias={multi_dict.analyzers.technical.alias}")
+    technical = multi_dict.analyzers.technical
+    technical_alias = (
+        technical.alias if isinstance(technical, LLMInference) else "<unknown>"
+    )
+    print(f"      analyzers.technical: alias={technical_alias}")
 
     print("\n" + "=" * 60)
     print("Use bind() or ExecutionSettings to connect to real endpoints.")

examples/03_tracing.py

@@ -15,8 +15,8 @@
 
 from plait.graph import visualize_graph
 from plait.module import LLMInference, Module
-from plait.tracing.proxy import Proxy
 from plait.tracing.tracer import Tracer
+from plait.values import Value
 
 # --- Sequential Pipeline ---
 
@@ -29,7 +29,7 @@ def __init__(self) -> None:
         self.step1 = LLMInference(alias="fast", system_prompt="Summarize.")
         self.step2 = LLMInference(alias="smart", system_prompt="Analyze.")
 
-    def forward(self, text: str) -> Proxy:
+    def forward(self, text: Value) -> Value:
         summary = self.step1(text)
         return self.step2(summary)
 
@@ -46,7 +46,7 @@ def __init__(self) -> None:
         self.b = LLMInference(alias="llm", system_prompt="Perspective B")
         self.c = LLMInference(alias="llm", system_prompt="Perspective C")
 
-    def forward(self, text: str) -> dict[str, Proxy]:
+    def forward(self, text: Value) -> dict[str, Value]:
         # All three depend only on input - can run in parallel
         return {"a": self.a(text), "b": self.b(text), "c": self.c(text)}
 
@@ -63,7 +63,7 @@ def __init__(self) -> None:
         self.branch_b = LLMInference(alias="fast", system_prompt="View B")
         self.synth = LLMInference(alias="smart", system_prompt="Synthesize")
 
-    def forward(self, text: str) -> Proxy:
+    def forward(self, text: Value) -> Value:
         a = self.branch_a(text)
         b = self.branch_b(text)
         return self.synth(a, b)  # Waits for both branches
@@ -82,7 +82,7 @@ def __init__(self) -> None:
         self.analyze_b = LLMInference(alias="llm", system_prompt="Aspect B")
         self.combine = LLMInference(alias="smart", system_prompt="Combine")
 
-    def forward(self, text: str) -> Proxy:
+    def forward(self, text: Value) -> Value:
         cleaned = self.preprocess(text)
         a = self.analyze_a(cleaned)  # Fan-out from cleaned
         b = self.analyze_b(cleaned)
@@ -108,7 +108,7 @@ def print_graph_info(name: str, graph: Any) -> None:
     # Sequential
     print("\n1. Two-Stage Pipeline")
     print("-" * 40)
-    seq_graph = tracer.trace(TwoStage(), "input text")
+    seq_graph = tracer.trace_values(TwoStage(), "input text")
     print(f"   Nodes: {len(seq_graph.nodes)}")
     print("   Execution order:")
     for i, node_id in enumerate(seq_graph.topological_order(), 1):
@@ -119,7 +119,7 @@ def print_graph_info(name: str, graph: Any) -> None:
     # Parallel
     print("\n2. Parallel Fan-out")
     print("-" * 40)
-    par_graph = tracer.trace(Parallel(), "input text")
+    par_graph = tracer.trace_values(Parallel(), "input text")
     print(f"   Nodes: {len(par_graph.nodes)}")
     print(f"   Outputs (independent): {len(par_graph.output_ids)}")
     input_id = par_graph.input_ids[0]
@@ -131,7 +131,7 @@ def print_graph_info(name: str, graph: Any) -> None:
     # Diamond
     print("\n3. Diamond Pattern (Fan-out + Fan-in)")
     print("-" * 40)
-    diamond_graph = tracer.trace(Diamond(), "input text")
+    diamond_graph = tracer.trace_values(Diamond(), "input text")
     print(f"   Nodes: {len(diamond_graph.nodes)}")
     print("   Structure:")
     print("      input")
@@ -145,7 +145,7 @@ def print_graph_info(name: str, graph: Any) -> None:
     # Multi-stage
     print("\n4. Multi-Stage Pipeline")
     print("-" * 40)
-    multi_graph = tracer.trace(MultiStage(), "input text")
+    multi_graph = tracer.trace_values(MultiStage(), "input text")
     print(f"   Nodes: {len(multi_graph.nodes)}")
     print("   Topological order:")
     for i, node_id in enumerate(multi_graph.topological_order(), 1):