Browse Source

enh: enable configuration for tavily extract depth

Luke 4 months ago
parent
commit
7917128ed3

+ 6 - 0
backend/open_webui/config.py

@@ -1950,6 +1950,12 @@ TAVILY_API_KEY = PersistentConfig(
     os.getenv("TAVILY_API_KEY", ""),
     os.getenv("TAVILY_API_KEY", ""),
 )
 )
 
 
+TAVILY_EXTRACT_DEPTH = PersistentConfig(
+    "TAVILY_EXTRACT_DEPTH",
+    "rag.web.search.tavily_extract_depth",
+    int(os.getenv("TAVILY_EXTRACT_DEPTH", "basic")),
+)
+
 JINA_API_KEY = PersistentConfig(
 JINA_API_KEY = PersistentConfig(
     "JINA_API_KEY",
     "JINA_API_KEY",
     "rag.web.search.jina_api_key",
     "rag.web.search.jina_api_key",

+ 2 - 0
backend/open_webui/main.py

@@ -212,6 +212,7 @@ from open_webui.config import (
     SERPSTACK_API_KEY,
     SERPSTACK_API_KEY,
     SERPSTACK_HTTPS,
     SERPSTACK_HTTPS,
     TAVILY_API_KEY,
     TAVILY_API_KEY,
+    TAVILY_EXTRACT_DEPTH,
     BING_SEARCH_V7_ENDPOINT,
     BING_SEARCH_V7_ENDPOINT,
     BING_SEARCH_V7_SUBSCRIPTION_KEY,
     BING_SEARCH_V7_SUBSCRIPTION_KEY,
     BRAVE_SEARCH_API_KEY,
     BRAVE_SEARCH_API_KEY,
@@ -614,6 +615,7 @@ app.state.config.RAG_WEB_SEARCH_TRUST_ENV = RAG_WEB_SEARCH_TRUST_ENV
 app.state.config.PLAYWRIGHT_WS_URI = PLAYWRIGHT_WS_URI
 app.state.config.PLAYWRIGHT_WS_URI = PLAYWRIGHT_WS_URI
 app.state.config.FIRECRAWL_API_BASE_URL = FIRECRAWL_API_BASE_URL
 app.state.config.FIRECRAWL_API_BASE_URL = FIRECRAWL_API_BASE_URL
 app.state.config.FIRECRAWL_API_KEY = FIRECRAWL_API_KEY
 app.state.config.FIRECRAWL_API_KEY = FIRECRAWL_API_KEY
+app.state.config.TAVILY_EXTRACT_DEPTH = TAVILY_EXTRACT_DEPTH
 
 
 app.state.EMBEDDING_FUNCTION = None
 app.state.EMBEDDING_FUNCTION = None
 app.state.ef = None
 app.state.ef = None

+ 2 - 0
backend/open_webui/retrieval/web/utils.py

@@ -33,6 +33,7 @@ from open_webui.config import (
     FIRECRAWL_API_BASE_URL,
     FIRECRAWL_API_BASE_URL,
     FIRECRAWL_API_KEY,
     FIRECRAWL_API_KEY,
     TAVILY_API_KEY,
     TAVILY_API_KEY,
+    TAVILY_EXTRACT_DEPTH,
 )
 )
 from open_webui.env import SRC_LOG_LEVELS
 from open_webui.env import SRC_LOG_LEVELS
 
 
@@ -613,6 +614,7 @@ def get_web_loader(
 
 
     if RAG_WEB_LOADER_ENGINE.value == "tavily":
     if RAG_WEB_LOADER_ENGINE.value == "tavily":
         web_loader_args["api_key"] = TAVILY_API_KEY.value
         web_loader_args["api_key"] = TAVILY_API_KEY.value
+        web_loader_args["extract_depth"] = TAVILY_EXTRACT_DEPTH.value
 
 
     # Create the appropriate WebLoader based on the configuration
     # Create the appropriate WebLoader based on the configuration
     WebLoaderClass = RAG_WEB_LOADER_ENGINES[RAG_WEB_LOADER_ENGINE.value]
     WebLoaderClass = RAG_WEB_LOADER_ENGINES[RAG_WEB_LOADER_ENGINE.value]