logprobs fixed

Author: mahirlabibdihan

weboperator/models/azure_openai.py

@@ -43,7 +43,7 @@ def wrapper(*args, **kwargs):  # type: ignore
                     return func(*args, **kwargs)
                 # Retry on specified errors
                 except Exception as e:
-                    print(f"Error {e}")
+                    print(f"Error {type(e)} -> {e} ")
                     # Increment retries
                     num_retries += 1
 
@@ -66,24 +66,71 @@ def chat(self, messages: list[dict], **kwargs) -> str | list[str]:
         Chat completion using the chat/completions endpoint.
         Supports multi-modal inputs (text + images) for vision models.
         """
-        response = AzureClient.chat.completions.create(
-            model=self.name,
-            messages=messages,
-            max_tokens=self.max_tokens,
-            temperature=self.temperature,
-            top_p=self.top_p,
-            # reasoning_effort=self.reasoning_effort,
-            n=kwargs.get("n", self.n),
-            logprobs=True,
-            top_logprobs=10,
-        )
+        try:
+            if "gpt-5" in self.name:
+                # For gpt-5 models, we might want to set different parameters
+                response = AzureClient.chat.completions.create(
+                    model=self.name,
+                    messages=messages,
+                    # max_completion_tokens=self.max_tokens,
+                    temperature=self.temperature,
+                    top_p=self.top_p,
+                    n=kwargs.get("n", self.n),
+                )
+            elif "o4" in self.name:
+                response = AzureClient.chat.completions.create(
+                    model=self.name,
+                    messages=messages,
+                    max_tokens=self.max_tokens,
+                    reasoning_effort="medium",
+                    stream=False,
+                )
+            else:
+                if kwargs.get("logprobs", False):   
+                    response = AzureClient.chat.completions.create(
+                        model=self.name,
+                        messages=messages,
+                        # max_tokens=self.max_tokens,
+                        temperature=self.temperature,
+                        top_p=self.top_p,
+                        # reasoning_effort=self.reasoning_effort,
+                        n=kwargs.get("n", self.n),
+                        logprobs=True,
+                        top_logprobs=10,
+                    )
+                else:
+                    response = AzureClient.chat.completions.create(
+                        model=self.name,
+                        messages=messages,
+                        # max_tokens=self.max_tokens,
+                        temperature=self.temperature,
+                        top_p=self.top_p,
+                        # reasoning_effort=self.reasoning_effort,
+                        n=kwargs.get("n", self.n),
+                    )
+        except openai.BadRequestError as e:
+            print(f"BadRequestError: {e}")
+            return "", []
+        except Exception as e:
+            raise e
 
         if len(response.choices) == 0:
             raise ValueError("No choices returned from the model.")
+        
+        predictions = [
+            choice.message.content.strip()
+            for choice in response.choices
+            if choice.message.content.strip()
+        ]
+        
+        if len(predictions) == 0:
+            raise ValueError("No valid predictions returned from the model.")
 
         top_logprobs = [
             choice.logprobs.content
             for choice in response.choices
             if hasattr(choice, "logprobs") and choice.logprobs
         ]
-        return response.choices[0].message.content.strip(), top_logprobs[0]
+        if kwargs.get("logprobs", False):
+            return predictions[0], top_logprobs[0]
+        return predictions[0], []

weboperator/models/openhf.py

@@ -82,16 +82,26 @@ def chat(self, messages: list[dict], **kwargs) -> str:
             except requests.exceptions.RequestException as e:
                 raise ConnectionError(f"Could not connect to HUGGING_FACE_API_SERVER: {e}")
 
-        response = OpenHFClient.chat.completions.create(
-            model=self.name,
-            messages=messages,
-            # max_tokens=self.max_tokens,
-            temperature=self.temperature,
-            top_p=self.top_p,
-            n=kwargs.get("n", self.n),
-            logprobs=True,
-            top_logprobs=10,
-        )
+        if kwargs.get("logprobs", False):    
+            response = OpenHFClient.chat.completions.create(
+                model=self.name,
+                messages=messages,
+                # max_tokens=self.max_tokens,
+                temperature=self.temperature,
+                top_p=self.top_p,
+                n=kwargs.get("n", self.n),
+                logprobs=True,
+                top_logprobs=10,
+            )
+        else:
+            response = OpenHFClient.chat.completions.create(
+                model=self.name,
+                messages=messages,
+                # max_tokens=self.max_tokens,
+                temperature=self.temperature,
+                top_p=self.top_p,
+                n=kwargs.get("n", self.n),
+            )
 
         # Raise OpenHFError if we get invalid response to trigger retry
         if not response or not hasattr(response, "choices") or not response.choices:

weboperator/models/openrouter.py

@@ -106,16 +106,26 @@ def chat(self, messages: list[dict], **kwargs) -> str:
         Chat completion using the chat/completions endpoint.
         Supports multi-modal inputs (text + images) for vision models.
         """
-        response = self.client.chat.completions.create(
-            model=self.name,
-            messages=messages,
-            # max_tokens=self.max_tokens,
-            temperature=self.temperature,
-            top_p=self.top_p,
-            n=kwargs.get("n", self.n),
-            logprobs=True,
-            top_logprobs=10,
-        )
+        if kwargs.get("logprobs", False):    
+            response = self.client.chat.completions.create(
+                model=self.name,
+                messages=messages,
+                # max_tokens=self.max_tokens,
+                temperature=self.temperature,
+                top_p=self.top_p,
+                n=kwargs.get("n", self.n),
+                logprobs=True,
+                top_logprobs=10,
+            )
+        else:
+            response = self.client.chat.completions.create(
+                model=self.name,
+                messages=messages,
+                # max_tokens=self.max_tokens,
+                temperature=self.temperature,
+                top_p=self.top_p,
+                n=kwargs.get("n", self.n),
+            )
         # print(response.choices[0])
         usage = getattr(response, "usage", None)
         if usage:

weboperator/webprm.py

@@ -201,7 +201,7 @@ def _get_brief_context(trajectory: List[Dict[str, Any]]) -> str:
         if len(trajectory) < 2:
             return "No previous actions."
 
-        # Show trajectory[-4:-2] for context (2 steps before current)
+        # Show up to the last 10 prior steps (excluding current)
         context_steps = trajectory[-10:-1] if len(trajectory) >= 10 else trajectory[:-1]
         # print(f"Context steps for evaluation: {context_steps}")
         if not context_steps:
@@ -641,7 +641,7 @@ def evaluate(
         # print("USER PROMPT: ", message[-1]["content"])
 
         for _ in range(3):  # Try up to 3 times to get a valid answer
-            response, scores = self.reward_model.chat(message)
+            response, scores = self.reward_model.chat(message, logprobs=True)
             generated_text = response
             if "# Checklist Evaluation" in generated_text:
                 # print(scores)