chatan 0.1.0__tar.gz → 0.1.1__tar.gz

{chatan-0.1.0 → chatan-0.1.1}/.gitignore

@@ -8,3 +8,5 @@ wheels/
 
 # Virtual environments
 .venv
+
+.DS_Store

{chatan-0.1.0 → chatan-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chatan
-Version: 0.1.0
+Version: 0.1.1
 Summary: Create synthetic datasets with LLM generators and samplers
 Project-URL: Documentation, https://github.com/cdreetz/chatan#readme
 Project-URL: Issues, https://github.com/cdreetz/chatan/issues

{chatan-0.1.0 → chatan-0.1.1}/src/chatan/__init__.py

@@ -1,6 +1,6 @@
 """Minos: Create synthetic datasets with LLM generators and samplers."""
 
-__version__ = "0.1.0"
+__version__ = "0.1.1"
 
 from .dataset import dataset
 from .generator import generator

{chatan-0.1.0 → chatan-0.1.1}/src/chatan/generator.py

@@ -58,14 +58,24 @@ class AnthropicGenerator(BaseGenerator):
 
 class GeneratorFunction:
     """Callable generator function for use in dataset schemas."""
-    
-    def __init__(self, generator: BaseGenerator, prompt_template: str):
+
+    def __init__(
+        self,
+        generator: BaseGenerator,
+        prompt_template: str,
+        variables: Optional[Dict[str, Any]] = None,
+    ):
         self.generator = generator
         self.prompt_template = prompt_template
-    
+        self.variables = variables or {}
+
     def __call__(self, context: Dict[str, Any]) -> str:
         """Generate content with context substitution."""
-        prompt = self.prompt_template.format(**context)
+        merged = dict(context)
+        for key, value in self.variables.items():
+            merged[key] = value(context) if callable(value) else value
+
+        prompt = self.prompt_template.format(**merged)
         result = self.generator.generate(prompt)
         return result.strip() if isinstance(result, str) else result
 
@@ -81,9 +91,19 @@ class GeneratorClient:
         else:
             raise ValueError(f"Unsupported provider: {provider}")
     
-    def __call__(self, prompt_template: str) -> GeneratorFunction:
-        """Create a generator function."""
-        return GeneratorFunction(self._generator, prompt_template)
+    def __call__(self, prompt_template: str, **variables) -> GeneratorFunction:
+        """Create a generator function.
+
+        Parameters
+        ----------
+        prompt_template:
+            Template string for the prompt.
+        **variables:
+            Optional variables to include when formatting the prompt. If a value
+            is callable it will be invoked with the row context when the
+            generator function is executed.
+        """
+        return GeneratorFunction(self._generator, prompt_template, variables)
 
 
 # Factory function

{chatan-0.1.0 → chatan-0.1.1}/tests/test_generator.py

@@ -289,6 +289,27 @@ class TestIntegration:
         assert result2 == "Response"
         assert mock_client.chat.completions.create.call_count == 2
 
+    @patch('openai.OpenAI')
+    def test_generator_function_with_variables(self, mock_openai):
+        """GeneratorFunction should accept default variables."""
+        mock_client = Mock()
+        mock_response = Mock()
+        mock_choice = Mock()
+        mock_choice.message.content = "Question about elephants"
+        mock_response.choices = [mock_choice]
+        mock_client.chat.completions.create.return_value = mock_response
+        mock_openai.return_value = mock_client
+
+        gen = generator("openai", "test-key")
+        func = gen("Question about {animal}", animal="elephants")
+        result = func({})
+
+        assert result == "Question about elephants"
+        mock_client.chat.completions.create.assert_called_once_with(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": "Question about elephants"}]
+        )
+
     def test_case_insensitive_provider(self):
         """Test that provider names are case insensitive."""
         with patch('chatan.generator.OpenAIGenerator') as mock_gen: