refactor: replace n parameter with sampling_params in generate()

  Allows per-request override of any sampling parameter (temperature,
  top_p, n, etc.) instead of just n. Preserves output_kind=FINAL_ONLY
  enforcement from post_init logic.

Author: gitlost-murali

src/forge/actors/generator.py

@@ -288,15 +288,19 @@ def split_keys(keys):
 
     @endpoint
     async def generate(
-        self, prompt: str, *, priority: int = 0, n: int | None = None
+        self,
+        prompt: str,
+        *,
+        priority: int = 0,
+        sampling_params: SamplingParams | None = None,
     ) -> list[Completion]:
         """Generate a response for the given prompt
 
         Args:
             prompt (str): The prompt to generate a response for.
             priority (int, optional): The priority of the request. Defaults to 0.
-            n (int, optional): Number of completions to generate. If not provided, uses the default
-                from self.sampling_params.n.
+            sampling_params (SamplingParams, optional): Sampling parameters to use for this request.
+                If not provided, uses self.sampling_params.
 
         Returns:
             list[Completion]: n completions from vLLM based on your prompt.
@@ -305,10 +309,10 @@ async def generate(
         t.start()
         record_metric("generator/generate/count_requests", 1, Reduce.SUM)
 
-        if n is not None and n != self.sampling_params.n:
-            params = self.sampling_params.__replace__(n=n)
-        else:
-            params = self.sampling_params
+        params = sampling_params or self.sampling_params
+        # Ensure output_kind is set to FINAL_ONLY (as required by post_init)
+        if params.output_kind != RequestOutputKind.FINAL_ONLY:
+            params = params.__replace__(output_kind=RequestOutputKind.FINAL_ONLY)
 
         self.request_id += 1 % sys.maxsize
         request_id = str(self.request_id)

tests/unit_tests/test_generator_config.py

@@ -132,27 +132,6 @@ def test_generator_yaml_config_loading(self):
             self.assertEqual(generator.sampling_params.n, 2)
             self.assertEqual(generator.sampling_params.max_tokens, 32)
 
-    @pytest.mark.skipif(
-        _import_error(),
-        reason="Import error, likely due to missing dependencies on CI.",
-    )
-    def test_generate_n_parameter_logic(self):
-        from forge.actors.generator import Generator
-
-        generator = Generator(sampling_params={"n": 2, "max_tokens": 16})
-        base_params = generator.sampling_params
-
-        def get_params_for(n_override: int | None):
-            if n_override in (None, base_params.n):
-                return base_params
-            return base_params.__replace__(n=n_override)
-
-        self.assertIs(get_params_for(None), base_params)
-        self.assertIs(get_params_for(2), base_params)
-        updated = get_params_for(4)
-        self.assertEqual(updated.n, 4)
-        self.assertIsNot(updated, base_params)
-
 
 if __name__ == "__main__":
     unittest.main()