wehub-resource-sync
50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
from langchain_core.tools import tool
|
|
from pydantic import BaseModel, Field
|
|
from tavily import AsyncTavilyClient
|
|
from typing import List, Optional, Dict
|
|
from copilotkit.langchain import copilotkit_emit_state
|
|
from langchain_core.runnables import RunnableConfig
|
|
|
|
tavily_client = AsyncTavilyClient()
|
|
|
|
|
|
class TavilyExtractInput(BaseModel):
|
|
urls: List[str] = Field(
|
|
description="List of a single or several URLs for extracting raw content to gather additional information"
|
|
)
|
|
state: Optional[Dict] = Field(description="State of the research")
|
|
|
|
|
|
@tool("tavily_extract", args_schema=TavilyExtractInput, return_direct=True)
|
|
async def tavily_extract(urls, state):
|
|
"""Perform full scrape to a provided list of urls."""
|
|
|
|
try:
|
|
response = await tavily_client.extract(urls=urls)
|
|
results = response["results"]
|
|
# Match and add raw_content to urls in state
|
|
tool_msg = "Extracted raw content to gather additional information from the following sources:\n"
|
|
for itm in results:
|
|
url = itm["url"]
|
|
raw_content = itm["raw_content"]
|
|
if url in state["sources"]:
|
|
state["sources"][url]["raw_content"] = raw_content
|
|
else:
|
|
state["sources"][url] = {"raw_content": raw_content}
|
|
tool_msg += f"{url}\n"
|
|
|
|
config = RunnableConfig()
|
|
state["logs"] = state.get("logs", [])
|
|
state["logs"].append(
|
|
{
|
|
"message": "🚀 Extracting additional content from valuable sources",
|
|
"done": True,
|
|
}
|
|
)
|
|
await copilotkit_emit_state(config, state)
|
|
return state, tool_msg
|
|
|
|
except Exception as e:
|
|
print(f"Error occurred during extract: {str(e)}")
|
|
return state, ""
|