foundry-mcp 0.7.0__py3-none-any.whl → 0.8.10__py3-none-any.whl

foundry_mcp/core/research/providers/tavily.py

@@ -0,0 +1,383 @@
+"""Tavily search provider for web search.
+
+This module implements TavilySearchProvider, which wraps the Tavily Search API
+to provide web search capabilities for the deep research workflow.
+
+Tavily API documentation: https://docs.tavily.com/
+
+Example usage:
+    provider = TavilySearchProvider(api_key="tvly-...")
+    sources = await provider.search("machine learning trends", max_results=5)
+"""
+
+import asyncio
+import logging
+import os
+from datetime import datetime
+from typing import Any, Optional
+
+import httpx
+
+from foundry_mcp.core.research.models import ResearchSource, SourceType
+from foundry_mcp.core.research.providers.base import (
+    AuthenticationError,
+    RateLimitError,
+    SearchProvider,
+    SearchProviderError,
+    SearchResult,
+)
+
+logger = logging.getLogger(__name__)
+
+# Tavily API constants
+TAVILY_API_BASE_URL = "https://api.tavily.com"
+TAVILY_SEARCH_ENDPOINT = "/search"
+DEFAULT_TIMEOUT = 30.0
+DEFAULT_MAX_RETRIES = 3
+DEFAULT_RATE_LIMIT = 1.0  # requests per second
+
+
+class TavilySearchProvider(SearchProvider):
+    """Tavily Search API provider for web search.
+
+    Wraps the Tavily Search API to provide web search capabilities.
+    Supports basic and advanced search depths, domain filtering,
+    and automatic content extraction.
+
+    Attributes:
+        api_key: Tavily API key (required)
+        base_url: API base URL (default: https://api.tavily.com)
+        timeout: Request timeout in seconds (default: 30.0)
+        max_retries: Maximum retry attempts for rate limits (default: 3)
+
+    Example:
+        provider = TavilySearchProvider(api_key="tvly-...")
+        sources = await provider.search(
+            "AI trends 2024",
+            max_results=5,
+            search_depth="advanced",
+        )
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        base_url: str = TAVILY_API_BASE_URL,
+        timeout: float = DEFAULT_TIMEOUT,
+        max_retries: int = DEFAULT_MAX_RETRIES,
+    ):
+        """Initialize Tavily search provider.
+
+        Args:
+            api_key: Tavily API key. If not provided, reads from TAVILY_API_KEY env var.
+            base_url: API base URL (default: https://api.tavily.com)
+            timeout: Request timeout in seconds (default: 30.0)
+            max_retries: Maximum retry attempts for rate limits (default: 3)
+
+        Raises:
+            ValueError: If no API key is provided or found in environment
+        """
+        self._api_key = api_key or os.environ.get("TAVILY_API_KEY")
+        if not self._api_key:
+            raise ValueError(
+                "Tavily API key required. Provide via api_key parameter "
+                "or TAVILY_API_KEY environment variable."
+            )
+
+        self._base_url = base_url.rstrip("/")
+        self._timeout = timeout
+        self._max_retries = max_retries
+        self._rate_limit_value = DEFAULT_RATE_LIMIT
+
+    def get_provider_name(self) -> str:
+        """Return the provider identifier.
+
+        Returns:
+            "tavily"
+        """
+        return "tavily"
+
+    @property
+    def rate_limit(self) -> Optional[float]:
+        """Return the rate limit in requests per second.
+
+        Returns:
+            1.0 (one request per second)
+        """
+        return self._rate_limit_value
+
+    async def search(
+        self,
+        query: str,
+        max_results: int = 10,
+        **kwargs: Any,
+    ) -> list[ResearchSource]:
+        """Execute a web search via Tavily API.
+
+        Args:
+            query: The search query string
+            max_results: Maximum number of results to return (default: 10, max: 20)
+            **kwargs: Additional Tavily options:
+                - search_depth: "basic" or "advanced" (default: "basic")
+                - include_domains: List of domains to include
+                - exclude_domains: List of domains to exclude
+                - include_answer: Whether to include AI answer (default: False)
+                - include_raw_content: Whether to include raw HTML (default: False)
+                - sub_query_id: SubQuery ID for source tracking
+
+        Returns:
+            List of ResearchSource objects
+
+        Raises:
+            AuthenticationError: If API key is invalid
+            RateLimitError: If rate limit exceeded after all retries
+            SearchProviderError: For other API errors
+        """
+        # Extract Tavily-specific options
+        search_depth = kwargs.get("search_depth", "basic")
+        include_domains = kwargs.get("include_domains", [])
+        exclude_domains = kwargs.get("exclude_domains", [])
+        include_answer = kwargs.get("include_answer", False)
+        include_raw_content = kwargs.get("include_raw_content", False)
+        sub_query_id = kwargs.get("sub_query_id")
+
+        # Clamp max_results to Tavily's limit
+        max_results = min(max_results, 20)
+
+        # Build request payload
+        payload = {
+            "api_key": self._api_key,
+            "query": query,
+            "max_results": max_results,
+            "search_depth": search_depth,
+            "include_answer": include_answer,
+            "include_raw_content": include_raw_content,
+        }
+
+        if include_domains:
+            payload["include_domains"] = include_domains
+        if exclude_domains:
+            payload["exclude_domains"] = exclude_domains
+
+        # Execute with retry logic
+        response_data = await self._execute_with_retry(payload)
+
+        # Parse results
+        return self._parse_response(response_data, sub_query_id)
+
+    async def _execute_with_retry(
+        self,
+        payload: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Execute API request with exponential backoff retry.
+
+        Args:
+            payload: Request payload
+
+        Returns:
+            Parsed JSON response
+
+        Raises:
+            AuthenticationError: If API key is invalid
+            RateLimitError: If rate limit exceeded after all retries
+            SearchProviderError: For other API errors
+        """
+        url = f"{self._base_url}{TAVILY_SEARCH_ENDPOINT}"
+        last_error: Optional[Exception] = None
+
+        for attempt in range(self._max_retries):
+            try:
+                async with httpx.AsyncClient(timeout=self._timeout) as client:
+                    response = await client.post(url, json=payload)
+
+                    # Handle authentication errors (not retryable)
+                    if response.status_code == 401:
+                        raise AuthenticationError(
+                            provider="tavily",
+                            message="Invalid API key",
+                        )
+
+                    # Handle rate limiting
+                    if response.status_code == 429:
+                        retry_after = self._parse_retry_after(response)
+                        if attempt < self._max_retries - 1:
+                            wait_time = retry_after or (2**attempt)
+                            logger.warning(
+                                f"Tavily rate limit hit, waiting {wait_time}s "
+                                f"(attempt {attempt + 1}/{self._max_retries})"
+                            )
+                            await asyncio.sleep(wait_time)
+                            continue
+                        raise RateLimitError(
+                            provider="tavily",
+                            retry_after=retry_after,
+                        )
+
+                    # Handle other errors
+                    if response.status_code >= 400:
+                        error_msg = self._extract_error_message(response)
+                        raise SearchProviderError(
+                            provider="tavily",
+                            message=f"API error {response.status_code}: {error_msg}",
+                            retryable=response.status_code >= 500,
+                        )
+
+                    return response.json()
+
+            except httpx.TimeoutException as e:
+                last_error = e
+                if attempt < self._max_retries - 1:
+                    wait_time = 2**attempt
+                    logger.warning(
+                        f"Tavily request timeout, retrying in {wait_time}s "
+                        f"(attempt {attempt + 1}/{self._max_retries})"
+                    )
+                    await asyncio.sleep(wait_time)
+                    continue
+
+            except httpx.RequestError as e:
+                last_error = e
+                if attempt < self._max_retries - 1:
+                    wait_time = 2**attempt
+                    logger.warning(
+                        f"Tavily request error: {e}, retrying in {wait_time}s "
+                        f"(attempt {attempt + 1}/{self._max_retries})"
+                    )
+                    await asyncio.sleep(wait_time)
+                    continue
+
+            except (AuthenticationError, RateLimitError, SearchProviderError):
+                raise
+
+        # All retries exhausted
+        raise SearchProviderError(
+            provider="tavily",
+            message=f"Request failed after {self._max_retries} attempts",
+            retryable=False,
+            original_error=last_error,
+        )
+
+    def _parse_retry_after(self, response: httpx.Response) -> Optional[float]:
+        """Parse Retry-After header from response.
+
+        Args:
+            response: HTTP response
+
+        Returns:
+            Seconds to wait, or None if not provided
+        """
+        retry_after = response.headers.get("Retry-After")
+        if retry_after:
+            try:
+                return float(retry_after)
+            except ValueError:
+                pass
+        return None
+
+    def _extract_error_message(self, response: httpx.Response) -> str:
+        """Extract error message from response.
+
+        Args:
+            response: HTTP response
+
+        Returns:
+            Error message string
+        """
+        try:
+            data = response.json()
+            return data.get("error", data.get("message", response.text[:200]))
+        except Exception:
+            return response.text[:200] if response.text else "Unknown error"
+
+    def _parse_response(
+        self,
+        data: dict[str, Any],
+        sub_query_id: Optional[str] = None,
+    ) -> list[ResearchSource]:
+        """Parse Tavily API response into ResearchSource objects.
+
+        Args:
+            data: Tavily API response JSON
+            sub_query_id: SubQuery ID for source tracking
+
+        Returns:
+            List of ResearchSource objects
+        """
+        sources: list[ResearchSource] = []
+        results = data.get("results", [])
+
+        for result in results:
+            # Create SearchResult from Tavily response
+            search_result = SearchResult(
+                url=result.get("url", ""),
+                title=result.get("title", "Untitled"),
+                snippet=result.get("content"),  # Tavily uses "content" for snippet
+                content=result.get("raw_content"),  # Full content if requested
+                score=result.get("score"),
+                published_date=self._parse_date(result.get("published_date")),
+                source=self._extract_domain(result.get("url", "")),
+                metadata={
+                    "tavily_score": result.get("score"),
+                },
+            )
+
+            # Convert to ResearchSource
+            research_source = search_result.to_research_source(
+                source_type=SourceType.WEB,
+                sub_query_id=sub_query_id,
+            )
+            sources.append(research_source)
+
+        return sources
+
+    def _parse_date(self, date_str: Optional[str]) -> Optional[datetime]:
+        """Parse date string from Tavily response.
+
+        Args:
+            date_str: ISO format date string
+
+        Returns:
+            Parsed datetime or None
+        """
+        if not date_str:
+            return None
+        try:
+            return datetime.fromisoformat(date_str.replace("Z", "+00:00"))
+        except ValueError:
+            return None
+
+    def _extract_domain(self, url: str) -> Optional[str]:
+        """Extract domain from URL.
+
+        Args:
+            url: Full URL
+
+        Returns:
+            Domain name or None
+        """
+        try:
+            from urllib.parse import urlparse
+
+            parsed = urlparse(url)
+            return parsed.netloc
+        except Exception:
+            return None
+
+    async def health_check(self) -> bool:
+        """Check if Tavily API is accessible.
+
+        Performs a lightweight search to verify API key and connectivity.
+
+        Returns:
+            True if provider is healthy, False otherwise
+        """
+        try:
+            # Perform minimal search to verify connectivity
+            await self.search("test", max_results=1)
+            return True
+        except AuthenticationError:
+            logger.error("Tavily health check failed: invalid API key")
+            return False
+        except Exception as e:
+            logger.warning(f"Tavily health check failed: {e}")
+            return False

foundry_mcp/core/research/workflows/__init__.py

@@ -5,18 +5,21 @@ This package provides the workflow classes for multi-model orchestration:
 - ConsensusWorkflow: Multi-model parallel consultation with synthesis
 - ThinkDeepWorkflow: Hypothesis-driven systematic investigation
 - IdeateWorkflow: Creative brainstorming with idea clustering
+- DeepResearchWorkflow: Multi-phase iterative deep research
 """
 
 from foundry_mcp.core.research.workflows.base import ResearchWorkflowBase
 from foundry_mcp.core.research.workflows.chat import ChatWorkflow
 from foundry_mcp.core.research.workflows.consensus import ConsensusWorkflow
-from foundry_mcp.core.research.workflows.thinkdeep import ThinkDeepWorkflow
+from foundry_mcp.core.research.workflows.deep_research import DeepResearchWorkflow
 from foundry_mcp.core.research.workflows.ideate import IdeateWorkflow
+from foundry_mcp.core.research.workflows.thinkdeep import ThinkDeepWorkflow
 
 __all__ = [
     "ResearchWorkflowBase",
     "ChatWorkflow",
     "ConsensusWorkflow",
-    "ThinkDeepWorkflow",
+    "DeepResearchWorkflow",
     "IdeateWorkflow",
+    "ThinkDeepWorkflow",
 ]

foundry_mcp/core/research/workflows/base.py

@@ -10,12 +10,17 @@ from dataclasses import dataclass
 from typing import Any, Optional
 
 from foundry_mcp.config import ResearchConfig
+from foundry_mcp.core.llm_config import ProviderSpec
 from foundry_mcp.core.providers import (
+    ContextWindowError,
     ProviderContext,
     ProviderHooks,
     ProviderRequest,
     ProviderResult,
     ProviderStatus,
+    is_context_window_error,
+    extract_token_counts,
+    create_context_window_guidance,
 )
 from foundry_mcp.core.providers.registry import available_providers, resolve_provider
 from foundry_mcp.core.research.memory import ResearchMemory
@@ -23,6 +28,25 @@ from foundry_mcp.core.research.memory import ResearchMemory
 logger = logging.getLogger(__name__)
 
 
+def _estimate_prompt_tokens(prompt: str, system_prompt: str | None = None) -> int:
+    """Estimate token count for a prompt using simple heuristic.
+
+    Uses ~4 characters per token as a rough estimate. This is conservative
+    and works reasonably well for English text.
+
+    Args:
+        prompt: User prompt
+        system_prompt: Optional system prompt
+
+    Returns:
+        Estimated token count
+    """
+    total_chars = len(prompt)
+    if system_prompt:
+        total_chars += len(system_prompt)
+    return total_chars // 4
+
+
 @dataclass
 class WorkflowResult:
     """Result of a workflow execution.
@@ -33,6 +57,9 @@ class WorkflowResult:
         provider_id: Provider that generated the response
         model_used: Model that generated the response
         tokens_used: Total tokens consumed
+        input_tokens: Tokens consumed by the prompt
+        output_tokens: Tokens generated in the response
+        cached_tokens: Tokens served from cache
         duration_ms: Execution duration in milliseconds
         metadata: Additional workflow-specific data
         error: Error message if success is False
@@ -43,6 +70,9 @@ class WorkflowResult:
     provider_id: Optional[str] = None
     model_used: Optional[str] = None
     tokens_used: Optional[int] = None
+    input_tokens: Optional[int] = None
+    output_tokens: Optional[int] = None
+    cached_tokens: Optional[int] = None
     duration_ms: Optional[float] = None
     metadata: dict[str, Any] = None
     error: Optional[str] = None
@@ -85,30 +115,48 @@ class ResearchWorkflowBase(ABC):
         """Resolve and cache a provider instance.
 
         Args:
-            provider_id: Provider ID to resolve (uses config default if None)
+            provider_id: Provider ID or full spec to resolve (uses config default if None)
+                         Supports both simple IDs ("codex") and full specs ("[cli]codex:gpt-5.2")
             hooks: Optional provider hooks
 
         Returns:
             ProviderContext instance or None if unavailable
         """
-        provider_id = provider_id or self.config.default_provider
+        provider_spec_str = provider_id or self.config.default_provider
+
+        # Check cache first (using full spec string as key)
+        if provider_spec_str in self._provider_cache:
+            return self._provider_cache[provider_spec_str]
 
-        # Check cache first
-        if provider_id in self._provider_cache:
-            return self._provider_cache[provider_id]
+        # Parse the provider spec to extract base provider ID
+        try:
+            spec = ProviderSpec.parse_flexible(provider_spec_str)
+        except ValueError as exc:
+            logger.warning("Invalid provider spec '%s': %s", provider_spec_str, exc)
+            return None
 
-        # Check availability
+        # Check availability using base provider ID
         available = available_providers()
-        if provider_id not in available:
-            logger.warning("Provider %s not available. Available: %s", provider_id, available)
+        if spec.provider not in available:
+            logger.warning(
+                "Provider %s (from spec '%s') not available. Available: %s",
+                spec.provider,
+                provider_spec_str,
+                available,
+            )
             return None
 
         try:
-            provider = resolve_provider(provider_id, hooks=hooks or ProviderHooks())
-            self._provider_cache[provider_id] = provider
+            # Resolve using base provider ID and pass model override if specified
+            provider = resolve_provider(
+                spec.provider,
+                hooks=hooks or ProviderHooks(),
+                model=spec.model,
+            )
+            self._provider_cache[provider_spec_str] = provider
             return provider
         except Exception as exc:
-            logger.error("Failed to resolve provider %s: %s", provider_id, exc)
+            logger.error("Failed to resolve provider %s: %s", spec.provider, exc)
             return None
 
     def _execute_provider(
@@ -149,11 +197,14 @@ class ResearchWorkflowBase(ABC):
             prompt=prompt,
             system_prompt=system_prompt,
             model=model,
-            timeout=timeout or 30.0,
+            timeout=timeout or self.config.default_timeout,
             temperature=temperature,
             max_tokens=max_tokens,
         )
 
+        # Estimate prompt tokens for error reporting
+        estimated_tokens = _estimate_prompt_tokens(prompt, system_prompt)
+
         try:
             result: ProviderResult = provider.generate(request)
 
@@ -172,10 +223,53 @@ class ResearchWorkflowBase(ABC):
                 provider_id=result.provider_id,
                 model_used=result.model_used,
                 tokens_used=result.tokens.total_tokens if result.tokens else None,
+                input_tokens=result.tokens.input_tokens if result.tokens else None,
+                output_tokens=result.tokens.output_tokens if result.tokens else None,
+                cached_tokens=result.tokens.cached_input_tokens if result.tokens else None,
                 duration_ms=result.duration_ms,
             )
 
