From 5d20a7a6c67d08d4f8d01d6d4f0fac2ff74f06db Mon Sep 17 00:00:00 2001
From: Jack Hopkins <jack@raveler.co.uk>
Date: Thu, 2 Nov 2023 18:47:11 +0000
Subject: [PATCH 1/2] Commit Subject Line: Add dill to requirements.txt

Commit Description:
---
 requirements.txt             |  3 ++-
 src/monkey_patch/monkey.py   | 14 -------------
 src/monkey_patch/register.py |  3 ++-
 src/monkey_patch/utils.py    | 39 +++++++++++++++++++++++++++++++++++-
 4 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index d733e7f..183136d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,4 +4,5 @@ numpy~=1.24.4
 python-dotenv==1.0.0
 bitarray==2.8.2
 pydantic==2.4.2
-fastapi~=0.104.0
\ No newline at end of file
+fastapi~=0.104.0
+dill~=0.3.7
\ No newline at end of file
diff --git a/src/monkey_patch/monkey.py b/src/monkey_patch/monkey.py
index 62cf613..6b9fdb8 100644
--- a/src/monkey_patch/monkey.py
+++ b/src/monkey_patch/monkey.py
@@ -181,19 +181,6 @@ def mock_func(*args, **kwargs):
             else:
                 return patched_func(*args, **kwargs)
 
-        def _get_args(func_args, kwarg_names, num_args):
-            num_pos_args = num_args - len(kwarg_names)  # Calculate number of positional arguments
-            args_for_call = func_args[:num_pos_args]
-            # Pop keyword arguments off the stack
-            kwargs_for_call = {}  # New dictionary to hold keyword arguments for the call
-            for name in reversed(kwarg_names):  # Reverse to match the order on the stack
-                try:
-                    kwargs_for_call[name] = func_args.pop()  # Pop the value off the stack
-                except IndexError:
-                    print(f"Debug: func_args is empty, can't pop for {name}")
-            func_args = func_args[:-num_pos_args]  # Remove the positional arguments from func_args
-            return args_for_call, func_args, kwargs_for_call
-
         return wrapper
 
     @staticmethod
@@ -227,7 +214,6 @@ def wrapper(*args, **kwargs):
                     raise TypeError(f"Output type was not valid. Expected an object of type {function_description.output_type_hint}, got '{output.generated_response}'")
                 output.generated_response = choice
                 output.distilled_model = False
-                
 
             datapoint = FunctionExample(args, kwargs, output.generated_response)
             if output.suitable_for_finetuning and not output.distilled_model:
diff --git a/src/monkey_patch/register.py b/src/monkey_patch/register.py
index eff6bfc..c1e7521 100644
--- a/src/monkey_patch/register.py
+++ b/src/monkey_patch/register.py
@@ -2,6 +2,7 @@
 from typing import get_type_hints, Literal, Optional
 
 from monkey_patch.models.function_description import FunctionDescription
+from monkey_patch.utils import get_source
 
 alignable_functions = {}
 
@@ -74,7 +75,7 @@ def get_class_definition(class_type):
                 elif hasattr(class_type, "__args__"):  # Access inner types
                     return [get_class_definition(arg) for arg in class_type.__args__ if arg is not None]
             elif inspect.isclass(class_type) and class_type.__module__ != "builtins":
-                return inspect.getsource(class_type)
+                return get_source(class_type)
             return class_type.__name__
 
         # Extract class definitions for input and output types
diff --git a/src/monkey_patch/utils.py b/src/monkey_patch/utils.py
index 7c2a382..fa71985 100644
--- a/src/monkey_patch/utils.py
+++ b/src/monkey_patch/utils.py
@@ -3,7 +3,7 @@
 import json
 import typing
 from typing import get_args, Literal
-
+import inspect
 
 def json_default(thing):
     try:
@@ -121,3 +121,40 @@ def get_key(args, kwargs) -> tuple:
     args_tuple = _deep_tuple(args)
     kwargs_tuple = _deep_tuple(kwargs)
     return args_tuple, kwargs_tuple
+
+def _get_source_ipython(func) -> str:
+    """
+    Get the source code of a function from IPython (to support Colab and Jupyter notebooks)
+    :param func: The function to get the source code from
+    :return: The source code of the function
+    """
+    # Get the IPython instance
+    from IPython import get_ipython
+    ipython = get_ipython()
+
+    # Get the input history
+    input_cells = ipython.history_manager.input_hist_parsed
+
+    class_name = func.__name__
+    source_code = None
+
+    for cell in input_cells:
+        if f"class {class_name}" in cell:
+            source_code = cell
+            break
+
+    # If found, print the source code
+    return source_code
+
+def get_source(func) -> str:
+    """
+    Get the source code of a function
+    Args:
+        func (function): the function to get the source code from
+    Returns:
+        source (str): the source code of the function
+    """
+    try:
+        return inspect.getsource(func)
+    except Exception:
+        return _get_source_ipython(func)
\ No newline at end of file

From d1cff15f07d4a15ded0a2885a7fa9b9b16dd8e72 Mon Sep 17 00:00:00 2001
From: Jack Hopkins <jack@raveler.co.uk>
Date: Thu, 2 Nov 2023 18:49:36 +0000
Subject: [PATCH 2/2] Commit Subject: Update requirements.txt

Commit Description:
---
 requirements.txt | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 183136d..d733e7f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,5 +4,4 @@ numpy~=1.24.4
 python-dotenv==1.0.0
 bitarray==2.8.2
 pydantic==2.4.2
-fastapi~=0.104.0
-dill~=0.3.7
\ No newline at end of file
+fastapi~=0.104.0
\ No newline at end of file