Use env variable to switch retrievers (#7)

langchain-ai · Oct 9, 2023 · 9b0beed · 9b0beed
1 parent e07a246
commit 9b0beed
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/main.py b/main.py
@@ -1,4 +1,5 @@
 """Main entrypoint for the app."""
+import os
 import asyncio
 from operator import itemgetter
 from typing import Dict, List, Optional, Sequence
@@ -151,8 +152,9 @@ def _get_relevant_documents(
 
 
 def get_base_retriever():
-    # return TavilySearchAPIRetriever(k=6, include_raw_content=True, include_images=True)
-    return BackupRetriever()
+    if (os.environ.get("USE_BACKUP", False)):
+        return BackupRetriever()
+    return TavilySearchAPIRetriever(k=6, include_raw_content=True, include_images=True)
 
 
 def _get_retriever():