Fix imports

3 files changed, 4 insertions, 10 deletions

diff --git a/rag/generator/ollama.py b/rag/generator/ollama.py
index ec6a55f..1e968f6 100644
--- a/rag/generator/ollama.py
+++ b/rag/generator/ollama.py
@@ -7,10 +7,7 @@ from loguru import logger as log
 from .prompt import Prompt
 from .abstract import AbstractGenerator
 
-try:
-    from rag.retriever.vector import Document
-except ModuleNotFoundError:
-    from retriever.vector import Document
+from rag.retriever.vector import Document
 
 SYSTEM_PROMPT = (
     "# System Preamble"
diff --git a/rag/generator/prompt.py b/rag/generator/prompt.py
index ed372c9..059592a 100644
--- a/rag/generator/prompt.py
+++ b/rag/generator/prompt.py
@@ -2,10 +2,7 @@ from dataclasses import dataclass
 from typing import List
 
 
-try:
-    from rag.retriever.vector import Document
-except ModuleNotFoundError:
-    from retriever.vector import Document
+from rag.retriever.vector import Document
 
 
 @dataclass
diff --git a/rag/ui.py b/rag/ui.py
index 37ab664..18e6bd1 100644
--- a/rag/ui.py
+++ b/rag/ui.py
@@ -4,8 +4,8 @@ from langchain_community.document_loaders.blob_loaders import Blob
 
 from dotenv import load_dotenv
 from rag.generator import get_generator, MODELS
-from generator.prompt import Prompt
-from retriever.retriever import Retriever
+from rag.generator.prompt import Prompt
+from rag.retriever.retriever import Retriever
 
 
 if __name__ == "__main__":