From 08d08aef260b58c5f70a411570b1d15ae3a2e260 Mon Sep 17 00:00:00 2001
From: Jason Liu <jason@jxnl.co>
Date: Sun, 22 Oct 2023 17:45:16 -0400
Subject: [PATCH] clean up

---
 instructor/distil.py         | 64 ++++++++++++++++++++----------------
 instructor/function_calls.py |  2 --
 2 files changed, 35 insertions(+), 31 deletions(-)

diff --git a/instructor/distil.py b/instructor/distil.py
index 179a9fb..94f02c8 100644
--- a/instructor/distil.py
+++ b/instructor/distil.py
@@ -171,36 +171,17 @@ class Instructions:
         base_model: BaseModel = type(resp)
 
         if finetune_format == FinetuneFormat.MESSAGES:
-            openai_function_call = openai_schema(base_model).openai_schema
-            func_def = get_signature_from_fn(fn).replace(fn.__name__, name)
-
-            str_args = ", ".join(map(str, args))
-            str_kwargs = (
-                ", ".join(f"{k}={json.dumps(v)}" for k, v in kwargs.items()) or None
+            openai_kwargs = self.openai_kwargs(name, fn, args, kwargs, base_model)
+            openai_kwargs.append(
+                {
+                    "role": "assistant",
+                    "function_call": {
+                        "name": base_model.__name__,
+                        "arguments": resp.model_dump_json(indent=self.indent),
+                    },
+                }
             )
-            call_args = ", ".join(filter(None, [str_args, str_kwargs]))
-
-            function_body = {
-                "messages": [
-                    {
-                        "role": "system",
-                        "content": f"Predict the results of this function:\n\n{func_def}",
-                    },
-                    {
-                        "role": "user",
-                        "content": f"Return `{name}({call_args})`",
-                    },
-                    {
-                        "role": "assistant",
-                        "function_call": {
-                            "name": openai_function_call["name"],
-                            "arguments": resp.model_dump_json(indent=self.indent),
-                        },
-                    },
-                ],
-                "functions": [openai_function_call],
-            }
-            self.logger.info(json.dumps(function_body))
+            self.logger.info(json.dumps(openai_kwargs))
 
         if finetune_format == FinetuneFormat.RAW:
             function_body = dict(
@@ -212,3 +193,28 @@ class Instructions:
                 schema=base_model.model_json_schema(),
             )
             self.logger.info(json.dumps(function_body))
+
+    def openai_kwargs(self, name, fn, args, kwargs, base_model):
+        openai_function_call = openai_schema(base_model).openai_schema
+        func_def = get_signature_from_fn(fn).replace(fn.__name__, name)
+
+        str_args = ", ".join(map(str, args))
+        str_kwargs = (
+            ", ".join(f"{k}={json.dumps(v)}" for k, v in kwargs.items()) or None
+        )
+        call_args = ", ".join(filter(None, [str_args, str_kwargs]))
+
+        function_body = {
+            "messages": [
+                {
+                    "role": "system",
+                    "content": f"Predict the results of this function:\n\n{func_def}",
+                },
+                {
+                    "role": "user",
+                    "content": f"Return `{name}({call_args})`",
+                },
+            ],
+            "functions": [openai_function_call],
+        }
+        return function_body
diff --git a/instructor/function_calls.py b/instructor/function_calls.py
index 6146ec1..5ef779b 100644
--- a/instructor/function_calls.py
+++ b/instructor/function_calls.py
@@ -188,8 +188,6 @@ class OpenAISchema(BaseModel):
                     f"the required parameters with correct types"
                 )
 
-        _remove_a_key(parameters, "title")
-        _remove_a_key(parameters, "additionalProperties")
         return {
             "name": schema["title"],
             "description": schema["description"],