quick chat exmaple formatting fixed, added promtp formatting using bl…

…ack, fixed test closure. Closes #38

examples/quick_chat.py

@@ -35,13 +35,13 @@ def format_message_history(message_history : List[Tuple[str, str]]) -> str:
 @ell.lm(model="gpt-4o-mini", temperature=0.3, max_tokens=20)
 def chat(message_history : List[Tuple[str, str]], *, personality : str):
 
-    return [
-        ell.system(f"""You are
-    {personality}. 
+        return [
+            ell.system(f"""Here is your description.
+{personality}. 
 
-    Your goal is to come up with a response to a chat. Only respond in one sentence (should be like a text message in informality.) Never use Emojis."""),
-        ell.user(format_message_history(message_history)),
-    ]
+Your goal is to come up with a response to a chat. Only respond in one sentence (should be like a text message in informality.) Never use Emojis."""),
+            ell.user(format_message_history(message_history)),
+        ]
 
 
 

pyproject.toml

@@ -37,6 +37,7 @@ requests = "^2.32.3"
 typing-extensions = "^4.12.2"
 
 
+black = "^24.8.0"
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.3.2"
 

src/ell/util/closure.py

@@ -40,6 +40,7 @@ def xD():
 import importlib.util
 import re
 from collections import deque
+import black
 
 DELIM = "$$$$$$$$$$$$$$$$$$$$$$$$$"
 FORBIDDEN_NAMES = ["ell", "lstr"]
@@ -94,12 +95,26 @@ def lexical_closure(
     CLOSURE_SOURCE[hash(func)] = dirty_src 
 
     dsrc = _clean_src(dirty_src_without_func)
+
+    # Format the sorce and dsrc soruce using Black
+    source = _format_source(source)
+    dsrc = _format_source(dsrc)
+
     fn_hash = _generate_function_hash(source, dsrc, func.__qualname__)
 
     _update_ell_func(outer_ell_func, source, dsrc, globals_and_frees['globals'], globals_and_frees['frees'], fn_hash, uses)
 
     return (dirty_src, (source, dsrc), ({fn_hash} if not initial_call and hasattr(outer_ell_func, "__ell_func__") else uses))
 
+
+def _format_source(source: str) -> str:
+    """Format the source code using Black."""
+    try:
+        return black.format_str(source, mode=black.Mode())
+    except:
+        # If Black formatting fails, return the original source
+        return source
+
 def _get_globals_and_frees(func: Callable) -> Dict[str, Dict]:
     """Get global and free variables for a function."""
     globals_dict = collections.OrderedDict(dill.detect.globalvars(func))
@@ -331,7 +346,10 @@ def get_referenced_names(code: str, module_name: str):
 
 def lexically_closured_source(func):
     _, fnclosure, uses = lexical_closure(func, initial_call=True, recursion_stack=[])
-    return fnclosure, uses
+    source, dsrc = fnclosure
+    formatted_source = _format_source(source)
+    formatted_dsrc = _format_source(dsrc)
+    return (formatted_source, formatted_dsrc), uses
 
 import ast
 

tests/test_closure.py

@@ -78,13 +78,13 @@ def test_get_referenced_names():
 
 def test_is_function_called():
     code = """
-    def foo():
-        pass
-    
-    def bar():
-        foo()
-    
-    x = 1 + 2
+def foo():
+    pass
+
+def bar():
+    foo()
+
+x = 1 + 2
     """
     assert is_function_called("foo", code)
     assert not is_function_called("bar", code)