fix model-id error caused by previous refactor (#46)

Author: jakelorocco

mellea/backends/formatter.py

@@ -5,6 +5,7 @@
 import re
 import sys
 from collections.abc import Iterable, Mapping
+from dataclasses import fields
 from typing import Any
 
 import jinja2
@@ -396,14 +397,13 @@ def _get_model_id(self) -> str:
             "model_id was neither a `str` nor `ModelIdentifier`"
         )
 
-        # Go through the ModelIdentifier's fields, find one that isn't `"None"` or `""`.
-        ids = [model_id.hf_model_name, model_id.ollama_name]
-        model_id = ""
-        for val in ids:
-            if val != "None" and val != "":
-                model_id = val  # type: ignore
-                break
-        return model_id
+        # Go through the ModelIdentifier's fields, find one that can be matched against.
+        for field in fields(model_id):
+            val = getattr(model_id, field.name)
+            if val is not None and val != "":
+                return val
+
+        return ""  # Cannot match against any model identifiers. Will ultimately use default.
 
 
 def _simplify_model_string(input: str) -> str:

test/test_formatter_baseclasses.py

@@ -7,7 +7,7 @@
 import pytest
 
 from mellea.backends.formatter import TemplateFormatter
-from mellea.backends.model_ids import IBM_GRANITE_3_2_8B
+from mellea.backends.model_ids import ModelIdentifier, IBM_GRANITE_3_2_8B
 from mellea.stdlib.base import (
     BasicContext,
     CBlock,
@@ -284,6 +284,22 @@ def test_fake_model_id(instr: Instruction):
         "default" in tmpl.name
     ), "there should always be a default instruction template"
 
+def test_custom_model_id():
+    model_id = ModelIdentifier(mlx_name="new-model-here")
+    tf = TemplateFormatter(model_id=model_id)
+    assert tf._get_model_id() == "new-model-here", "getting the model id should always give a string if one exists"
+
+def test_empty_model_id(instr: Instruction):
+    model_id = ModelIdentifier()
+    tf = TemplateFormatter(model_id=model_id)
+    assert tf._get_model_id() == ""
+
+    tmpl = tf._load_template(instr.format_for_llm())
+    assert tmpl.name is not None
+    assert (
+        "default" in tmpl.name
+    ), "there should always be a default instruction template"
+
 
 def test_template_caching(instr: Instruction):
     """Caching shouldn't be interacted with this way by users.