Improve gemini support

PrefectHQ · Sep 13, 2024 · a0ef300 · a0ef300
1 parent ad2b455
commit a0ef300
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/src/controlflow/tools/tools.py b/src/controlflow/tools/tools.py
@@ -51,10 +51,17 @@ class Tool(ControlFlowModel):
     metadata: dict = {}
 
     fn: Callable = Field(None, exclude=True)
+    _lc_tool: Optional[langchain_core.tools.BaseTool] = None
 
     def to_lc_tool(self) -> dict:
-        payload = self.model_dump(include={"name", "description", "parameters"})
-        return dict(type="function", function=payload)
+        if self._lc_tool is None:
+            self._lc_tool = langchain_core.tools.StructuredTool.from_function(
+                self.fn,
+                name=self.name,
+                description=self.description,
+            )
+
+        return self._lc_tool
 
     @prefect_task(task_run_name="Tool call: {self.name}")
     def run(self, input: dict):