From e264bf77645a801976c7d183fba681be91d16cfd Mon Sep 17 00:00:00 2001
From: Reinier van der Leer <pwuts@agpt.co>
Date: Sat, 8 Jun 2024 15:28:52 +0200
Subject: [PATCH] `forge.llm.providers.schema` + `code_flow_executor` lint-fix
 and cleanup

---
 .../components/code_flow_executor/__init__.py |  6 +--
 forge/forge/llm/providers/schema.py           | 48 ++++++++++---------
 2 files changed, 29 insertions(+), 25 deletions(-)

diff --git a/forge/forge/components/code_flow_executor/__init__.py b/forge/forge/components/code_flow_executor/__init__.py
index 0e8e36466e..517b37f1f7 100644
--- a/forge/forge/components/code_flow_executor/__init__.py
+++ b/forge/forge/components/code_flow_executor/__init__.py
@@ -1,3 +1,3 @@
-from .code_flow_executor import (
-    CodeFlowExecutionComponent
-)
+from .code_flow_executor import CodeFlowExecutionComponent
+
+__all__ = ["CodeFlowExecutionComponent"]
diff --git a/forge/forge/llm/providers/schema.py b/forge/forge/llm/providers/schema.py
index d7a432674a..1b31606aaa 100644
--- a/forge/forge/llm/providers/schema.py
+++ b/forge/forge/llm/providers/schema.py
@@ -144,38 +144,42 @@ class CompletionModelFunction(BaseModel):
         )
         return f"{self.name}: {self.description}. Params: ({params})"
 
-    def fmt_header(self, impl="pass", force_async=False) -> str:
+    def fmt_function_stub(self, impl: str = "pass") -> str:
         """
-        Formats and returns the function header as a string with types and descriptions.
+        Formats and returns a function stub as a string with types and descriptions.
 
         Returns:
             str: The formatted function header.
         """
-        def indent(content: str, spaces: int = 4):
-            return " " * spaces + content.replace("\n", "\n" + " " * spaces)
+        from forge.llm.prompting.utils import indent
 
         params = ", ".join(
-            f"{name}: {p.python_type}{f'= {str(p.default)}' if p.default else ' = None' if not p.required else ''}"
+            f"{name}: {p.python_type}"
+            + (
+                f" = {str(p.default)}"
+                if p.default
+                else " = None"
+                if not p.required
+                else ""
+            )
             for name, p in self.parameters.items()
         )
-        func = "async def" if self.is_async or force_async else "def"
-        return_str = f" -> {self.return_type}" if self.return_type else ""
-        return f"{func} {self.name}({params}){return_str}:\n" + indent(
-            (
-                '"""\n'
-                f"{self.description}\n\n"
-                "Params:\n"
-                + indent(
-                    "\n".join(
-                        f"{name}: {param.description}"
-                        for name, param in self.parameters.items()
-                        if param.description
-                    )
+        _def = "async def" if self.is_async else "def"
+        _return = f" -> {self.return_type}" if self.return_type else ""
+        return f"{_def} {self.name}({params}){_return}:\n" + indent(
+            '"""\n'
+            f"{self.description}\n\n"
+            "Params:\n"
+            + indent(
+                "\n".join(
+                    f"{name}: {param.description}"
+                    for name, param in self.parameters.items()
+                    if param.description
                 )
-                + "\n"
-                '"""\n'
-                f"{impl}"
-            ),
+            )
+            + "\n"
+            '"""\n'
+            f"{impl}"
         )
 
     def validate_call(