From b1ff0c55422d7b0af2c379679b8721014ef36926 Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Tue, 18 Jun 2024 01:37:32 +0200
Subject: Wip rewrite

---
 rag/retriever/rerank/abstract.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

(limited to 'rag/retriever/rerank/abstract.py')

diff --git a/rag/retriever/rerank/abstract.py b/rag/retriever/rerank/abstract.py
index b96b70a..f32ee77 100644
--- a/rag/retriever/rerank/abstract.py
+++ b/rag/retriever/rerank/abstract.py
@@ -1,6 +1,8 @@
 from abc import abstractmethod
+from typing import List
 
-from rag.generator.prompt import Prompt
+from rag.memory import Message
+from rag.retriever.vector import Document
 
 
 class AbstractReranker(type):
@@ -13,5 +15,9 @@ class AbstractReranker(type):
         return cls._instances[cls]
 
     @abstractmethod
-    def rank(self, prompt: Prompt) -> Prompt:
-        return prompt
+    def rerank_documents(self, query: str, documents: List[Document]) -> List[Document]:
+        pass
+
+    @abstractmethod
+    def rerank_messages(self, query: str, messages: List[Message]) -> List[Message]:
+        pass
-- 
cgit v1.2.3-70-g09d2