bug: b64 encoded identifiers (#33)

- if a generator identifier string ends in , the split function (on ) splits into too large of an array
- instead, just split on the first  as the rest of the string is comma delimtited

Author: briangreunke

rigging/generator/base.py

@@ -26,8 +26,7 @@
 
 @t.runtime_checkable
 class LazyGenerator(t.Protocol):
-    def __call__(self) -> type[Generator]:
-        ...
+    def __call__(self) -> type[Generator]: ...
 
 
 g_providers: dict[str, type[Generator] | LazyGenerator] = {}
@@ -382,16 +381,14 @@ def chat(
         self,
         messages: t.Sequence[MessageDict],
         params: GenerateParams | None = None,
-    ) -> ChatPipeline:
-        ...
+    ) -> ChatPipeline: ...
 
     @t.overload
     def chat(
         self,
         messages: t.Sequence[Message] | MessageDict | Message | str | None = None,
         params: GenerateParams | None = None,
-    ) -> ChatPipeline:
-        ...
+    ) -> ChatPipeline: ...
 
     def chat(
         self,
@@ -460,17 +457,15 @@ def chat(
     generator: Generator,
     messages: t.Sequence[MessageDict],
     params: GenerateParams | None = None,
-) -> ChatPipeline:
-    ...
+) -> ChatPipeline: ...
 
 
 @t.overload
 def chat(
     generator: Generator,
     messages: t.Sequence[Message] | MessageDict | Message | str | None = None,
     params: GenerateParams | None = None,
-) -> ChatPipeline:
-    ...
+) -> ChatPipeline: ...
 
 
 def chat(
@@ -597,7 +592,7 @@ def get_generator(identifier: str, *, params: GenerateParams | None = None) -> G
     if "," in model:
         try:
             model, kwargs_str = model.split(",", 1)
-            kwargs = dict(arg.split("=") for arg in kwargs_str.split(","))
+            kwargs = dict(arg.split("=", 1) for arg in kwargs_str.split(","))
         except Exception as e:
             raise InvalidModelSpecifiedError(identifier) from e
 

tests/test_generator_ids.py

@@ -75,3 +75,9 @@ def test_register_generator() -> None:
     register_generator("echo", EchoGenerator)
     generator = get_generator("echo!test")
     assert isinstance(generator, EchoGenerator)
+
+
+def test_get_generator_b64() -> None:
+    generator = get_generator("litellm!test_model,api_key=ZXhhbXBsZXRleHQ=")
+    assert isinstance(generator, LiteLLMGenerator)
+    assert generator.model == "test_model"