feat: support images in the search results

This commit is contained in:
Li Xin 2025-04-19 09:57:02 +08:00
parent 3733d346d7
commit 6ffe46e39b
8 changed files with 365 additions and 35 deletions

View File

@ -7,10 +7,12 @@ You are a professional reporter responsible for writing clear, comprehensive rep
# Role
You should act as an objective and analytical reporter who:
- Uses the same language as the initial question
- Presents facts accurately and impartially
- Organizes information logically
- Highlights key findings and insights
- Uses clear and concise language
- To enrich the report, includes relevant images from the previous steps
- Relies strictly on provided information
- Never fabricates or assumes information
- Clearly distinguishes between facts and analysis
@ -37,6 +39,7 @@ Structure your report in the following format:
- Include relevant subsections as needed
- Present information in a structured, easy-to-follow manner
- Highlight unexpected or particularly noteworthy details
- **Including images from the previous steps in the report is very helpful.**
5. **Survey Note** (for more comprehensive reports)
- A more detailed, academic-style analysis
@ -64,6 +67,7 @@ Structure your report in the following format:
- Use proper markdown syntax
- Include headers for sections
- Prioritize using Markdown tables for data presentation and comparison
- **Including images from the previous steps in the report is very helpful.**
- Use tables whenever presenting comparative data, statistics, features, or options
- Structure tables with clear headers and aligned columns
- Add emphasis for important points
@ -87,7 +91,7 @@ Structure your report in the following format:
- Keep tables concise and focused on key information
- Use proper Markdown table syntax:
```
```markdown
| Header 1 | Header 2 | Header 3 |
|----------|----------|----------|
| Data 1 | Data 2 | Data 3 |
@ -96,7 +100,7 @@ Structure your report in the following format:
- For feature comparison tables, use this format:
```
```markdown
| Feature/Option | Description | Pros | Cons |
|----------------|-------------|------|------|
| Feature 1 | Description | Pros | Cons |
@ -111,3 +115,5 @@ Structure your report in the following format:
- Place all citations in the "Key Citations" section at the end, not inline in the text
- For each citation, use the format: `- [Source Title](URL)`
- Include an empty line between each citation for better readability
- Include images using `![Image Description](image_url)` in a separate section.
- The included images should **only** be from the information gathered **from the previous steps**. **Never** include images that are not from the previous steps.

View File

