diff --git a/agents/python/src/lib/mcp_integration.py b/agents/python/src/lib/mcp_integration.py index 6faf85d..73afe54 100644 --- a/agents/python/src/lib/mcp_integration.py +++ b/agents/python/src/lib/mcp_integration.py @@ -278,7 +278,8 @@ async def _call_mcp_tool(tool_name: str, arguments: Dict[str, Any]) -> Any: async def search_knowledge_base( query: str, count: int = 5, - search_effort: str = "fast" + search_effort: str = "fast", + source_indexes: Optional[List[str]] = None ) -> List[Dict[str, Any]]: """ Search the knowledge base via MCP server. @@ -287,19 +288,25 @@ async def search_knowledge_base( query: Search query count: Number of results to return search_effort: Search effort level ('fast', 'medium', or 'deep') + source_indexes: Optional list of source indexes to search in priority order + (e.g., ['tako'], ['web'], ['tako', 'web']) Returns: List of search results with metadata """ - result = await _call_mcp_tool("knowledge_search", { + args = { "query": query, "api_token": TAKO_API_TOKEN, "count": count, "search_effort": search_effort, "country_code": "US", "locale": "en-US" - }) + } + if source_indexes: + args["source_indexes"] = source_indexes + + result = await _call_mcp_tool("knowledge_search", args) if result and "results" in result: formatted_results = [] diff --git a/agents/python/src/lib/search.py b/agents/python/src/lib/search.py index ce1999c..2333ec6 100644 --- a/agents/python/src/lib/search.py +++ b/agents/python/src/lib/search.py @@ -164,14 +164,16 @@ async def search_node(state: AgentState, config: RunnableConfig): fallback_tasks = [] fallback_logs = [] - # Fallback web searches for fast questions + # Fallback Tako web searches for fast questions if fast_questions: - logger.info("No Tako results found, falling back to web searches for Tako questions") - fallback_web_queries = [q["question"] for q in fast_questions[:2]] - for query in fallback_web_queries: - fallback_logs.append(("web", query)) - state["logs"].append({"message": f"Fallback web search: {query}", "done": False}) - fallback_tasks.extend([async_tavily_search(query) for query in fallback_web_queries]) + logger.info("No Tako results found, falling back to Tako web index for questions") + for q_obj in fast_questions[:2]: + fallback_logs.append(("tako_web", q_obj["question"])) + state["logs"].append({"message": f"Tako web search: {q_obj['question']}", "done": False}) + fallback_tasks.extend([ + search_knowledge_base(q["question"], search_effort="fast", source_indexes=["web"]) + for q in fast_questions[:2] + ]) # Deep search for prediction market questions if prediction_market_questions: @@ -192,13 +194,8 @@ async def search_node(state: AgentState, config: RunnableConfig): for i, result in enumerate(fallback_results): task_type, _ = fallback_logs[i] if isinstance(result, Exception): - if task_type == "web": - search_results.append({"error": str(result)}) - else: - tako_results.append({"error": str(result)}) - elif task_type == "web": - search_results.append(result) - elif result: # deep Tako result + tako_results.append({"error": str(result)}) + elif result: # Tako result (web or deep) # Add resources immediately for streaming existing_urls = {r.get("url") for r in state["resources"]} existing_titles = {r.get("title", "").lower() for r in state["resources"] if r.get("resource_type") == "tako_chart"} diff --git a/src/app/Main.tsx b/src/app/Main.tsx index ccdee9c..0793f13 100644 --- a/src/app/Main.tsx +++ b/src/app/Main.tsx @@ -13,8 +13,8 @@ const CHAT_SUGGESTIONS = [ message: "Tell me about Vietnam's economy", }, { - title: "Nvidia vs Intel", - message: "Compare the performance of Nvidia and Intel over the last 10 years", + title: "Amazon vs Walmart", + message: "Compare the performance of Amazon and Walmart over the last 10 years", }, { title: "Commodities Prices", @@ -25,12 +25,16 @@ const CHAT_SUGGESTIONS = [ message: "What is the trend of global military spending since the Ukraine war started?", }, { - title: "Performance of AI companies", - message: "How has the performance of AI company's stock prices and funding evolved since 2020 (also include traffic trends)?", + title: "Semiconductor Companies", + message: "How has the performance of semiconductor companies (like NVIDIA, AMD, TSMC, Intel) evolved since 2020 in terms of stock prices, revenue, and market share?", }, { title: "Rent & Inflation", message: "How has the rent, inflation, and average wages trended in top US cities?", + }, + { + title: "Crime & Public Health", + message: "How do crime rates (homicides, violent crime) and public health issues (fentanyl overdoses, drug use) correlate with population changes in major US cities like San Francisco, Chicago, and New York?", } ];