fix integrations

Author: kylemontgomery1

rllm/engine/rollout/openai_engine.py

@@ -34,6 +34,7 @@ def __init__(self, model: str, tokenizer=None, api_retries: int = 3, base_url: s
     async def chat_completion(self, messages: list[dict], **kwargs) -> ModelOutput:
         sampling_params = self.sampling_params.copy()
         sampling_params.update(kwargs)
+        sampling_params.pop("model", None)
         retries = self.api_retries
         while retries > 0:
             try:

rllm/integrations/smolagents.py

@@ -38,7 +38,7 @@ class SmolTool:
 
 # Import BaseAgent from rLLM for wrapper classes
 from rllm.agents.agent import Step, Trajectory
-from rllm.engine.rollout.rollout_engine import ModelOutput, RolloutEngine
+from rllm.engine import ModelOutput
 
 logger = logging.getLogger(__name__)
 
@@ -94,7 +94,7 @@ async def arun(
 
         self.logger.log_task(
             content=self.task.strip(),
-            subtitle=f"{type(self.model).__name__} - {(self.model.model_id if hasattr(self.model, 'model_id') else '')}",
+            subtitle=f"{type(self.model).__name__}",
             level=LogLevel.INFO,
             title=self.name if hasattr(self, "name") else None,
         )
@@ -128,7 +128,7 @@ async def arun(
             total_output_tokens = 0
             correct_token_usage = True
             for step in self.memory.steps:
-                if isinstance(step, ActionStep | PlanningStep):
+                if isinstance(step, ActionStep) or isinstance(step, PlanningStep):
                     if step.token_usage is None:
                         correct_token_usage = False
                         break
@@ -312,15 +312,12 @@ class RLLMOpenAIModel(SmolModel):
     - Skips MessageRole.TOOL_CALL messages (handled as part of assistant messages)
     """
 
-    def __init__(self, rollout_engine: RolloutEngine, **kwargs):
+    def __init__(self, rollout_engine=None, **kwargs):
         """
         Initialize the RLLM-integrated OpenAI model.
 
         Args:
             rollout_engine: rLLM's RolloutEngine instance
-            application_id: Unique identifier for the application
-            model_id: The model identifier (for compatibility)
-            sampling_params: Sampling parameters for generation
             **kwargs: Additional arguments (ignored, for compatibility)
         """
         self.rollout_engine = rollout_engine
@@ -330,6 +327,9 @@ def __init__(self, rollout_engine: RolloutEngine, **kwargs):
         # Store kwargs for potential future use
         self.kwargs = kwargs
 
+        if not rollout_engine:
+            raise ValueError("rollout_engine is required for RLLMOpenAIModel. Pass an instance of rLLM's RolloutEngine.")
+
     async def generate_async(self, messages: list[dict[str, Any]], stop_sequences: list[str] | None = None, response_format: dict[str, str] | None = None, tools_to_call_from: list | None = None, **kwargs) -> Any:
         """
         Async version of generate that can be called from async contexts.
@@ -342,7 +342,7 @@ async def generate_async(self, messages: list[dict[str, Any]], stop_sequences: l
                 # Handle ChatMessage objects from SmolAgent using the helper method
                 prompt = self._convert_smolagent_messages_to_openai(messages)
 
-            model_output: ModelOutput = await self.rollout_engine.get_model_response(prompt, **kwargs)
+            model_output: ModelOutput = await self.rollout_engine.get_model_response(prompt, max_tokens=kwargs.pop("max_tokens", 4096), **kwargs)
 
             # Extract text and token usage from ModelOutput
             response_text = model_output.text

rllm/integrations/strands.py

@@ -9,45 +9,22 @@
 from strands.types.tools import ToolSpec
 
 from rllm.agents.agent import Step, Trajectory
-from rllm.engine.rollout import ModelOutput, RolloutEngine
+from rllm.engine import ModelOutput, RolloutEngine
 
 T = TypeVar("T", bound=BaseModel)
 
 
 class RLLMModel(Model):
     """Model class that uses rLLM's RolloutEngine for inference."""
 
-    def __init__(self, rollout_engine: RolloutEngine, model_id: str = "gpt-4", **model_config):
+    def __init__(self, rollout_engine: RolloutEngine, **kwargs):
         """Initialize the RLLMModel.
 
         Args:
             rollout_engine: The rLLM RolloutEngine instance to use for inference
-            model_id: The model ID to use
-            **model_config: Additional model configuration
         """
         self.rollout_engine = rollout_engine
-        self.config = {"model_id": model_id, "params": model_config}
-
-    def update_config(self, **model_config: Any) -> None:
-        """Update the model configuration.
-
-        Args:
-            **model_config: Configuration overrides.
-        """
-        if "model_id" in model_config:
-            self.config["model_id"] = model_config.pop("model_id")
-
-        if "params" not in self.config:
-            self.config["params"] = {}
-        self.config["params"].update(model_config)
-
-    def get_config(self) -> dict[str, Any]:
-        """Get the model configuration.
-
-        Returns:
-            The model's configuration.
-        """
-        return self.config.copy()
+        self.kwargs = kwargs
 
     async def structured_output(self, output_model: type[T], prompt: Messages, system_prompt: str | None = None, **kwargs: Any) -> AsyncGenerator[dict[str, T | Any], None]:
         """Get structured output from the model.
@@ -73,9 +50,7 @@ async def structured_output(self, output_model: type[T], prompt: Messages, syste
             messages[-1]["content"] = f"{original_content}\n\nPlease respond with a JSON object that matches this schema: {output_model.model_json_schema()}"
 
         # Get response from rollout engine
-        model_output: ModelOutput = await self.rollout_engine.get_model_response(messages, model=self.config["model_id"], **self.config.get("params", {}), **kwargs)
-
-        response_text = model_output.text
+        response_text = (await self.rollout_engine.get_model_response(messages, **kwargs)).text
 
         try:
             # Try to parse the response as JSON and convert to the output model
@@ -129,7 +104,7 @@ async def stream(
         yield {"contentBlockStart": {"start": {}}}
 
         # Get response from rollout engine
-        model_output: ModelOutput = await self.rollout_engine.get_model_response(chat_messages, model=self.config["model_id"], **self.config.get("params", {}), **kwargs)
+        model_output: ModelOutput = await self.rollout_engine.get_model_response(chat_messages, **kwargs)
 
         # Extract text from ModelOutput
         response_text = model_output.text
@@ -203,11 +178,10 @@ def _convert_messages_to_chat_format(self, messages: Messages, system_prompt: st
 
 
 class StrandsAgent(Agent):
-    def __init__(self, model: str, **kwargs):
+    def __init__(self, model: RLLMModel, **kwargs):
         """Initialize StrandsAgent with trajectory tracking.
 
         Args:
-            model: The model to use (can be a string or Model instance)
             **kwargs: Additional arguments to pass to the base Agent class
         """