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.outputs import RequestOutput
 from vllm.sampling_params import RequestOutputKind, SamplingParams
+from vllm.tokenizers import cached_tokenizer_from_config
+from vllm.tool_parsers import ToolParserManager
 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,42 @@ 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 (with encode/decode methods). 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 = cached_tokenizer_from_config(
+                    model_config=self.vllm_config.model_config,
+                )
+            parser_cls = ToolParserManager.get_tool_parser(self.tool_call_parser)  # type: ignore[union-attr]
+            parser = parser_cls(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 +588,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 +635,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 +658,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,
@@ -157,8 +156,6 @@ async def test_tool_parsing_multi_turn(policy, tokenizer):
     logger.info(f"Final answer: {final.text}")
     assert "579" in final.text, "Expected 123 + 456 = 579"
 
-    logger.info("✅ test_tool_parsing_multi_turn passed!")
-
 
 @requires_cuda
 @pytest.mark.asyncio
@@ -195,5 +192,3 @@ async def test_content_without_tool_calls(policy, tokenizer):
     assert (
         completion.content == completion.text
     ), "Content should equal text when no tools"
-
-    logger.info("✅ test_content_without_tool_calls passed!")

tests/unit_tests/test_generator.py

@@ -7,12 +7,9 @@
 """Unit tests for Generator's _to_completions and _extract_tool_calls logic."""
 
 import json
-from dataclasses import dataclass
-from typing import Any
 from unittest.mock import MagicMock
 
 import pytest
-
 from vllm.outputs import CompletionOutput, RequestOutput
 
 
@@ -26,20 +23,15 @@ def _import_error():
         return True
 
 
-@dataclass
-class TokenizerWrapper:
-    """Wrapper to mimic vLLM's tokenizer structure (has .tokenizer attr)."""
-
-    tokenizer: Any
-
-
 class _StubTokenizer:
     """Minimal stub tokenizer for initializing the Hermes tool parser in tests.
 
     The Hermes tool parser from vLLM requires:
     - get_vocab(): Returns vocab dict mapping tokens to ids
     - vocab: Direct vocab attribute
     - eos_token_id: End of sequence token id
+    - encode(text, add_special_tokens=False): Encode text to token ids
+    - decode(token_ids): Decode token ids to text
     - <tool_call> and </tool_call> tokens in vocab (for streaming support)
     """
 
@@ -50,22 +42,32 @@ def __init__(self):
             "<tool_call>": 1,
             "</tool_call>": 2,
         }
+        self._id_to_token = {v: k for k, v in self.vocab.items()}
         self.eos_token_id = 0
 
     def get_vocab(self) -> dict[str, int]:
         """Return vocabulary dict (required by Hermes tool parser)."""
         return self.vocab
 
+    def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
+        """Encode text to token ids. Returns ids for known tokens, empty otherwise."""
+        if text in self.vocab:
+            return [self.vocab[text]]
+        return [ord(c) for c in text]
+
+    def decode(self, token_ids: list[int]) -> str:
+        """Decode token ids to text."""
+        return "".join(self._id_to_token.get(tid, chr(tid)) for tid in token_ids)
 
-@pytest.fixture(scope="module")
-def tokenizer_wrapper():
-    """Create a tokenizer wrapper with Hermes-compatible tokenizer (has tool_call tokens)."""
 
-    return TokenizerWrapper(_StubTokenizer())
+@pytest.fixture(scope="module")
+def stub_tokenizer():
+    """Create a stub tokenizer compatible with Hermes tool parser."""
+    return _StubTokenizer()
 
 
 @pytest.fixture
-def generator_with_hermes(tokenizer_wrapper):
+def generator_with_hermes(stub_tokenizer):
     """Create Generator with hermes parser properly initialized."""
     from forge.actors.generator import Generator
 
@@ -74,7 +76,7 @@ def generator_with_hermes(tokenizer_wrapper):
         sampling_params={"max_tokens": 64},
         tool_call_parser="hermes",
     )
-    generator._tool_parser = generator._init_tool_parser(tokenizer_wrapper)
+    generator._tool_parser = generator._init_tool_parser(stub_tokenizer)
     generator.generator_version = 1
 
     return generator
@@ -115,7 +117,7 @@ def make_mock_request_output(
 class TestInitToolParser:
     """Test the _init_tool_parser method of Generator."""
 
-    def test_init_hermes_parser(self, tokenizer_wrapper):
+    def test_init_hermes_parser(self, stub_tokenizer):
         """Test that passing tool_call_parser='hermes' initializes the parser."""
         from forge.actors.generator import Generator
 
@@ -125,7 +127,7 @@ def test_init_hermes_parser(self, tokenizer_wrapper):
             tool_call_parser="hermes",
         )
 
-        parser = generator._init_tool_parser(tokenizer_wrapper)
+        parser = generator._init_tool_parser(stub_tokenizer)
 
         assert parser is not None
         assert hasattr(parser, "extract_tool_calls")
@@ -142,7 +144,7 @@ def test_init_parser_none_when_not_configured(self):
 
         assert generator.tool_call_parser is None
 
-    def test_init_parser_invalid_parser_name(self, tokenizer_wrapper):
+    def test_init_parser_invalid_parser_name(self, stub_tokenizer):
         """Test that invalid parser name returns None."""
         from forge.actors.generator import Generator
 
@@ -152,7 +154,7 @@ def test_init_parser_invalid_parser_name(self, tokenizer_wrapper):
             tool_call_parser="nonexistent_parser",
         )
 
-        parser = generator._init_tool_parser(tokenizer_wrapper)
+        parser = generator._init_tool_parser(stub_tokenizer)
         assert parser is None
 
 
@@ -298,7 +300,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 +321,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 +347,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