From ed166c390b634f9dcded10edfac87b35fa09fdd8 Mon Sep 17 00:00:00 2001
From: Vasily Vasinov <des.elyon@gmail.com>
Date: Tue, 2 Jul 2024 15:14:58 -0600
Subject: [PATCH] Rename TextRerankRagModule to TextChunkRerankRagModule

---
 CHANGELOG.md                                                | 2 +-
 docs/griptape-framework/engines/rag-engines.md              | 2 +-
 griptape/engines/rag/modules/__init__.py                    | 4 ++--
 ...erank_rag_module.py => text_chunks_rerank_rag_module.py} | 2 +-
 ..._rag_module.py => test_text_chunks_rerank_rag_module.py} | 6 +++---
 5 files changed, 8 insertions(+), 8 deletions(-)
 rename griptape/engines/rag/modules/retrieval/{text_rerank_rag_module.py => text_chunks_rerank_rag_module.py} (87%)
 rename tests/unit/engines/rag/modules/retrieval/{test_text_rerank_rag_module.py => test_text_chunks_rerank_rag_module.py} (70%)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 78676ff88..91c8fe89f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -42,7 +42,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - RAG modules:
     - Retrieval:
       - `VectorStoreRetrievalRagModule` for retrieving text chunks from vector stores.
-      - `TextRerankRagModule` for re-ranking retrieved results.
+      - `TextChunksRerankRagModule` for re-ranking retrieved results.
     - Response:
       - `MetadataBeforeResponseRagModule` for appending metadata.
       - `RulesetsBeforeResponseRagModule` for appending rulesets.
diff --git a/docs/griptape-framework/engines/rag-engines.md b/docs/griptape-framework/engines/rag-engines.md
index 4bd5750ef..c03bf0f47 100644
--- a/docs/griptape-framework/engines/rag-engines.md
+++ b/docs/griptape-framework/engines/rag-engines.md
@@ -20,7 +20,7 @@ No modules implemented yet.
 
 #### Retrieval
 - `TextRetrievalRagModule` is for retrieving text chunks.
-- `TextRerankRagModule` is for re-ranking retrieved results.
+- `TextChunksRerankRagModule` is for re-ranking retrieved results.
 
 #### Response
 - `MetadataBeforeResponseRagModule` is for appending metadata.
diff --git a/griptape/engines/rag/modules/__init__.py b/griptape/engines/rag/modules/__init__.py
index 000762243..eb514331d 100644
--- a/griptape/engines/rag/modules/__init__.py
+++ b/griptape/engines/rag/modules/__init__.py
@@ -2,7 +2,7 @@
 from .query.base_query_rag_module import BaseQueryRagModule
 from .retrieval.base_retrieval_rag_module import BaseRetrievalRagModule
 from .retrieval.base_rerank_rag_module import BaseRerankRagModule
-from .retrieval.text_rerank_rag_module import TextRerankRagModule
+from .retrieval.text_chunks_rerank_rag_module import TextChunksRerankRagModule
 from .retrieval.vector_store_retrieval_rag_module import VectorStoreRetrievalRagModule
 from .response.base_before_response_rag_module import BaseBeforeResponseRagModule
 from .response.base_after_response_rag_module import BaseAfterResponseRagModule
@@ -17,7 +17,7 @@
     "BaseQueryRagModule",
     "BaseRetrievalRagModule",
     "BaseRerankRagModule",
-    "TextRerankRagModule",
+    "TextChunksRerankRagModule",
     "VectorStoreRetrievalRagModule",
     "BaseBeforeResponseRagModule",
     "BaseAfterResponseRagModule",
diff --git a/griptape/engines/rag/modules/retrieval/text_rerank_rag_module.py b/griptape/engines/rag/modules/retrieval/text_chunks_rerank_rag_module.py
similarity index 87%
rename from griptape/engines/rag/modules/retrieval/text_rerank_rag_module.py
rename to griptape/engines/rag/modules/retrieval/text_chunks_rerank_rag_module.py
index 5ea4c7a18..b78c3a880 100644
--- a/griptape/engines/rag/modules/retrieval/text_rerank_rag_module.py
+++ b/griptape/engines/rag/modules/retrieval/text_chunks_rerank_rag_module.py
@@ -7,6 +7,6 @@
 
 
 @define(kw_only=True)
-class TextRerankRagModule(BaseRerankRagModule):
+class TextChunksRerankRagModule(BaseRerankRagModule):
     def run(self, context: RagContext) -> Sequence[BaseArtifact]:
         return self.rerank_driver.run(context.query, context.text_chunks)
diff --git a/tests/unit/engines/rag/modules/retrieval/test_text_rerank_rag_module.py b/tests/unit/engines/rag/modules/retrieval/test_text_chunks_rerank_rag_module.py
similarity index 70%
rename from tests/unit/engines/rag/modules/retrieval/test_text_rerank_rag_module.py
rename to tests/unit/engines/rag/modules/retrieval/test_text_chunks_rerank_rag_module.py
index 57fc38e64..838f02e37 100644
--- a/tests/unit/engines/rag/modules/retrieval/test_text_rerank_rag_module.py
+++ b/tests/unit/engines/rag/modules/retrieval/test_text_chunks_rerank_rag_module.py
@@ -2,10 +2,10 @@
 import pytest
 from griptape.drivers import CohereRerankDriver
 from griptape.engines.rag import RagContext
-from griptape.engines.rag.modules import TextRerankRagModule
+from griptape.engines.rag.modules import TextChunksRerankRagModule
 
 
-class TestTextRerankRagModule:
+class TestTextChunksRerankRagModule:
     @pytest.fixture
     def mock_client(self, mocker):
         mock_client = mocker.patch("cohere.Client").return_value
@@ -14,7 +14,7 @@ def mock_client(self, mocker):
         return mock_client
 
     def test_run(self, mock_client):
-        module = TextRerankRagModule(rerank_driver=CohereRerankDriver(api_key="api-key"))
+        module = TextChunksRerankRagModule(rerank_driver=CohereRerankDriver(api_key="api-key"))
         result = module.run(RagContext(query="test"))
 
         assert len(result) == 2