fix: port tool parsing to v1 Generator after rebase

The rebase onto main dropped tool parsing changes from generator.py
because main refactored it into a thin version-detection wrapper.
This ports the tool parsing logic to the v1 Generator:

- Add tool_call_parser field and _init_tool_parser() method
- Add _extract_tool_calls() using vLLM's ToolParserManager
- Update _to_completions() to populate tool_calls/content on Completion
- Fix tests to match v1 Generator interface (_to_completions takes prompt)
- Fix integration test to import Generator directly (Policy doesn't exist)

Author: gitlost-murali

src/forge/actors/vllm/v1/generator.py

@@ -38,8 +38,15 @@
 from torchstore.api import _controller as get_torchstore_controller
 from vllm.engine.arg_utils import EngineArgs
 from vllm.entrypoints.llm import UsageContext
+from vllm.entrypoints.openai.protocol import (
+    ChatCompletionRequest,
+    ExtractedToolCallInformation,
+    ToolCall,
+)
+from vllm.entrypoints.openai.tool_parsers import ToolParserManager
 from vllm.outputs import RequestOutput
 from vllm.sampling_params import RequestOutputKind, SamplingParams
+from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs
 from vllm.v1.engine.async_llm import AsyncLLM
 
 logger = logging.getLogger(__name__)
@@ -78,6 +85,7 @@ class Generator(ForgeActor):
     sampling_params: SamplingParams | Mapping = field(default_factory=SamplingParams)
     prefetch_weights_to_shm: bool = True
     n_fetcher_procs: int = 8
+    tool_call_parser: str | None = None
 
     def __post_init__(self):
         super().__init__()
@@ -91,6 +99,8 @@ def __post_init__(self):
             self.engine_args = EngineArgs(**self.engine_args)
         self.vllm_config = self.engine_args.create_engine_config(UsageContext.LLM_CLASS)
 
+        self._tool_parser = None  # Will hold ToolParser instance if configured
+
         if isinstance(self.sampling_params, Mapping):
             self.sampling_params = SamplingParams.from_optional(**self.sampling_params)
             self.sampling_params.output_kind = RequestOutputKind.FINAL_ONLY
@@ -273,9 +283,44 @@ async def setup(self, host_mesh, worker_registry, gpu_ids: list[str]):
             )
         logger.info(f"Retrieved workers from registry: {self.workers}")
 
+        if self.tool_call_parser is not None:
+            self._tool_parser = self._init_tool_parser()
+
         if self.prefetch_weights_to_shm:
             self._spawn_fetchers()
 
