mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-05-20 07:01:03 +00:00
Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| d9aa92afaa | |||
| 29be360954 | |||
| 3ed70e11d5 | |||
| 55ce399969 |
@@ -11,6 +11,7 @@ static/browser_history/*.gif
|
||||
|
||||
# Virtual environments
|
||||
.venv
|
||||
venv/
|
||||
|
||||
# Environment variables
|
||||
.env
|
||||
|
||||
@@ -7,7 +7,8 @@ Server script for running the DeerFlow API.
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
|
||||
import signal
|
||||
import sys
|
||||
import uvicorn
|
||||
|
||||
# Configure logging
|
||||
@@ -18,6 +19,17 @@ logging.basicConfig(
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def handle_shutdown(signum, frame):
|
||||
"""Handle graceful shutdown on SIGTERM/SIGINT"""
|
||||
logger.info("Received shutdown signal. Starting graceful shutdown...")
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
# Register signal handlers
|
||||
signal.signal(signal.SIGTERM, handle_shutdown)
|
||||
signal.signal(signal.SIGINT, handle_shutdown)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Parse command line arguments
|
||||
parser = argparse.ArgumentParser(description="Run the DeerFlow API server")
|
||||
@@ -50,16 +62,18 @@ if __name__ == "__main__":
|
||||
|
||||
# Determine reload setting
|
||||
reload = False
|
||||
|
||||
# Command line arguments override defaults
|
||||
if args.reload:
|
||||
reload = True
|
||||
|
||||
logger.info("Starting DeerFlow API server")
|
||||
uvicorn.run(
|
||||
"src.server:app",
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
reload=reload,
|
||||
log_level=args.log_level,
|
||||
)
|
||||
try:
|
||||
logger.info(f"Starting DeerFlow API server on {args.host}:{args.port}")
|
||||
uvicorn.run(
|
||||
"src.server:app",
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
reload=reload,
|
||||
log_level=args.log_level,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to start server: {str(e)}")
|
||||
sys.exit(1)
|
||||
|
||||
+2
-2
@@ -50,10 +50,10 @@ def background_investigation_node(
|
||||
logger.info("background investigation node is running.")
|
||||
configurable = Configuration.from_runnable_config(config)
|
||||
query = state["messages"][-1].content
|
||||
if SELECTED_SEARCH_ENGINE == SearchEngine.TAVILY:
|
||||
if SELECTED_SEARCH_ENGINE == SearchEngine.TAVILY.value:
|
||||
searched_content = LoggedTavilySearch(
|
||||
max_results=configurable.max_search_results
|
||||
).invoke({"query": query})
|
||||
).invoke(query)
|
||||
background_investigation_results = None
|
||||
if isinstance(searched_content, list):
|
||||
background_investigation_results = [
|
||||
|
||||
@@ -68,7 +68,7 @@ def mock_web_search_tool():
|
||||
yield mock
|
||||
|
||||
|
||||
@pytest.mark.parametrize("search_engine", [SearchEngine.TAVILY, "other"])
|
||||
@pytest.mark.parametrize("search_engine", [SearchEngine.TAVILY.value, "other"])
|
||||
def test_background_investigation_node_tavily(
|
||||
mock_state,
|
||||
mock_tavily_search,
|
||||
@@ -93,10 +93,8 @@ def test_background_investigation_node_tavily(
|
||||
results = json.loads(update["background_investigation_results"])
|
||||
assert isinstance(results, list)
|
||||
|
||||
if search_engine == SearchEngine.TAVILY:
|
||||
mock_tavily_search.return_value.invoke.assert_called_once_with(
|
||||
{"query": "test query"}
|
||||
)
|
||||
if search_engine == SearchEngine.TAVILY.value:
|
||||
mock_tavily_search.return_value.invoke.assert_called_once_with("test query")
|
||||
assert len(results) == 2
|
||||
assert results[0]["title"] == "Test Title 1"
|
||||
assert results[0]["content"] == "Test Content 1"
|
||||
@@ -111,7 +109,7 @@ def test_background_investigation_node_malformed_response(
|
||||
mock_state, mock_tavily_search, patch_config_from_runnable_config, mock_config
|
||||
):
|
||||
"""Test background_investigation_node with malformed Tavily response"""
|
||||
with patch("src.graph.nodes.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY):
|
||||
with patch("src.graph.nodes.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value):
|
||||
# Mock a malformed response
|
||||
mock_tavily_search.return_value.invoke.return_value = "invalid response"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user