Agente de web search construido e implementado(Tavily ejemplo de .env en readme)
This commit is contained in:
37
backend/app/agents/web_search/__init__.py
Normal file
37
backend/app/agents/web_search/__init__.py
Normal file
@@ -0,0 +1,37 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from .agent import agent
|
||||
from .models import WebSearchResponse, WebSearchState
|
||||
|
||||
|
||||
async def search_web(
|
||||
query: str,
|
||||
max_results: int = 5,
|
||||
include_raw_content: bool = False,
|
||||
) -> WebSearchResponse:
|
||||
"""
|
||||
Execute web search using Tavily MCP server.
|
||||
|
||||
Args:
|
||||
query: Search query string
|
||||
max_results: Maximum number of results to return (1-10)
|
||||
include_raw_content: Whether to include full content in results
|
||||
|
||||
Returns:
|
||||
WebSearchResponse with results and summary
|
||||
"""
|
||||
state = WebSearchState(
|
||||
user_query=query,
|
||||
max_results=max_results,
|
||||
include_raw_content=include_raw_content,
|
||||
)
|
||||
|
||||
prompt = (
|
||||
f"Search the web for: {query}\n\n"
|
||||
f"Return the top {max_results} most relevant results. "
|
||||
"Provide a concise summary of the key findings."
|
||||
)
|
||||
|
||||
# Ejecutar agente con Tavily API directa
|
||||
result = await agent.run(prompt, deps=state)
|
||||
return result.output
|
||||
72
backend/app/agents/web_search/agent.py
Normal file
72
backend/app/agents/web_search/agent.py
Normal file
@@ -0,0 +1,72 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic_ai import Agent, RunContext
|
||||
from pydantic_ai.models.openai import OpenAIChatModel
|
||||
from pydantic_ai.providers.azure import AzureProvider
|
||||
from tavily import TavilyClient
|
||||
|
||||
from app.core.config import settings
|
||||
|
||||
from .models import WebSearchResponse, WebSearchState, SearchResult
|
||||
|
||||
|
||||
provider = AzureProvider(
|
||||
azure_endpoint=settings.AZURE_OPENAI_ENDPOINT,
|
||||
api_version=settings.AZURE_OPENAI_API_VERSION,
|
||||
api_key=settings.AZURE_OPENAI_API_KEY,
|
||||
)
|
||||
model = OpenAIChatModel(model_name="gpt-4o", provider=provider)
|
||||
|
||||
|
||||
tavily_client = TavilyClient(api_key=settings.TAVILY_API_KEY)
|
||||
|
||||
agent = Agent(
|
||||
model=model,
|
||||
name="WebSearchAgent",
|
||||
deps_type=WebSearchState,
|
||||
output_type=WebSearchResponse,
|
||||
system_prompt=(
|
||||
"You are a web search assistant powered by Tavily. "
|
||||
"Use the tavily_search tool to find relevant, up-to-date information. "
|
||||
"Return a structured WebSearchResponse with results and a concise summary. "
|
||||
"Always cite your sources with URLs."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@agent.tool
|
||||
def tavily_search(ctx: RunContext[WebSearchState], query: str) -> list[SearchResult]:
|
||||
"""Search the web using Tavily API for up-to-date information."""
|
||||
response = tavily_client.search(
|
||||
query=query,
|
||||
max_results=ctx.deps.max_results,
|
||||
search_depth="basic",
|
||||
include_raw_content=ctx.deps.include_raw_content,
|
||||
)
|
||||
|
||||
results = []
|
||||
for item in response.get("results", []):
|
||||
results.append(
|
||||
SearchResult(
|
||||
title=item.get("title", ""),
|
||||
url=item.get("url", ""),
|
||||
content=item.get("content", ""),
|
||||
score=item.get("score"),
|
||||
)
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
@agent.output_validator
|
||||
def finalize_response(
|
||||
ctx: RunContext[WebSearchState],
|
||||
response: WebSearchResponse,
|
||||
) -> WebSearchResponse:
|
||||
"""Post-process and validate the search response"""
|
||||
return response.model_copy(
|
||||
update={
|
||||
"query": ctx.deps.user_query,
|
||||
"total_results": len(response.results),
|
||||
}
|
||||
)
|
||||
29
backend/app/agents/web_search/models.py
Normal file
29
backend/app/agents/web_search/models.py
Normal file
@@ -0,0 +1,29 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class SearchResult(BaseModel):
|
||||
"""Individual search result from web search"""
|
||||
|
||||
title: str
|
||||
url: str
|
||||
content: str
|
||||
score: float | None = None
|
||||
|
||||
|
||||
class WebSearchState(BaseModel):
|
||||
"""State passed to agent tools via deps"""
|
||||
|
||||
user_query: str
|
||||
max_results: int = Field(default=5, ge=1, le=10)
|
||||
include_raw_content: bool = False
|
||||
|
||||
|
||||
class WebSearchResponse(BaseModel):
|
||||
"""Structured response from the web search agent"""
|
||||
|
||||
query: str
|
||||
results: list[SearchResult]
|
||||
summary: str
|
||||
total_results: int
|
||||
Reference in New Issue
Block a user