+    def _init_tool_parser(self, tokenizer=None):  # type: ignore[no-untyped-def]
+        """Initialize the tool parser based on configuration.
+
+        Args:
+            tokenizer: Optional tokenizer wrapper (with .tokenizer attr). If not provided,
+                one is created from vllm_config. Passing explicitly is useful for testing.
+
+        Returns:
+            Initialized ToolParser instance, or None if tool parsing is not configured.
+        """
+        try:
+            if tokenizer is None:
+                tokenizer = init_tokenizer_from_configs(
+                    model_config=self.vllm_config.model_config,
+                    scheduler_config=self.vllm_config.scheduler_config,
+                    lora_config=self.vllm_config.lora_config,
+                )
+            parser_cls = ToolParserManager.get_tool_parser(self.tool_call_parser)  # type: ignore[union-attr]
+            parser = parser_cls(tokenizer.tokenizer)
+            logger.info(f"Initialized tool parser: {self.tool_call_parser}")
+            return parser
+        except KeyError:
+            available = list(ToolParserManager.tool_parsers.keys())
+            logger.error(
+                f"Unknown tool parser: '{self.tool_call_parser}'. "
+                f"Available parsers: {available}"
+            )
+            return None
+        except Exception as e:
+            logger.error(f"Failed to initialize tool parser: {e}")
+            return None
+
     def _spawn_fetchers(self):
         """Spawn weight fetchers that prefetch weights from torchstore to shared memory.
 
@@ -545,6 +590,38 @@ def _extract_logprobs(self, output) -> torch.Tensor | None:
             )
         return None
 
+    def _extract_tool_calls(self, model_output: str) -> ExtractedToolCallInformation:
+        """Extract tool calls from model output using the configured tool parser.
+
+        Args:
+            model_output: Raw text output from the model.
+
+        Returns:
+            ExtractedToolCallInformation with parsed tool calls and remaining content.
+        """
+        if self._tool_parser is None:
+            return ExtractedToolCallInformation(
+                tools_called=False, tool_calls=[], content=model_output
+            )
+
+        try:
+            dummy_request = ChatCompletionRequest(
+                model=self.vllm_config.model_config.model,
+                messages=[{"role": "user", "content": ""}],
+                seed=42,  # to calm the linter
+            )
+
+            extracted = self._tool_parser.extract_tool_calls(
+                model_output, dummy_request
+            )
+
+            return extracted
+        except Exception as e:
+            logger.warning(f"Failed to parse tool calls: {e}")
+            return ExtractedToolCallInformation(
+                tools_called=False, tool_calls=[], content=model_output
+            )
+
     def _to_completions(
         self, request_output: RequestOutput, prompt: str
     ) -> list[Completion]:
@@ -560,6 +637,14 @@ def _to_completions(
         completions = []
 
         for output in request_output.outputs:
+            tool_calls: list[ToolCall] = []
+            content: str | None = None
+
+            if self._tool_parser is not None:
+                extracted = self._extract_tool_calls(output.text)
+                tool_calls = extracted.tool_calls
+                content = extracted.content
+
             completion = Completion(
                 prompt=to_prompt(prompt),
                 text=output.text,
@@ -575,6 +660,8 @@ def _to_completions(
                 stop_reason=output.finish_reason,
                 generator_version=self.generator_version,
                 metadata={"num_cached_tokens": request_output.num_cached_tokens},
+                tool_calls=tool_calls,
+                content=content,
             )
             completions.append(completion)
 

tests/integration_tests/test_tool_parsing.py

@@ -21,8 +21,7 @@
 import pytest
 import pytest_asyncio
 import torch
-
-from forge.rl import Policy
+from forge.actors.generator import Generator
 from huggingface_hub import snapshot_download
 from vllm.transformers_utils.tokenizer import get_tokenizer
 
@@ -89,7 +88,7 @@ def tokenizer():
 async def policy(model_path):
     """Create and teardown policy service for each test."""
     logger.info("Setting up policy service...")
-    policy = await Policy.options(
+    policy = await Generator.options(
         procs=1,
         num_replicas=1,
         with_gpus=True,

tests/unit_tests/test_generator.py

@@ -12,7 +12,6 @@
 from unittest.mock import MagicMock
 
 import pytest
-
 from vllm.outputs import CompletionOutput, RequestOutput
 
 
@@ -298,7 +297,7 @@ def test_to_completions_without_tool_parser(self):
             outputs=[{"text": "The answer is 4.", "token_ids": [10, 20, 30]}],
         )
 
-        completions = generator._to_completions(request_output)
+        completions = generator._to_completions(request_output, request_output.prompt)
 
         assert len(completions) == 1
         completion = completions[0]
@@ -319,7 +318,7 @@ def test_to_completions_no_tool_call_with_parser(self, generator_with_hermes):
             ],
         )
 
-        completions = generator._to_completions(request_output)
+        completions = generator._to_completions(request_output, request_output.prompt)
 
         assert len(completions) == 1
         completion = completions[0]
@@ -345,7 +344,7 @@ def test_to_completions_multiple_outputs(self, generator_with_hermes):
             ],
         )
 
-        completions = generator._to_completions(request_output)
+        completions = generator._to_completions(request_output, request_output.prompt)
 
         assert len(completions) == 2
         # First completion has tool call