Bugfix `logprob` in workflows.

src/workflows/configs.py

@@ -8,9 +8,11 @@ including model configurations, workflow settings, and other package-wide consta
 AVAILABLE_MODELS = {
     "OpenAI/gpt-4o": {
         "model": "gpt-4o-2024-11-20",
+        "logprobs": True,
     },
     "OpenAI/gpt-4o-mini": {
         "model": "gpt-4o-mini-2024-07-18",
+        "logprobs": True,
     },
     "OpenAI/gpt-3.5-turbo": {
         "model": "gpt-3.5-turbo-0125",
@@ -26,12 +28,15 @@ AVAILABLE_MODELS = {
     },
     "Cohere/command-r": {
         "model": "command-r-08-2024",
+        "logprobs": True,
     },
     "Cohere/command-r-plus": {
         "model": "command-r-plus-08-2024",
+        "logprobs": True,
     },
     "Cohere/command-r7b": {
         "model": "command-r7b-12-2024",
+        "logprobs": False,
     },
 }
 

src/workflows/executors.py

@@ -231,7 +231,7 @@ def execute_model_step(
     if return_full_content:
         result["content"] = api_response["content"]
     if logprobs:
-        result["logprob"] = api_response["log_prob"]
+        result["logprob"] = api_response.get("logprob")
     return result
 
 

src/workflows/llms.py

@@ -129,8 +129,6 @@ def _llm_completion(
     Raises:
         ValueError: If logprobs=True with Anthropic models
     """
-    if model not in AVAILABLE_MODELS:
-        raise ValueError(f"Model {model} not supported")
     model_name = AVAILABLE_MODELS[model]["model"]
     provider = model.split("/")[0]
     if provider == "Cohere":
@@ -173,13 +171,19 @@ def completion(
     Raises:
         ValueError: If logprobs=True with Anthropic models
     """
+    if model not in AVAILABLE_MODELS:
+        raise ValueError(f"Model {model} not supported")
+    if logprobs and not AVAILABLE_MODELS[model].get("logprobs", False):
+        logger.warning(f"{model} does not support logprobs feature, setting logprobs to False")
+        logprobs = False
+
     # Check cache first
     cached_response = llm_cache.get(model, system, prompt, response_format, temperature)
-    if cached_response is not None:
+    if cached_response and (not logprobs or cached_response.get("logprob")):
         logger.info(f"Cache hit for model {model}")
         return cached_response
 
-    logger.info(f"Cache miss for model {model}, calling API")
+    logger.info(f"Cache miss for model {model}, calling API. Logprobs: {logprobs}")
 
     # Continue with the original implementation for cache miss
     response = _llm_completion(model, system, prompt, response_format, temperature, logprobs)

src/workflows/qb_agents.py

@@ -1,14 +1,18 @@
 import time
 from typing import Any, Iterable, TypedDict
 
+from loguru import logger
+
 from .executors import WorkflowOutput, execute_workflow
 from .structs import TossupWorkflow, Workflow
 
 
-def _get_workflow_response(workflow: Workflow, available_vars: dict[str, Any]) -> tuple[WorkflowOutput, float]:
+def _get_workflow_response(
+    workflow: Workflow, available_vars: dict[str, Any], logprob_step: bool | str = False
+) -> tuple[WorkflowOutput, float]:
     """Get response from executing a complete workflow."""
     start_time = time.time()
-    workflow_output = execute_workflow(workflow, available_vars, return_full_content=True)
+    workflow_output = execute_workflow(workflow, available_vars, return_full_content=True, logprob_step=logprob_step)
     response_time = time.time() - start_time
     return workflow_output, response_time
 
@@ -78,15 +82,17 @@ class QuizBowlTossupAgent:
         """
         for i, question_text in enumerate(question_runs):
             # Execute the complete workflow
+            answer_var_step = self.workflow.outputs["answer"].split(".")[0]
             workflow_output, response_time = _get_workflow_response(
-                self.workflow, {self.external_input_variable: question_text}
+                self.workflow, {self.external_input_variable: question_text}, logprob_step=answer_var_step
             )
             final_outputs = workflow_output["final_outputs"]
-            buzz = self.workflow.buzzer.run(final_outputs["confidence"], logprob=final_outputs.get("logprob"))
+            buzz = self.workflow.buzzer.run(final_outputs["confidence"], logprob=workflow_output["logprob"])
             result: TossupResult = {
                 "position": i + 1,
                 "answer": final_outputs["answer"],
                 "confidence": final_outputs["confidence"],
+                "logprob": workflow_output["logprob"],
                 "buzz": buzz,
                 "question_fragment": question_text,
                 "step_contents": workflow_output["step_contents"],

src/workflows/structs.py

@@ -6,6 +6,8 @@ from typing import Any, Literal, Optional
 import numpy as np
 from pydantic import BaseModel, Field, model_validator
 
+from .configs import AVAILABLE_MODELS
+
 """
 Core data structures for defining workflows and their components.
 
@@ -259,6 +261,13 @@ class Workflow(BaseModel):
         """Get all model selections for all steps."""
         return {step_id: step.get_full_model_name() for step_id, step in self.steps.items()}
 
+    def get_output_model_selections(self) -> dict[str, str]:
+        """Get all output model selections for all steps."""
+        return {
+            output_var: target_var.split(".")[0] if target_var else None
+            for output_var, target_var in self.outputs.items()
+        }
+
     # Step update method
 
     def add_step(self, step: ModelStep) -> "Workflow":
@@ -305,14 +314,19 @@ class Buzzer(BaseModel):
         use_enum_values = True
         frozen = True
 
+    def update(self, **kwargs) -> "Buzzer":
+        """Update the buzzer with the given kwargs."""
+        return self.model_copy(update=kwargs)
+
     def run(self, confidence: float, prob: float | None = None, logprob: float | None = None) -> bool:
         """Run the buzzer logic."""
         if logprob is not None and prob is not None:
             raise ValueError("Cannot provide both logprob and prob")
-        if logprob is not None:
-            prob = np.exp(logprob)
         if self.prob_threshold is None:
             return confidence >= self.confidence_threshold
+        if logprob is None and prob is None:
+            raise ValueError("Must provide either logprob or prob if prob_threshold is not None")
+        prob = prob or float(np.exp(logprob))
         if self.method == BuzzerMethod.AND:
             return confidence >= self.confidence_threshold and prob >= self.prob_threshold
         elif self.method == BuzzerMethod.OR:
@@ -333,6 +347,24 @@ class TossupWorkflow(Workflow):
 
     buzzer: Buzzer
 
+    def get_answer_model(self, answer_var: str | None = None) -> str | None:
+        answer_var = answer_var or self.outputs["answer"]
+        if answer_var is None:
+            return None
+        step_id = answer_var.split(".")[0]
+        return self.steps[step_id].get_full_model_name()
+
+    def is_token_probs_supported(self, answer_var: str | None = None) -> bool:
+        model_name = self.get_answer_model(answer_var)
+        if model_name is None:
+            return True
+        return AVAILABLE_MODELS[model_name].get("logprobs", False)
+
     def update_buzzer(self, buzzer: Buzzer) -> "TossupWorkflow":
         """Update the buzzer."""
         return self.model_copy(update={"buzzer": buzzer})
+
+    def refresh_buzzer(self) -> "TossupWorkflow":
+        if not self.is_token_probs_supported():
+            return self.update_buzzer(self.buzzer.update(prob_threshold=None, method="AND"))
+        return self

src/workflows/utils.py

@@ -168,7 +168,7 @@ def topological_sort(dependencies: dict[str, set[str]]) -> list[str]:
 
     nodes = list(dependencies.keys())
     dependents: dict[str, list[str]] = {node: [] for node in nodes}
-    in_degree: dict[str, int] = {node: 0 for node in nodes}
+    in_degree: dict[str, int] = dict.fromkeys(nodes, 0)
 
     # Calculate in-degrees and build dependents list
     for node, deps in dependencies.items():