turn self route using env (#121)

* turn self route using env * add SELF_ROUTE_ENABLED to env.example --------- Co-authored-by: Saisakul Chernbumroong <[email protected]>
uktrade · Feb 26, 2025 · 9d45501 · 9d45501
1 parent 9bf3b6d
commit 9d45501
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/.env.example b/.env.example
@@ -102,4 +102,5 @@ REDBOX_API_KEY = myapi
 # AUTHBROKER_CLIENT_SECRET=REPLACE_WITH_GITLAB_SECRET
 # AUTHBROKER_URL=https://sso.trade.gov.uk
 
-ENABLE_METADATA_EXTRACTION = True
+ENABLE_METADATA_EXTRACTION = True
+SELF_ROUTE_ENABLED = False
diff --git a/redbox-core/redbox/models/chain.py b/redbox-core/redbox/models/chain.py
@@ -5,6 +5,8 @@
 from typing import Annotated, Literal, NotRequired, Required, TypedDict, get_args, get_origin
 from uuid import UUID, uuid4
 
+import environ
+from dotenv import load_dotenv
 from langchain_core.documents import Document
 from langchain_core.messages import AnyMessage
 from langgraph.graph.message import add_messages
@@ -14,6 +16,9 @@
 from redbox.models import prompts
 from redbox.models.settings import ChatLLMBackend
 
+load_dotenv()
+env = environ.Env()
+
 
 class ChainChatMessage(TypedDict):
     role: Literal["user", "ai", "system"]
@@ -29,7 +34,7 @@ class AISettings(BaseModel):
 
     # Prompts and LangGraph settings
     max_document_tokens: int = 1_000_000
-    self_route_enabled: bool = False
+    self_route_enabled: bool = env.bool("SELF_ROUTE_ENABLED", default=False)
     map_max_concurrency: int = 128
     stuff_chunk_context_ratio: float = 0.75
     recursion_limit: int = 50