Refactor

1 files changed, 15 insertions, 0 deletions

diff --git a/rag/generator/__init__.py b/rag/generator/__init__.py
new file mode 100644
index 0000000..7da603c
--- /dev/null
+++ b/rag/generator/__init__.py
@@ -0,0 +1,15 @@
+from typing import Type
+
+from .abstract import AbstractGenerator
+from .ollama import Ollama
+from .cohere import Cohere
+
+
+def get_generator(model: str) -> Type[AbstractGenerator]:
+    match model:
+        case "ollama":
+            return Ollama()
+        case "cohere":
+            return Cohere()
+        case _:
+            exit(1)