Refactor

1 files changed, 14 insertions, 0 deletions

diff --git a/rag/generator/prompt.py b/rag/generator/prompt.py
new file mode 100644
index 0000000..ed372c9
--- /dev/null
+++ b/rag/generator/prompt.py
@@ -0,0 +1,14 @@
+from dataclasses import dataclass
+from typing import List
+
+
+try:
+    from rag.retriever.vector import Document
+except ModuleNotFoundError:
+    from retriever.vector import Document
+
+
+@dataclass
+class Prompt:
+    query: str
+    documents: List[Document]