Compare commits

..

2 Commits

Author SHA1 Message Date
hetao dcfa5e9dee feat: config max_search_results for search engine via api 2025-05-18 12:22:41 +08:00
Henry Li 8abfd2d9ae feat: implement UI 2025-05-18 11:52:25 +08:00
5 changed files with 16 additions and 158 deletions
-1
View File
@@ -11,7 +11,6 @@ static/browser_history/*.gif
# Virtual environments
.venv
venv/
# Environment variables
.env
+11 -25
View File
@@ -7,8 +7,7 @@ Server script for running the DeerFlow API.
import argparse
import logging
import signal
import sys
import uvicorn
# Configure logging
@@ -19,17 +18,6 @@ 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")
@@ -62,18 +50,16 @@ if __name__ == "__main__":
# Determine reload setting
reload = False
# Command line arguments override defaults
if args.reload:
reload = True
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)
logger.info("Starting DeerFlow API server")
uvicorn.run(
"src.server:app",
host=args.host,
port=args.port,
reload=reload,
log_level=args.log_level,
)
+2 -2
View File
@@ -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.value:
if SELECTED_SEARCH_ENGINE == SearchEngine.TAVILY:
searched_content = LoggedTavilySearch(
max_results=configurable.max_search_results
).invoke(query)
).invoke({"query": query})
background_investigation_results = None
if isinstance(searched_content, list):
background_investigation_results = [
+3 -2
View File
@@ -44,12 +44,13 @@ def get_llm_by_type(
return llm
# Initialize LLMs for different purposes - now these will be cached
basic_llm = get_llm_by_type("basic")
# In the future, we will use reasoning_llm and vl_llm for different purposes
# reasoning_llm = get_llm_by_type("reasoning")
# vl_llm = get_llm_by_type("vision")
if __name__ == "__main__":
# Initialize LLMs for different purposes - now these will be cached
basic_llm = get_llm_by_type("basic")
print(basic_llm.invoke("Hello"))
-128
View File
@@ -1,128 +0,0 @@
import json
import pytest
from unittest.mock import patch, MagicMock
# 在这里 mock 掉 get_llm_by_type,避免 ValueError
with patch("src.llms.llm.get_llm_by_type", return_value=MagicMock()):
from langgraph.types import Command
from src.graph.nodes import background_investigation_node
from src.config import SearchEngine
from langchain_core.messages import HumanMessage
# Mock data
MOCK_SEARCH_RESULTS = [
{"title": "Test Title 1", "content": "Test Content 1"},
{"title": "Test Title 2", "content": "Test Content 2"},
]
@pytest.fixture
def mock_state():
return {
"messages": [HumanMessage(content="test query")],
"background_investigation_results": None,
}
@pytest.fixture
def mock_configurable():
mock = MagicMock()
mock.max_search_results = 5
return mock
@pytest.fixture
def mock_config():
# 你可以根据实际需要返回一个 MagicMock 或 dict
return MagicMock()
@pytest.fixture
def patch_config_from_runnable_config(mock_configurable):
with patch(
"src.graph.nodes.Configuration.from_runnable_config",
return_value=mock_configurable,
):
yield
@pytest.fixture
def mock_tavily_search():
with patch("src.graph.nodes.LoggedTavilySearch") as mock:
instance = mock.return_value
instance.invoke.return_value = [
{"title": "Test Title 1", "content": "Test Content 1"},
{"title": "Test Title 2", "content": "Test Content 2"},
]
yield mock
@pytest.fixture
def mock_web_search_tool():
with patch("src.graph.nodes.get_web_search_tool") as mock:
instance = mock.return_value
instance.invoke.return_value = [
{"title": "Test Title 1", "content": "Test Content 1"},
{"title": "Test Title 2", "content": "Test Content 2"},
]
yield mock
@pytest.mark.parametrize("search_engine", [SearchEngine.TAVILY.value, "other"])
def test_background_investigation_node_tavily(
mock_state,
mock_tavily_search,
mock_web_search_tool,
search_engine,
patch_config_from_runnable_config,
mock_config,
):
"""Test background_investigation_node with Tavily search engine"""
with patch("src.graph.nodes.SELECTED_SEARCH_ENGINE", search_engine):
result = background_investigation_node(mock_state, mock_config)
# Verify the result structure
assert isinstance(result, Command)
assert result.goto == "planner"
# Verify the update contains background_investigation_results
update = result.update
assert "background_investigation_results" in update
# Parse and verify the JSON content
results = json.loads(update["background_investigation_results"])
assert isinstance(results, list)
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"
else:
mock_web_search_tool.return_value.invoke.assert_called_once_with(
"test query"
)
assert len(results) == 2
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.value):
# Mock a malformed response
mock_tavily_search.return_value.invoke.return_value = "invalid response"
result = background_investigation_node(mock_state, mock_config)
# Verify the result structure
assert isinstance(result, Command)
assert result.goto == "planner"
# Verify the update contains background_investigation_results
update = result.update
assert "background_investigation_results" in update
# Parse and verify the JSON content
results = json.loads(update["background_investigation_results"])
assert results is None