+        except ContextWindowError:
+            # Re-raise context window errors directly
+            raise
+
         except Exception as exc:
+            # Check if this is a context window error
+            if is_context_window_error(exc):
+                # Extract token counts from error message if available
+                prompt_tokens, max_context = extract_token_counts(str(exc))
+
+                # Use estimated tokens if not extracted
+                if prompt_tokens is None:
+                    prompt_tokens = estimated_tokens
+
+                # Log detailed context window error
+                logger.error(
+                    "Context window exceeded: prompt_tokens=%s, max_tokens=%s, "
+                    "estimated_tokens=%d, provider=%s, error=%s",
+                    prompt_tokens,
+                    max_context,
+                    estimated_tokens,
+                    provider_id,
+                    str(exc),
+                )
+
+                # Generate actionable guidance
+                guidance = create_context_window_guidance(
+                    prompt_tokens=prompt_tokens,
+                    max_tokens=max_context,
+                    provider_id=provider_id,
+                )
+
+                # Raise specific context window error with details
+                raise ContextWindowError(
+                    guidance,
+                    provider=provider_id,
+                    prompt_tokens=prompt_tokens,
+                    max_tokens=max_context,
+                ) from exc
+
+            # Non-context-window error - log and return error result
             logger.error("Provider execution failed: %s", exc)
             return WorkflowResult(
                 success=False,