mirror of
https://github.com/CopilotKit/CopilotKit.git
synced 2026-09-14 16:26:20 +08:00
Jordan Ritter
87 lines
3.9 KiB
Python
87 lines
3.9 KiB
Python
import asyncio
|
|
from copilotkit.langchain import copilotkit_emit_state
|
|
from datetime import datetime
|
|
from dotenv import load_dotenv
|
|
import json
|
|
from langchain_core.tools import tool
|
|
from pydantic import BaseModel, Field
|
|
from tavily import AsyncTavilyClient
|
|
from typing import List, Dict, Optional
|
|
from langchain_core.runnables import RunnableConfig
|
|
load_dotenv('.env')
|
|
tavily_client = AsyncTavilyClient()
|
|
|
|
# Add Tavily's arguments to enhance the web search tool's capabilities
|
|
class TavilyQuery(BaseModel):
|
|
query: str = Field(description="Web search query")
|
|
topic: str = Field(
|
|
description="Type of search, MUST be 'general' or 'news'. Choose 'news' ONLY when the company you searching is publicly traded and is likely to be featured on popular news")
|
|
days: int = Field(description="Number of days back to run 'news' search")
|
|
domains: Optional[List[str]] = Field(default=None,
|
|
description="List of domains to include in the research. Useful when trying to gather information from trusted and relevant domains")
|
|
|
|
|
|
# Define the args_schema for the tavily_search tool using a multi-query approach, enabling more precise queries for Tavily.
|
|
class TavilySearchInput(BaseModel):
|
|
sub_queries: List[TavilyQuery] = Field(description="Set of sub-queries that can be answered in isolation")
|
|
state: Optional[Dict] = Field(description="State of the research, will be provided later")
|
|
|
|
|
|
@tool("tavily_search", args_schema=TavilySearchInput, return_direct=True)
|
|
async def tavily_search(sub_queries: List[TavilyQuery], state):
|
|
"""Perform searches for each sub-query using the Tavily search tool concurrently."""
|
|
# Define a coroutine function to perform a single search with error handling
|
|
async def perform_search(itm, index):
|
|
try:
|
|
# Add date to the query as we need the most recent results
|
|
query_with_date = f"{itm.query} {datetime.now().strftime('%m-%Y')}"
|
|
# state["logs"][index]["message"] = f"🌐 Searched: '{query.query}'",
|
|
topic = itm.topic if itm.topic in ['general','news'] else "general"
|
|
tavily_response = await tavily_client.search(query=query_with_date, topic=topic, days=itm.days, max_results=10)
|
|
state["logs"][index]["done"] = True
|
|
tavily_response['results'] = [search for search in tavily_response['results'] if search['score'] > 0.45]
|
|
await copilotkit_emit_state(config, state)
|
|
return tavily_response['results']
|
|
except Exception as e:
|
|
# Handle any exceptions, log them, and return an empty list
|
|
print(f"Error occurred during search for query '{itm.query}': {str(e)}")
|
|
state["logs"][index]["done"] = True
|
|
await copilotkit_emit_state(config, state)
|
|
return []
|
|
|
|
config = RunnableConfig()
|
|
state["logs"] = state.get("logs", [])
|
|
# Log search queries
|
|
for query in sub_queries:
|
|
state["logs"].append({
|
|
"message": f"🌐 Searching the web: '{query.query}'",
|
|
"done": False
|
|
})
|
|
await copilotkit_emit_state(config, state)
|
|
|
|
# Run all the search tasks in parallel
|
|
search_tasks = [perform_search(query, i) for i, query in enumerate(sub_queries)]
|
|
search_responses = await asyncio.gather(*search_tasks)
|
|
|
|
|
|
# Combine the results from all the responses
|
|
tool_msg = "In search, found the following new documents:\n"
|
|
sources = state.get('sources', {})
|
|
for i, response in enumerate(search_responses):
|
|
for source in response:
|
|
if not sources or source['url'] not in sources:
|
|
sources[source['url']] = source
|
|
tool_msg += json.dumps(source)
|
|
|
|
state["logs"][i]["done"] = True
|
|
await copilotkit_emit_state(config, state)
|
|
|
|
for key,val in sources.items():
|
|
if not sources[key].get('title',None):
|
|
sources[key]['title'] = 'No Title, Invalid Link'
|
|
|
|
|
|
state['sources'] = sources
|
|
|
|
return state, tool_msg
|