From cc4e164f9cc0bfc20565de5ac313014b28e19bd0 Mon Sep 17 00:00:00 2001
From: Jason Schleifer <jason.schleifer@gmail.com>
Date: Tue, 11 Feb 2025 15:36:14 +1300
Subject: [PATCH] added LocalRerankDriver

---
 __init__.py                            |  2 ++
 nodes/drivers/gtUILocalRerankDriver.py | 25 +++++++++++++++++++++++++
 2 files changed, 27 insertions(+)
 create mode 100644 nodes/drivers/gtUILocalRerankDriver.py

diff --git a/__init__.py b/__init__.py
index 437ed78..3fe7f08 100644
--- a/__init__.py
+++ b/__init__.py
@@ -168,6 +168,7 @@ from .nodes.drivers.gtUILeonardoImageGenerationDriver import (
 # - LM Studio
 from .nodes.drivers.gtUILMStudioChatPromptDriver import gtUILMStudioChatPromptDriver
 from .nodes.drivers.gtUILMStudioEmbeddingDriver import gtUILMStudioEmbeddingDriver
+from .nodes.drivers.gtUILocalRerankDriver import gtUILocalRerankDriver
 
 # - Local
 from .nodes.drivers.gtUILocalVectorStoreDriver import gtUILocalVectorStoreDriver
@@ -393,6 +394,7 @@ NODE_CLASS_MAPPINGS = {
     "Griptape Embedding Driver: Voyage AI": gtUIVoyageAiEmbeddingDriver,
     # RERANK DRIVER
     "Griptape Rerank Driver: Cohere": gtUICohereRerankDriver,
+    "Griptape Rerank Driver: Local": gtUILocalRerankDriver,
     # VECTOR STORE DRIVERS
     "Griptape Vector Store Driver: Amazon OpenSearch": gtUIAmazonOpenSearchVectorStoreDriver,
     "Griptape Vector Store Driver: Azure MongoDB": gtUIAzureMongoDbVectorStoreDriver,
diff --git a/nodes/drivers/gtUILocalRerankDriver.py b/nodes/drivers/gtUILocalRerankDriver.py
new file mode 100644
index 0000000..e210dee
--- /dev/null
+++ b/nodes/drivers/gtUILocalRerankDriver.py
@@ -0,0 +1,25 @@
+from typing import Any, Tuple
+
+from griptape.drivers.rerank.local import LocalRerankDriver
+
+from .gtUIBaseRerankDriver import gtUIBaseRerankDriver
+
+
+class gtUILocalRerankDriver(gtUIBaseRerankDriver):
+    DESCRIPTION = "Griptape Local Rerank Driver"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        inputs = super().INPUT_TYPES()
+        inputs["required"].update()
+        inputs["optional"].update()
+
+        return inputs
+
+    CATEGORY = "Griptape/Agent Drivers/Rerank"
+
+    RETURN_TYPES = ("RERANK_DRIVER",)
+
+    def create(self, **kwargs) -> Tuple[Any, ...]:
+        driver = LocalRerankDriver()
+        return (driver,)