backend: Filter web search domains for Tavily (#713)

* wip

* Add domain filtering for Tavily search

* pr review fixes

Author: tianjing-li

src/backend/config/settings.py

@@ -146,7 +146,7 @@ class CompassSettings(BaseSettings, BaseModel):
     )
 
 
-class WebSearchSettings(BaseSettings, BaseModel):
+class TavilySearchSettings(BaseSettings, BaseModel):
     model_config = SETTINGS_CONFIG
     api_key: Optional[str] = Field(
         default=None, validation_alias=AliasChoices("TAVILY_API_KEY", "api_key")
@@ -185,7 +185,7 @@ class ToolSettings(BaseSettings, BaseModel):
     python_interpreter: Optional[PythonToolSettings] = Field(
         default=PythonToolSettings()
     )
-    web_search: Optional[WebSearchSettings] = Field(default=WebSearchSettings())
+    tavily: Optional[TavilySearchSettings] = Field(default=TavilySearchSettings())
     wolfram_alpha: Optional[WolframAlphaSettings] = Field(
         default=WolframAlphaSettings()
     )

src/backend/crud/agent_tool_metadata.py

@@ -25,7 +25,7 @@ def create_agent_tool_metadata(
 
 def get_agent_tool_metadata_by_id(
     db: Session, agent_tool_metadata_id: str
-) -> AgentToolMetadata:
+) -> AgentToolMetadata | None:
     """
     Get a agent tool metadata by its ID.
 
@@ -61,6 +61,35 @@ def get_all_agent_tool_metadata_by_agent_id(
     )
 
 
+def get_agent_tool_metadata(
+    db: Session,
+    agent_id: str,
+    tool_name: str,
+    user_id: str,
+) -> AgentToolMetadata | None:
+    """
+    Get a agent tool metadata.
+
+    Args:
+        db (Session): Database session.
+        agent_id (str): Agent ID.
+        tool_name (str): Tool name.
+        user_id (str): User ID.
+
+    Returns:
+        AgentToolMetadata: Agent tool metadata.
+    """
+    return (
+        db.query(AgentToolMetadata)
+        .filter(
+            AgentToolMetadata.agent_id == agent_id,
+            AgentToolMetadata.tool_name == tool_name,
+            AgentToolMetadata.user_id == user_id,
+        )
+        .first()
+    )
+
+
 def update_agent_tool_metadata(
     db: Session,
     agent_tool_metadata: AgentToolMetadata,

src/backend/routers/agent.py

@@ -384,6 +384,43 @@ async def update_agent(
     return agent
 
 
+@router.delete("/{agent_id}", response_model=DeleteAgent)
+async def delete_agent(
+    agent_id: str,
+    session: DBSessionDep,
+    ctx: Context = Depends(get_context),
+) -> DeleteAgent:
+    """
+    Delete an agent by ID.
+
+    Args:
+        agent_id (str): Agent ID.
+        session (DBSessionDep): Database session.
+        ctx (Context): Context object.
+
+    Returns:
+        DeleteAgent: Empty response.
+
+    Raises:
+        HTTPException: If the agent with the given ID is not found.
+    """
+    user_id = ctx.get_user_id()
+    ctx.with_event_type(MetricsMessageType.ASSISTANT_DELETED)
+    agent = validate_agent_exists(session, agent_id, user_id)
+    agent_schema = Agent.model_validate(agent)
+    ctx.with_agent(agent_schema)
+    ctx.with_metrics_agent(agent_to_metrics_agent(agent))
+
+    deleted = agent_crud.delete_agent(session, agent_id, user_id)
+    if not deleted:
+        raise HTTPException(status_code=401, detail="Could not delete Agent.")
+
+    return DeleteAgent()
+
+
+# Agent Tool Metadata endpoints
+
+
 async def handle_tool_metadata_update(
     agent: Agent,
     new_agent: Agent,
@@ -453,43 +490,6 @@ async def update_or_create_tool_metadata(
         create_agent_tool_metadata(session, agent.id, create_metadata_req, ctx)
 
 
-@router.delete("/{agent_id}", response_model=DeleteAgent)
-async def delete_agent(
-    agent_id: str,
-    session: DBSessionDep,
-    ctx: Context = Depends(get_context),
-) -> DeleteAgent:
-    """
-    Delete an agent by ID.
-
-    Args:
-        agent_id (str): Agent ID.
-        session (DBSessionDep): Database session.
-        ctx (Context): Context object.
-
-    Returns:
-        DeleteAgent: Empty response.
-
-    Raises:
-        HTTPException: If the agent with the given ID is not found.
-    """
-    user_id = ctx.get_user_id()
-    ctx.with_event_type(MetricsMessageType.ASSISTANT_DELETED)
-    agent = validate_agent_exists(session, agent_id, user_id)
-    agent_schema = Agent.model_validate(agent)
-    ctx.with_agent(agent_schema)
-    ctx.with_metrics_agent(agent_to_metrics_agent(agent))
-
-    deleted = agent_crud.delete_agent(session, agent_id, user_id)
-    if not deleted:
-        raise HTTPException(status_code=401, detail="Could not delete Agent.")
-
-    return DeleteAgent()
-
-
-# Tool Metadata Endpoints
-
-
 @router.get("/{agent_id}/tool-metadata", response_model=list[AgentToolMetadataPublic])
 async def list_agent_tool_metadata(
     agent_id: str, session: DBSessionDep, ctx: Context = Depends(get_context)
@@ -539,7 +539,7 @@ def create_agent_tool_metadata(
         ctx (Context): Context object.
 
     Returns:
-        AgentToolMetadata: Created agent tool metadata.
+        AgentToolMetadataPublic: Created agent tool metadata.
 
     Raises:
         HTTPException: If the agent tool metadata creation fails.

src/backend/tools/tavily.py

@@ -4,28 +4,65 @@
 from tavily import TavilyClient
 
 from backend.config.settings import Settings
+from backend.crud import agent_tool_metadata as agent_tool_metadata_crud
+from backend.database_models.database import DBSessionDep
 from backend.model_deployments.base import BaseDeployment
+from backend.schemas.context import Context
 from backend.tools.base import BaseTool
 
 
 class TavilyInternetSearch(BaseTool):
     NAME = "web_search"
-    TAVILY_API_KEY = Settings().tools.web_search.api_key
+    TAVILY_API_KEY = Settings().tools.tavily.api_key
+    POST_RERANK_MAX_RESULTS = 6
 
     def __init__(self):
         self.client = TavilyClient(api_key=self.TAVILY_API_KEY)
-        self.num_results = 6
 
     @classmethod
     def is_available(cls) -> bool:
         return cls.TAVILY_API_KEY is not None
 
+    def get_filtered_domains(self, session: DBSessionDep, ctx: Context) -> list[str]:
+        agent_id = ctx.get_agent_id()
+        user_id = ctx.get_user_id()
+
+        if not agent_id or not user_id:
+            # Default for Tavily is None
+            return []
+
+        agent_tool_metadata = agent_tool_metadata_crud.get_agent_tool_metadata(
+            db=session,
+            agent_id=agent_id,
+            tool_name=self.NAME,
+            user_id=user_id,
+        )
+
+        if not agent_tool_metadata:
+            # Default for Tavily is None
+            return []
+
+        return [
+            artifact["domain"]
+            for artifact in agent_tool_metadata.artifacts
+            if "domain" in artifact
+        ]
+
     async def call(
-        self, parameters: dict, ctx: Any, **kwargs: Any
+        self, parameters: dict, ctx: Any, session: DBSessionDep, **kwargs: Any
     ) -> List[Dict[str, Any]]:
+        # Gather search parameters
         query = parameters.get("query", "")
+        # Get domains set on Agent tool metadata artifacts
+        filter_domains = None
+        domains = self.get_filtered_domains(session, ctx)
+
+        if domains:
+            filter_domains = domains
+
+        # Do search
         result = self.client.search(
-            query=query, search_depth="advanced", include_raw_content=True
+            query=query, search_depth="advanced", include_raw_content=True, include_domains=filter_domains
         )
 
         if "results" not in result:
@@ -96,7 +133,7 @@ async def rerank_page_snippets(
                 seen_urls.append(result["url"])
                 reranked.append(result)
 
-        return reranked[: self.num_results]
+        return reranked[:self.POST_RERANK_MAX_RESULTS]
 
     def to_langchain_tool(self) -> TavilySearchResults:
         internet_search = TavilySearchResults()