@ -17,14 +17,15 @@ You are dedicated to conducting thorough investigations and providing comprehens
- Combine the information gathered from the search results and the crawled content.
- Ensure the response is clear, concise, and directly addresses the problem.
- Track and attribute all information sources with their respective URLs for proper citation.
- Including images from the search results or the crawled content in the report is very helpful.
# Output Format
- Provide a structured response in markdown format.
- Include the following sections:
- **Problem Statement**: Restate the problem for clarity.
- **Search Results**: Summarize the key findings from the **web_search_tool** search. Track the sources of information but DO NOT include inline citations in the text.
- **Crawled Content**: Summarize the key findings from the **crawl_tool**. Track the sources of information but DO NOT include inline citations in the text.
- **Search Results**: Summarize the key findings from the **web_search_tool** search. Track the sources of information but DO NOT include inline citations in the text. Include images if relevant.
- **Crawled Content**: Summarize the key findings from the **crawl_tool**. Track the sources of information but DO NOT include inline citations in the text. Include images if relevant.
- **Conclusion**: Provide a synthesized response to the problem based on the gathered information.
- **References**: List all sources used with their complete URLs in link reference format at the end of the document. Make sure to include an empty line between each reference for better readability. Use this format for each reference:
```markdown
@ -46,4 +47,6 @@ You are dedicated to conducting thorough investigations and providing comprehens
- Only invoke `crawl_tool` when essential information cannot be obtained from search results alone.
- Always include source attribution for all information. This is critical for the final report's citations.
- When presenting information from multiple sources, clearly indicate which source each piece of information comes from.
- Include images using `![Image Description](image_url)` in a separate section.
- The included images should **only** be from the information gathered **from the search results or the crawled content**. **Never** include images that are not from the search results or the crawled content.
- Always use the same language as the initial question.

View File

@ -1,25 +1,31 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import json
import logging
import os
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.tools import DuckDuckGoSearchResults
from langchain_community.tools import BraveSearch
from langchain_community.tools import BraveSearch, DuckDuckGoSearchResults
from langchain_community.tools.arxiv import ArxivQueryRun
from langchain_community.utilities import ArxivAPIWrapper, BraveSearchWrapper
from src.config import SEARCH_MAX_RESULTS
from src.tools.tavily_search.tavily_search_results_with_images import (
TavilySearchResultsWithImages,
)
from .decorators import create_logged_tool
logger = logging.getLogger(__name__)
LoggedTavilySearch = create_logged_tool(TavilySearchResults)
LoggedTavilySearch = create_logged_tool(TavilySearchResultsWithImages)
tavily_search_tool = LoggedTavilySearch(
name="web_search",
max_results=SEARCH_MAX_RESULTS,
include_raw_content=True,
include_images=True,
include_image_descriptions=True,
)
LoggedDuckDuckGoSearch = create_logged_tool(DuckDuckGoSearchResults)
@ -45,3 +51,7 @@ arxiv_search_tool = LoggedArxivSearch(
load_all_available_meta=True,
),
)
if __name__ == "__main__":
results = tavily_search_tool.invoke("cute panda")
print(json.dumps(results, indent=2, ensure_ascii=False))

View File

@ -0,0 +1,4 @@
from .tavily_search_api_wrapper import EnhancedTavilySearchAPIWrapper
from .tavily_search_results_with_images import TavilySearchResultsWithImages
__all__ = ["EnhancedTavilySearchAPIWrapper", "TavilySearchResultsWithImages"]

View File

@ -0,0 +1,115 @@
import json
from typing import Dict, List, Optional
import aiohttp
import requests
from langchain_community.utilities.tavily_search import TAVILY_API_URL
from langchain_community.utilities.tavily_search import (
TavilySearchAPIWrapper as OriginalTavilySearchAPIWrapper,
)
class EnhancedTavilySearchAPIWrapper(OriginalTavilySearchAPIWrapper):
def raw_results(
self,
query: str,
max_results: Optional[int] = 5,
search_depth: Optional[str] = "advanced",
include_domains: Optional[List[str]] = [],
exclude_domains: Optional[List[str]] = [],
include_answer: Optional[bool] = False,
include_raw_content: Optional[bool] = False,
include_images: Optional[bool] = False,
include_image_descriptions: Optional[bool] = False,
) -> Dict:
params = {
"api_key": self.tavily_api_key.get_secret_value(),
"query": query,
"max_results": max_results,
"search_depth": search_depth,
"include_domains": include_domains,
"exclude_domains": exclude_domains,
"include_answer": include_answer,
"include_raw_content": include_raw_content,
"include_images": include_images,
"include_image_descriptions": include_image_descriptions,
}
response = requests.post(
# type: ignore
f"{TAVILY_API_URL}/search",
json=params,
)
response.raise_for_status()
return response.json()
async def raw_results_async(
self,
query: str,
max_results: Optional[int] = 5,
search_depth: Optional[str] = "advanced",
include_domains: Optional[List[str]] = [],
exclude_domains: Optional[List[str]] = [],
include_answer: Optional[bool] = False,
include_raw_content: Optional[bool] = False,
include_images: Optional[bool] = False,
include_image_descriptions: Optional[bool] = False,
) -> Dict:
"""Get results from the Tavily Search API asynchronously."""
# Function to perform the API call
async def fetch() -> str:
params = {
"api_key": self.tavily_api_key.get_secret_value(),
"query": query,
"max_results": max_results,
"search_depth": search_depth,
"include_domains": include_domains,
"exclude_domains": exclude_domains,
"include_answer": include_answer,
"include_raw_content": include_raw_content,
"include_images": include_images,
"include_image_descriptions": include_image_descriptions,
}
async with aiohttp.ClientSession() as session:
async with session.post(f"{TAVILY_API_URL}/search", json=params) as res:
if res.status == 200:
data = await res.text()
return data
else:
raise Exception(f"Error {res.status}: {res.reason}")
results_json_str = await fetch()
return json.loads(results_json_str)
def clean_results_with_images(
self, raw_results: Dict[str, List[Dict]]
) -> List[Dict]:
results = raw_results["results"]
"""Clean results from Tavily Search API."""
clean_results = []
for result in results:
clean_result = {
"type": "page",
"title": result["title"],
"url": result["url"],
"content": result["content"],
"score": result["score"],
}
if raw_content := result.get("raw_content"):
clean_result["raw_content"] = raw_content
clean_results.append(clean_result)
images = raw_results["images"]
for image in images:
clean_result = {
"type": "image",
"image_url": image["url"],
"image_description": image["description"],
}
clean_results.append(clean_result)
return clean_results
if __name__ == "__main__":
wrapper = EnhancedTavilySearchAPIWrapper()
results = wrapper.raw_results("cute panda", include_images=True)
print(json.dumps(results, indent=2, ensure_ascii=False))

View File

@ -0,0 +1,148 @@
import json
from typing import Dict, List, Optional, Tuple, Union
from langchain.callbacks.manager import (
AsyncCallbackManagerForToolRun,
CallbackManagerForToolRun,
)
from langchain_community.tools.tavily_search.tool import TavilySearchResults
from pydantic import Field
from src.tools.tavily_search.tavily_search_api_wrapper import (
EnhancedTavilySearchAPIWrapper,
)
class TavilySearchResultsWithImages(TavilySearchResults): # type: ignore[override, override]
"""Tool that queries the Tavily Search API and gets back json.
Setup:
Install ``langchain-openai`` and ``tavily-python``, and set environment variable ``TAVILY_API_KEY``.
.. code-block:: bash
pip install -U langchain-community tavily-python
export TAVILY_API_KEY="your-api-key"
Instantiate:
.. code-block:: python
from langchain_community.tools import TavilySearchResults
tool = TavilySearchResults(
max_results=5,
include_answer=True,
include_raw_content=True,
include_images=True,
include_image_descriptions=True,
# search_depth="advanced",
# include_domains = []
# exclude_domains = []
)
Invoke directly with args:
.. code-block:: python
tool.invoke({'query': 'who won the last french open'})
.. code-block:: json
{
"url": "https://www.nytimes.com...",
"content": "Novak Djokovic won the last French Open by beating Casper Ruud ..."
}
Invoke with tool call:
.. code-block:: python
tool.invoke({"args": {'query': 'who won the last french open'}, "type": "tool_call", "id": "foo", "name": "tavily"})
.. code-block:: python
ToolMessage(
content='{ "url": "https://www.nytimes.com...", "content": "Novak Djokovic won the last French Open by beating Casper Ruud ..." }',
artifact={
'query': 'who won the last french open',
'follow_up_questions': None,
'answer': 'Novak ...',
'images': [
'https://www.amny.com/wp-content/uploads/2023/06/AP23162622181176-1200x800.jpg',
...
],
'results': [
{
'title': 'Djokovic ...',
'url': 'https://www.nytimes.com...',
'content': "Novak...",
'score': 0.99505633,
'raw_content': 'Tennis\nNovak ...'
},
...
],
'response_time': 2.92
},
tool_call_id='1',
name='tavily_search_results_json',
)
""" # noqa: E501
include_image_descriptions: bool = False
"""Include a image descriptions in the response.
Default is False.
"""
api_wrapper: EnhancedTavilySearchAPIWrapper = Field(default_factory=EnhancedTavilySearchAPIWrapper) # type: ignore[arg-type]
def _run(
self,
query: str,
run_manager: Optional[CallbackManagerForToolRun] = None,
) -> Tuple[Union[List[Dict[str, str]], str], Dict]:
"""Use the tool."""
# TODO: remove try/except, should be handled by BaseTool
try:
raw_results = self.api_wrapper.raw_results(
query,
self.max_results,
self.search_depth,
self.include_domains,
self.exclude_domains,
self.include_answer,
self.include_raw_content,
self.include_images,
self.include_image_descriptions,
)
except Exception as e:
return repr(e), {}
cleaned_results = self.api_wrapper.clean_results_with_images(raw_results)
print("sync", json.dumps(cleaned_results, indent=2, ensure_ascii=False))
return cleaned_results, raw_results
async def _arun(
self,
query: str,
run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
) -> Tuple[Union[List[Dict[str, str]], str], Dict]:
"""Use the tool asynchronously."""
try:
raw_results = await self.api_wrapper.raw_results_async(
query,
self.max_results,
self.search_depth,
self.include_domains,
self.exclude_domains,
self.include_answer,
self.include_raw_content,
self.include_images,
self.include_image_descriptions,
)
except Exception as e:
return repr(e), {}
cleaned_results = self.api_wrapper.clean_results_with_images(raw_results)
print("async", json.dumps(cleaned_results, indent=2, ensure_ascii=False))
return cleaned_results, raw_results

View File

@ -94,26 +94,46 @@ function ActivityListItem({ messageId }: { messageId: string }) {
}
const __pageCache = new LRUCache<string, string>({ max: 100 });
type SearchResult =
| {
type: "page";
title: string;
url: string;
content: string;
}
| {
type: "image";
image_url: string;
image_description: string;
};
function WebSearchToolCall({ toolCall }: { toolCall: ToolCallRuntime }) {
const searchResults = useMemo<
{ title: string; url: string; content: string }[]
>(() => {
let results: { title: string; url: string; content: string }[] | undefined =
undefined;
const searchResults = useMemo<SearchResult[]>(() => {
let results: SearchResult[] | undefined = undefined;
try {
results = toolCall.result ? parse(toolCall.result) : undefined;
} catch {
results = undefined;
}
if (Array.isArray(results)) {
results.forEach((result: { url: string; title: string }) => {
__pageCache.set(result.url, result.title);
results.forEach((result) => {
if (result.type === "page") {
__pageCache.set(result.url, result.title);
}
});
} else {
results = [];
}
console.info(results);
return results;
}, [toolCall.result]);
const pageResults = useMemo(
() => searchResults?.filter((result) => result.type === "page"),
[searchResults],
);
const imageResults = useMemo(
() => searchResults?.filter((result) => result.type === "image"),
[searchResults],
);
return (
<section>
<div className="font-medium italic">
@ -128,30 +148,48 @@ function WebSearchToolCall({ toolCall }: { toolCall: ToolCallRuntime }) {
</span>
</RainbowText>
</div>
{searchResults && (
<div className="px-5">
<div className="px-5">
{pageResults && (
<ul className="mt-2 flex flex-wrap gap-4">
{searchResults.map((searchResult, i) => (
<motion.li
key={`search-result-${i}`}
className="text-muted-foreground flex max-w-40 gap-2 rounded-md bg-slate-100 px-2 py-1 text-sm"
initial={{ opacity: 0, y: 10, scale: 0.66 }}
animate={{ opacity: 1, y: 0, scale: 1 }}
transition={{
duration: 0.2,
delay: i * 0.1,
ease: "easeOut",
}}
>
<FavIcon url={searchResult.url} title={searchResult.title} />
<a href={searchResult.url} target="_blank">
{searchResult.title}
{pageResults
.filter((result) => result.type === "page")
.map((searchResult, i) => (
<motion.li
key={`search-result-${i}`}
className="text-muted-foreground flex max-w-40 gap-2 rounded-md bg-slate-100 px-2 py-1 text-sm"
initial={{ opacity: 0, y: 10, scale: 0.66 }}
animate={{ opacity: 1, y: 0, scale: 1 }}
transition={{
duration: 0.2,
delay: i * 0.1,
ease: "easeOut",
}}
>
<FavIcon url={searchResult.url} title={searchResult.title} />
<a href={searchResult.url} target="_blank">
{searchResult.title}
</a>
</motion.li>
))}
{imageResults.map((searchResult, i) => (
<li key={`search-result-${i}`}>
<a
className="flex flex-col gap-2 opacity-75 transition-opacity duration-300 hover:opacity-100"
href={searchResult.image_url}
target="_blank"
>
<div
className="h-40 w-40 max-w-full rounded-md bg-slate-100 bg-cover bg-center bg-no-repeat"
style={{
backgroundImage: `url(${searchResult.image_url})`,
}}
/>
</a>
</motion.li>
</li>
))}
</ul>
</div>
)}
)}
</div>
</section>
);
}

View File

@ -190,6 +190,12 @@ textarea {
@apply list-decimal pl-4;
}
img {
display: block;
max-width: 100%;
margin: 0 auto;
}
table {
@apply w-full;
table-layout: fixed;