From 39bdc10aef9996a4f10f35d2e91af69c5ef3cdfe Mon Sep 17 00:00:00 2001 From: Naufal Reky Ardhana Date: Tue, 11 Aug 2026 00:09:09 +0700 Subject: [PATCH] fix: handle current perplexity SSE responses --- mcp_service.py | 46 +++- src/core/perplexity_client.py | 472 ++++++++++++++++++++++++-------- src/services/chunk_extractor.py | 37 ++- src/services/sse_parser.py | 12 +- tests/test_chunk_extractor.py | 37 ++- tests/test_mcp_tools.py | 58 +++- tests/test_perplexity_client.py | 363 +++++++++++++++++++++++- tests/test_sse_parser.py | 12 +- 8 files changed, 887 insertions(+), 150 deletions(-) diff --git a/mcp_service.py b/mcp_service.py index a96a073..e0d26ff 100644 --- a/mcp_service.py +++ b/mcp_service.py @@ -7,8 +7,10 @@ import logging import os from typing import Literal, Optional + from mcp.server.fastmcp import FastMCP from mcp.server.transport_security import TransportSecuritySettings + from src.core.perplexity_client import PerplexityClient from src.prompts import PROGRAMMING_RESEARCH_PROMPTS, VALID_CATEGORIES @@ -62,6 +64,15 @@ def get_transport_security() -> TransportSecuritySettings: _client: Optional[PerplexityClient] = None DEFAULT_MODEL = "gpt56_terra_thinking" +SEARCH_TIMEOUT_SECONDS = 45.0 +RESEARCH_TIMEOUT_SECONDS = 180.0 +MIN_TIMEOUT_SECONDS = 5.0 +MAX_TIMEOUT_SECONDS = 600.0 + + +class EmptyResponseError(RuntimeError): + """Raised when upstream finishes without any usable answer text.""" + ResearchCategory = Literal[ "academic", @@ -102,14 +113,31 @@ def _build_response( include_related: bool, ) -> dict: """Build a response dict including only requested optional fields.""" - result: dict = {"text": response.text or "No response received."} + if not response.text: + raise EmptyResponseError("Perplexity returned no answer text.") + + result: dict = {"text": response.text} if include_citations: result["citations"] = response.citations if include_related: result["related_queries"] = response.related_queries + if response.partial: + result["partial"] = True + result["warning"] = response.warning or "Returning a partial answer." return result +def _validate_timeout(timeout_seconds: float) -> float: + """Validate the bounded timeout exposed through MCP tool schemas.""" + if isinstance(timeout_seconds, bool) or not isinstance( + timeout_seconds, (int, float) + ): + raise ValueError("timeout_seconds must be a number between 5 and 600") + if not MIN_TIMEOUT_SECONDS <= timeout_seconds <= MAX_TIMEOUT_SECONDS: + raise ValueError("timeout_seconds must be between 5 and 600 seconds") + return float(timeout_seconds) + + def _error_dict(error: Exception) -> dict: """Build an error response with minimal shape.""" msg = f"[Error] {type(error).__name__}: {error}" @@ -124,7 +152,8 @@ def perplexity_search( include_citations: bool = False, include_related: bool = False, model_preference: str = DEFAULT_MODEL, - mode: Literal["copilot", "search"] = "copilot", + mode: Literal["copilot", "search"] = "search", + timeout_seconds: float = SEARCH_TIMEOUT_SECONDS, ) -> dict: """Search via Perplexity AI. Returns answer text; citations opt-in. @@ -137,15 +166,16 @@ def perplexity_search( include_related: If True, include related query suggestions. model_preference: AI model to use. mode: "copilot" for comprehensive answers, "search" for quick results. + timeout_seconds: Maximum request time, from 5 to 600 seconds. Returns: Dict with "text" always present. "citations" and "related_queries" included only when their respective flags are True. """ - selected_sources: list[str] = list(sources) if sources else ["web"] - search_focus = "academic" if selected_sources == ["scholar"] else "internet" - try: + timeout = _validate_timeout(timeout_seconds) + selected_sources: list[str] = list(sources) if sources else ["web"] + search_focus = "academic" if selected_sources == ["scholar"] else "internet" client = get_client() response = client.ask( query=query, @@ -153,6 +183,7 @@ def perplexity_search( model_preference=model_preference, search_focus=search_focus, sources=selected_sources, + timeout_seconds=timeout, ) return _build_response(response, include_citations, include_related) except Exception as e: @@ -166,6 +197,7 @@ def perplexity_research( include_citations: bool = False, include_related: bool = False, model_preference: str = DEFAULT_MODEL, + timeout_seconds: float = RESEARCH_TIMEOUT_SECONDS, ) -> dict: """Research a topic with category-specific prompts. Citations opt-in. @@ -179,6 +211,7 @@ def perplexity_research( include_citations: If True, include source citations. Default False. include_related: If True, include related query suggestions. model_preference: AI model to use. + timeout_seconds: Maximum request time, from 5 to 600 seconds. Returns: Dict with "text" always present. "citations" and "related_queries" @@ -192,6 +225,7 @@ def perplexity_research( research_prompt = prompt_template.format(topic=topic) try: + timeout = _validate_timeout(timeout_seconds) client = get_client() response = client.ask( query=research_prompt, @@ -199,6 +233,7 @@ def perplexity_research( model_preference=model_preference, search_focus="internet", sources=["web", "scholar"], + timeout_seconds=timeout, ) return _build_response(response, include_citations, include_related) except Exception as e: @@ -210,6 +245,7 @@ def perplexity_research( if config.mcp_transport_mode == "http": import uvicorn + from src.core.mcp_auth import MCPAuthMiddleware app = mcp.streamable_http_app() diff --git a/src/core/perplexity_client.py b/src/core/perplexity_client.py index b7b18f0..2b88224 100644 --- a/src/core/perplexity_client.py +++ b/src/core/perplexity_client.py @@ -6,12 +6,22 @@ """ import json +import logging import os +import time import uuid -from typing import Generator, Optional, Any from dataclasses import dataclass, field -from dotenv import load_dotenv +from typing import Any, Generator, Optional + from curl_cffi import requests as cffi_requests +from curl_cffi.const import CurlOpt +from curl_cffi.requests.exceptions import Timeout as CurlTimeout +from dotenv import load_dotenv + +from src.services.sse_parser import PerplexitySSEParser + +logger = logging.getLogger(__name__) +DEFAULT_REQUEST_TIMEOUT_SECONDS = 1800.0 @dataclass @@ -23,6 +33,253 @@ class PerplexityResponse: media_items: list[dict] = field(default_factory=list) related_queries: list[str] = field(default_factory=list) raw_events: list[dict] = field(default_factory=list) + partial: bool = False + warning: Optional[str] = None + + +def _add_citations( + response: PerplexityResponse, + web_results: Any, + seen_urls: set[str], +) -> None: + """Append valid web citations once, keyed by URL.""" + if not isinstance(web_results, list): + return + + for web_result in web_results: + if not isinstance(web_result, dict): + continue + title = web_result.get("name", "") + url = web_result.get("url", "") + if not title or not url or url in seen_urls: + continue + seen_urls.add(url) + response.citations.append( + { + "title": title, + "url": url, + "snippet": web_result.get("snippet", ""), + } + ) + + +def _process_legacy_final( + event: dict, + response: PerplexityResponse, + seen_urls: set[str], +) -> None: + """Parse the legacy nested FINAL event format.""" + response.related_queries = event.get("related_queries", []) + text_json = event.get("text", "") + if not text_json: + return + + try: + steps = json.loads(text_json) + except (json.JSONDecodeError, TypeError): + return + if not isinstance(steps, list): + return + + for step in steps: + if not isinstance(step, dict): + continue + content = step.get("content", {}) + if not isinstance(content, dict): + continue + + if step.get("step_type") == "SEARCH_RESULTS": + _add_citations(response, content.get("web_results", []), seen_urls) + + if step.get("step_type") == "FINAL": + answer_value = content.get("answer", "") + if isinstance(answer_value, str) and answer_value: + try: + answer_data = json.loads(answer_value) + except json.JSONDecodeError: + response.text = answer_value + else: + if isinstance(answer_data, dict): + if isinstance(answer_data.get("answer"), str): + response.text = answer_data["answer"] + _add_citations( + response, + answer_data.get("web_results", []), + seen_urls, + ) + for item in answer_data.get("structured_answer", []): + if ( + isinstance(item, dict) + and item.get("type") == "markdown" + and isinstance(item.get("text"), str) + ): + response.text = item["text"] + + for item in content.get("structured_answer", []): + if ( + isinstance(item, dict) + and item.get("type") == "markdown" + and isinstance(item.get("text"), str) + ): + response.text = item["text"] + + +def _set_chunk(chunks: list[str], index: int, value: Any) -> None: + """Set a streamed chunk by index, extending sparse snapshots safely.""" + while len(chunks) <= index: + chunks.append("") + chunks[index] = str(value) if value is not None else "" + + +def _process_answer_block( + block: dict, + chunk_state: dict[str, list[str]], + answer_state: dict[str, str], +) -> bool: + """Apply a current or legacy answer block and report completion.""" + intended_usage = block.get("intended_usage", "") + if not PerplexitySSEParser.is_markdown_block(intended_usage): + return False + + complete = False + has_direct_answer = False + markdown = block.get("markdown_block") + if isinstance(markdown, dict): + answer = markdown.get("answer") + chunks = markdown.get("chunks") + if isinstance(answer, str) and answer: + answer_state[intended_usage] = answer + has_direct_answer = True + elif isinstance(chunks, list): + chunk_state[intended_usage] = [str(chunk) for chunk in chunks] + complete = markdown.get("progress") == "DONE" + + diff_block = block.get("diff_block") + if not isinstance(diff_block, dict): + return complete and intended_usage in ("ask_text", "ask_text_markdown") + + for patch in diff_block.get("patches", []): + if not isinstance(patch, dict): + continue + path = patch.get("path", "") + value = patch.get("value") + op = patch.get("op") + + if path == "" and isinstance(value, dict): + complete = complete or value.get("progress") == "DONE" + if has_direct_answer: + continue + answer = value.get("answer") + chunks = value.get("chunks") + if isinstance(answer, str) and answer: + answer_state[intended_usage] = answer + elif isinstance(chunks, list): + chunk_state[intended_usage] = [str(chunk) for chunk in chunks] + elif path == "/progress": + complete = complete or value == "DONE" + elif ( + not has_direct_answer + and path == "/answer" + and isinstance(value, str) + and value + ): + answer_state[intended_usage] = value + elif ( + not has_direct_answer + and path.startswith("/chunks/") + and op in ("add", "replace") + ): + try: + index = int(path.rsplit("/", 1)[-1]) + except ValueError: + continue + chunks = chunk_state.setdefault(intended_usage, []) + _set_chunk(chunks, index, value) + + return complete and intended_usage in ("ask_text", "ask_text_markdown") + + +def _current_answer_text( + chunk_state: dict[str, list[str]], + answer_state: dict[str, str], +) -> str: + """Choose the combined answer, falling back to ordered section blocks.""" + + def block_text(intended_usage: str) -> str: + return answer_state.get(intended_usage) or "".join( + chunk_state.get(intended_usage, []) + ) + + for combined_usage in ("ask_text", "ask_text_markdown"): + text = block_text(combined_usage) + if text: + return text + + def section_sort_key(intended_usage: str) -> tuple[int, int, str]: + section = intended_usage.removeprefix("ask_text_").removesuffix("_markdown") + try: + return (0, int(section), intended_usage) + except ValueError: + return (1, 0, intended_usage) + + section_usages = sorted( + ( + usage + for usage in set(chunk_state) | set(answer_state) + if usage.startswith("ask_text_") and usage.endswith("_markdown") + ), + key=section_sort_key, + ) + return "".join(block_text(usage) for usage in section_usages) + + +def _process_current_blocks( + event: dict, + response: PerplexityResponse, + chunk_state: dict[str, list[str]], + answer_state: dict[str, str], + seen_urls: set[str], +) -> bool: + """Parse current block-only events and report terminal answer progress.""" + complete = False + for block in event.get("blocks", []) or []: + if not isinstance(block, dict): + continue + intended_usage = block.get("intended_usage", "") + + if intended_usage == "web_results": + direct_results = block.get("web_results") + _add_citations(response, direct_results, seen_urls) + direct_block = block.get("web_result_block") + if isinstance(direct_block, dict): + _add_citations( + response, + direct_block.get("web_results", []), + seen_urls, + ) + diff_block = block.get("diff_block") + if isinstance(diff_block, dict): + for patch in diff_block.get("patches", []): + if not isinstance(patch, dict): + continue + value = patch.get("value") + if isinstance(value, dict): + _add_citations( + response, + value.get("web_results", []), + seen_urls, + ) + if patch.get("path", "").startswith("/web_results/"): + _add_citations(response, [value], seen_urls) + elif isinstance(value, list): + _add_citations(response, value, seen_urls) + + complete = _process_answer_block(block, chunk_state, answer_state) or complete + + current_text = _current_answer_text(chunk_state, answer_state) + if current_text: + response.text = current_text + return complete class PerplexityClient: @@ -211,6 +468,7 @@ def ask_stream( search_focus: str = "internet", sources: Optional[list[str]] = None, is_incognito: bool = False, + timeout_seconds: float = DEFAULT_REQUEST_TIMEOUT_SECONDS, ) -> Generator[dict, None, None]: """ Send a query to Perplexity and stream the response. @@ -222,6 +480,7 @@ def ask_stream( search_focus: Search focus (internet, academic, etc.) sources: List of sources to search is_incognito: If True, query won't appear in Perplexity dashboard + timeout_seconds: Upstream inactivity timeout in seconds Yields: Parsed SSE event dictionaries @@ -247,29 +506,32 @@ def ask_stream( cookies=cookies, json=payload, impersonate="edge", - timeout=1800, + timeout=timeout_seconds, + curl_options={CurlOpt.TIMEOUT_MS: int(timeout_seconds * 1000)}, stream=True, ) - - if response.status_code != 200: - raise Exception( - f"Request failed with status {response.status_code}: {response.text}" - ) - - # Parse SSE stream - buffer = "" - for chunk in response.iter_content(): - if chunk: - buffer += chunk.decode("utf-8", errors="ignore") - - while "\n" in buffer: - line, buffer = buffer.split("\n", 1) - line = line.strip() - - if line.startswith("data:"): - event = self._parse_sse_line(line) - if event: - yield event + try: + if response.status_code != 200: + raise Exception( + f"Request failed with status {response.status_code}: {response.text}" + ) + + # Parse SSE stream + buffer = "" + for chunk in response.iter_content(): + if chunk: + buffer += chunk.decode("utf-8", errors="ignore") + + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) + line = line.strip() + + if line.startswith("data:"): + event = self._parse_sse_line(line) + if event: + yield event + finally: + response.close() def ask( self, @@ -279,6 +541,7 @@ def ask( search_focus: str = "internet", sources: Optional[list[str]] = None, is_incognito: bool = False, + timeout_seconds: float = DEFAULT_REQUEST_TIMEOUT_SECONDS, ) -> PerplexityResponse: """ Send a query to Perplexity and return the complete response. @@ -290,115 +553,82 @@ def ask( search_focus: Search focus (internet, academic, etc.) sources: List of sources to search is_incognito: If True, query won't appear in Perplexity dashboard + timeout_seconds: Maximum total request time in seconds Returns: PerplexityResponse with parsed results """ result = PerplexityResponse() - - for event in self.ask_stream( + chunk_state: dict[str, list[str]] = {} + answer_state: dict[str, str] = {} + seen_urls: set[str] = set() + started_at = time.monotonic() + deadline = started_at + timeout_seconds + termination = "stream_ended" + stream = self.ask_stream( query=query, mode=mode, model_preference=model_preference, search_focus=search_focus, sources=sources, is_incognito=is_incognito, - ): - result.raw_events.append(event) - - step_type = event.get("step_type", "") - - # Handle FINAL event - contains the complete response as nested JSON - if step_type == "FINAL": - # Get related queries from top level - result.related_queries = event.get("related_queries", []) - - # Parse the nested text field which contains all step data as JSON - text_json = event.get("text", "") - if text_json: - try: - steps = json.loads(text_json) - for step in steps: - step_content = step.get("content", {}) - inner_step_type = step.get("step_type", "") - - # Extract citations from SEARCH_RESULTS - if inner_step_type == "SEARCH_RESULTS": - web_results = step_content.get("web_results", []) - for wr in web_results: - if wr.get("name") and wr.get("url"): - result.citations.append( - { - "title": wr.get("name", ""), - "url": wr.get("url", ""), - "snippet": wr.get("snippet", ""), - } - ) - - # Extract answer from inner FINAL step - if inner_step_type == "FINAL": - answer_str = step_content.get("answer", "") - if answer_str: - try: - answer_data = json.loads(answer_str) - if "answer" in answer_data: - result.text = answer_data["answer"] - # Also extract citations from web_results in answer - for wr in answer_data.get( - "web_results", [] - ): - if wr.get("name") and wr.get("url"): - result.citations.append( - { - "title": wr.get("name", ""), - "url": wr.get("url", ""), - "snippet": wr.get( - "snippet", "" - ), - } - ) - # Extract from structured_answer - for item in answer_data.get( - "structured_answer", [] - ): - if ( - item.get("type") == "markdown" - and item.get("text") - ): - result.text = item["text"] - except json.JSONDecodeError: - # answer might be plain text - result.text = answer_str - - # Also check for structured_answer in other steps - if "structured_answer" in step_content: - for item in step_content.get("structured_answer", []): - if item.get("type") == "markdown" and item.get( - "text" - ): - result.text = item["text"] - except json.JSONDecodeError: - pass - - # Handle streaming blocks with text chunks (fallback if no FINAL) - if not step_type and not result.text: - blocks = event.get("blocks", []) - for block in blocks: - if block.get("intended_usage") in ( - "ask_text_0_markdown", - "ask_text", - ): - diff_block = block.get("diff_block", {}) - patches = diff_block.get("patches", []) - for patch in patches: - if patch.get("op") == "replace" and patch.get("value"): - value = patch["value"] - if isinstance(value, dict) and value.get("chunks"): - result.text = "".join(value["chunks"]) - elif patch.get("op") == "add" and patch.get("value"): - # Append chunk - if result.text: - result.text += patch["value"] + timeout_seconds=timeout_seconds, + ) + + try: + for event in stream: + if not isinstance(event, dict): + continue + result.raw_events.append(event) + + step_type = event.get("step_type", "") + if step_type == "FINAL": + _process_legacy_final(event, result, seen_urls) + + block_complete = _process_current_blocks( + event, + result, + chunk_state, + answer_state, + seen_urls, + ) + + if step_type == "FINAL" or event.get("text_completed") is True: + termination = "completed" + break + if block_complete: + termination = "answer_done" + break + if time.monotonic() >= deadline: + raise CurlTimeout( + f"Perplexity request exceeded {timeout_seconds:g} seconds" + ) + except CurlTimeout as error: + if not result.text: + raise CurlTimeout( + "Perplexity request timed out with no answer text received " + f"within {timeout_seconds:g} seconds" + ) from error + result.partial = True + result.warning = ( + f"Timed out after {timeout_seconds:g} seconds; " + "returning partial answer." + ) + termination = "timeout_partial" + finally: + close = getattr(stream, "close", None) + if callable(close): + close() + logger.info( + "Perplexity request finished model=%s mode=%s duration=%.2fs " + "events=%d termination=%s partial=%s", + model_preference, + mode, + time.monotonic() - started_at, + len(result.raw_events), + termination, + result.partial, + ) return result diff --git a/src/services/chunk_extractor.py b/src/services/chunk_extractor.py index 3292f11..6eaf11d 100644 --- a/src/services/chunk_extractor.py +++ b/src/services/chunk_extractor.py @@ -8,10 +8,8 @@ from typing import Iterator, Optional from src.models.perplexity_models import ( - PerplexitySSEEvent, PerplexityBlock, StreamingState, - ChunkAggregator, ) from src.services.sse_parser import PerplexitySSEParser @@ -73,12 +71,41 @@ def _process_block(self, block: PerplexityBlock) -> Iterator[str]: Yields: Text chunks from the block's patches. """ - if not block.diff_block: - return - # Get or create aggregator for this block aggregator = self.state.get_or_create_aggregator(block.intended_usage) + # Current Perplexity responses send a complete markdown snapshot + # directly, without JSON patches. ``answer`` is authoritative when + # present; ``chunks`` remains the fallback for in-progress snapshots. + if block.markdown_block: + markdown = block.markdown_block + snapshot: Optional[str] = None + if markdown.answer: + snapshot = markdown.answer + elif markdown.chunks: + snapshot = "".join(markdown.chunks) + + if snapshot is not None: + current = aggregator.get_full_text() + if snapshot != current: + new_text = ( + snapshot[len(current) :] + if snapshot.startswith(current) + else snapshot + ) + aggregator.chunks = [snapshot] + if new_text: + yield new_text + + if markdown.progress == "DONE": + aggregator.is_complete = True + + if snapshot is not None: + return + + if not block.diff_block: + return + # Apply each patch and yield new content for patch in block.diff_block.patches: new_text = aggregator.apply_patch(patch) diff --git a/src/services/sse_parser.py b/src/services/sse_parser.py index 923f29e..0a316ca 100644 --- a/src/services/sse_parser.py +++ b/src/services/sse_parser.py @@ -4,15 +4,14 @@ Parses raw SSE events from Perplexity API and extracts structured blocks. """ -import json import logging -from typing import Optional, Iterator +from typing import Iterator, Optional from src.models.perplexity_models import ( - PerplexitySSEEvent, - PerplexityBlock, DiffBlock, JSONPatch, + PerplexityBlock, + PerplexitySSEEvent, ) logger = logging.getLogger(__name__) @@ -104,6 +103,7 @@ def is_markdown_block(intended_usage: str) -> bool: Check if a block is a markdown answer block. Markdown answer blocks have intended_usage matching: + - "ask_text" (current combined answer) - "ask_text_markdown" (combined answer) - "ask_text_N_markdown" where N is a number (individual sections) @@ -116,8 +116,8 @@ def is_markdown_block(intended_usage: str) -> bool: if not intended_usage: return False - # Match "ask_text_markdown" or "ask_text_N_markdown" - if intended_usage == "ask_text_markdown": + # Match the current plain block and legacy markdown block names. + if intended_usage in ("ask_text", "ask_text_markdown"): return True if intended_usage.startswith("ask_text_") and intended_usage.endswith( "_markdown" diff --git a/tests/test_chunk_extractor.py b/tests/test_chunk_extractor.py index a60f324..ff96af2 100644 --- a/tests/test_chunk_extractor.py +++ b/tests/test_chunk_extractor.py @@ -5,21 +5,21 @@ Covers all public methods and edge cases with 90%+ coverage target. """ -import pytest -from unittest.mock import Mock, MagicMock, patch from typing import Iterator, Optional +from unittest.mock import MagicMock, Mock, patch + +import pytest -from src.services.chunk_extractor import ChunkExtractor, extract_chunks_from_events from src.models.perplexity_models import ( - PerplexitySSEEvent, - PerplexityBlock, + ChunkAggregator, DiffBlock, JSONPatch, MarkdownBlock, - ChunkAggregator, + PerplexityBlock, + PerplexitySSEEvent, StreamingState, ) - +from src.services.chunk_extractor import ChunkExtractor, extract_chunks_from_events # ============================================================================ # Fixtures @@ -173,6 +173,29 @@ def test_init_creates_parser(self): class TestChunkExtractorProcessEvent: """Test ChunkExtractor.process_event() method.""" + def test_process_current_direct_markdown_answer(self): + """A direct ask_text markdown snapshot must produce the final answer.""" + extractor = ChunkExtractor() + event_data = { + "text_completed": False, + "blocks": [ + { + "intended_usage": "ask_text", + "markdown_block": { + "progress": "DONE", + "chunks": ["stale chunk"], + "answer": "Current schema answer", + }, + } + ], + } + + chunks = list(extractor.process_event(event_data)) + + assert chunks == ["Current schema answer"] + assert extractor.get_full_text() == "Current schema answer" + assert extractor.is_complete() is True + def test_process_event_with_markdown_blocks(self, perplexity_block_with_diff): """Test processing event with markdown blocks yields chunks.""" extractor = ChunkExtractor() diff --git a/tests/test_mcp_tools.py b/tests/test_mcp_tools.py index c12f305..86387f0 100644 --- a/tests/test_mcp_tools.py +++ b/tests/test_mcp_tools.py @@ -9,7 +9,7 @@ - Error responses return {"text": "[Error] ..."} """ -from unittest.mock import patch, MagicMock +from unittest.mock import MagicMock, patch import pytest @@ -26,6 +26,8 @@ def mock_response(): {"title": "Source B", "url": "https://b.example", "snippet": "snippet b"}, ] resp.related_queries = ["query 1", "query 2"] + resp.partial = False + resp.warning = None return resp @@ -79,6 +81,9 @@ def test_default_sources_web(self, mock_client): call_kwargs = mock_client.ask.call_args.kwargs assert call_kwargs["sources"] == ["web"] assert call_kwargs["search_focus"] == "internet" + assert call_kwargs["mode"] == "search" + assert call_kwargs["model_preference"] == "gpt56_terra_thinking" + assert call_kwargs["timeout_seconds"] == 45 def test_scholar_only_uses_academic_focus(self, mock_client): mcp_service.perplexity_search("q", sources=["scholar"]) @@ -101,6 +106,9 @@ def test_default_returns_text_only(self, mock_client): assert "text" in result assert "citations" not in result assert "related_queries" not in result + call_kwargs = mock_client.ask.call_args.kwargs + assert call_kwargs["mode"] == "copilot" + assert call_kwargs["timeout_seconds"] == 180 def test_academic_category_resolves(self, mock_client): """New 'academic' category must be registered and load the academic template.""" @@ -130,6 +138,43 @@ def test_include_citations_opt_in(self, mock_client): assert "citations" in result +class TestTimeoutAndPartialResponse: + """Verify user-facing timeout controls and partial-result metadata.""" + + @pytest.mark.parametrize("timeout_seconds", [4, 601]) + def test_search_rejects_timeout_outside_safe_range(self, timeout_seconds): + result = mcp_service.perplexity_search("q", timeout_seconds=timeout_seconds) + + assert set(result) == {"text"} + assert result["text"].startswith("[Error] ValueError:") + assert "between 5 and 600" in result["text"] + + def test_custom_search_timeout_is_forwarded(self, mock_client): + mcp_service.perplexity_search("q", timeout_seconds=30) + + assert mock_client.ask.call_args.kwargs["timeout_seconds"] == 30 + + def test_partial_response_is_marked(self, mock_client, mock_response): + mock_response.partial = True + mock_response.warning = "Timed out after 45 seconds; returning partial answer." + + result = mcp_service.perplexity_search("q") + + assert result == { + "text": "The answer is 42.", + "partial": True, + "warning": "Timed out after 45 seconds; returning partial answer.", + } + + def test_empty_response_is_an_explicit_error(self, mock_client, mock_response): + mock_response.text = "" + + result = mcp_service.perplexity_search("q") + + assert set(result) == {"text"} + assert result["text"].startswith("[Error] EmptyResponseError:") + + class TestErrorShape: """Error responses still match the minimal shape.""" @@ -170,3 +215,14 @@ def test_old_tools_removed(self): def test_new_tools_present(self): assert hasattr(mcp_service, "perplexity_search") assert hasattr(mcp_service, "perplexity_research") + + def test_fastmcp_schema_exposes_latency_defaults(self): + """OpenCode must see the intended mode/model/timeout tool defaults.""" + tools = {tool.name: tool for tool in mcp_service.mcp._tool_manager.list_tools()} + search = tools["perplexity_search"].parameters["properties"] + research = tools["perplexity_research"].parameters["properties"] + + assert search["model_preference"]["default"] == "gpt56_terra_thinking" + assert search["mode"]["default"] == "search" + assert search["timeout_seconds"]["default"] == 45.0 + assert research["timeout_seconds"]["default"] == 180.0 diff --git a/tests/test_perplexity_client.py b/tests/test_perplexity_client.py index e0e62ee..63057d2 100644 --- a/tests/test_perplexity_client.py +++ b/tests/test_perplexity_client.py @@ -11,9 +11,12 @@ """ import json -import pytest -from unittest.mock import patch, MagicMock from dataclasses import fields +from unittest.mock import MagicMock, patch + +import pytest +from curl_cffi.const import CurlOpt +from curl_cffi.requests.exceptions import Timeout as CurlTimeout from src.core.perplexity_client import PerplexityClient, PerplexityResponse @@ -836,6 +839,8 @@ def test_response_has_all_expected_fields(self): "media_items", "related_queries", "raw_events", + "partial", + "warning", } assert field_names == expected_fields @@ -902,5 +907,359 @@ def test_headers_contains_no_none_values(self): assert value is not None, f"Header {key} contains None" +class TestCurrentSSEAnswerParsing: + """Regression tests for the current block-only Perplexity SSE schema.""" + + def _create_client(self): + env_vars = {"PERPLEXITY_SESSION_TOKEN": "test_session_token"} + with patch.dict("os.environ", env_vars, clear=True): + with patch("src.core.perplexity_client.load_dotenv"): + return PerplexityClient() + + def test_direct_ask_text_answer_is_authoritative_and_stops_stream(self): + """A DONE ask_text.answer must win over chunks and stop consumption.""" + client = self._create_client() + stream_closed = [] + + def events(): + try: + yield { + "blocks": [ + { + "intended_usage": "ask_text", + "markdown_block": { + "progress": "DONE", + "chunks": ["stale chunk"], + "answer": "Current schema answer", + }, + } + ] + } + raise AssertionError("stream was consumed after terminal answer") + finally: + stream_closed.append(True) + + with patch.object(client, "ask_stream", return_value=events()): + response = client.ask("question") + + assert response.text == "Current schema answer" + assert response.partial is False + assert stream_closed == [True] + + def test_direct_answer_wins_over_conflicting_diff_snapshot(self): + """A stale diff snapshot must never replace markdown_block.answer.""" + client = self._create_client() + event = { + "blocks": [ + { + "intended_usage": "ask_text", + "markdown_block": { + "progress": "DONE", + "chunks": ["stale chunk"], + "answer": "Authoritative answer", + }, + "diff_block": { + "field": "markdown_block", + "patches": [ + { + "op": "replace", + "path": "", + "value": { + "progress": "DONE", + "chunks": ["stale diff chunk"], + "answer": "Stale diff answer", + }, + } + ], + }, + } + ] + } + + with patch.object(client, "ask_stream", return_value=iter([event])): + response = client.ask("question") + + assert response.text == "Authoritative answer" + + def test_numbered_sections_wait_for_stream_completion_and_sort_numerically(self): + """DONE sections must not truncate later sections or sort 10 before 2.""" + client = self._create_client() + events = iter( + [ + { + "blocks": [ + { + "intended_usage": "ask_text_2_markdown", + "markdown_block": { + "progress": "DONE", + "answer": "two", + }, + } + ] + }, + { + "blocks": [ + { + "intended_usage": "ask_text_10_markdown", + "markdown_block": { + "progress": "DONE", + "answer": "ten", + }, + } + ] + }, + { + "text_completed": True, + "blocks": [ + { + "intended_usage": "ask_text_1_markdown", + "markdown_block": { + "progress": "DONE", + "answer": "one", + }, + } + ], + }, + ] + ) + + with patch.object(client, "ask_stream", return_value=events): + response = client.ask("question") + + assert response.text == "onetwoten" + + def test_current_web_results_are_deduplicated_by_url(self): + """Current web_results diff blocks must populate unique citations.""" + client = self._create_client() + events = iter( + [ + { + "blocks": [ + { + "intended_usage": "web_results", + "diff_block": { + "field": "web_result_block", + "patches": [ + { + "op": "replace", + "path": "", + "value": { + "progress": "DONE", + "web_results": [ + { + "name": "Primary source", + "url": "https://example.com/source", + "snippet": "Evidence", + }, + { + "name": "Duplicate source", + "url": "https://example.com/source", + "snippet": "Duplicate", + }, + ], + }, + } + ], + }, + } + ] + }, + { + "blocks": [ + { + "intended_usage": "ask_text", + "markdown_block": { + "progress": "DONE", + "chunks": [], + "answer": "Answer with source", + }, + } + ] + }, + ] + ) + + with patch.object(client, "ask_stream", return_value=events): + response = client.ask("question") + + assert response.text == "Answer with source" + assert response.citations == [ + { + "title": "Primary source", + "url": "https://example.com/source", + "snippet": "Evidence", + } + ] + + def test_incremental_web_result_patch_adds_citation(self): + """Indexed web_results patches must be parsed as citation items.""" + client = self._create_client() + events = iter( + [ + { + "blocks": [ + { + "intended_usage": "web_results", + "diff_block": { + "field": "web_result_block", + "patches": [ + { + "op": "add", + "path": "/web_results/0", + "value": { + "name": "Incremental source", + "url": "https://example.com/incremental", + "snippet": "Incremental evidence", + }, + } + ], + }, + } + ] + }, + { + "blocks": [ + { + "intended_usage": "ask_text", + "markdown_block": { + "progress": "DONE", + "answer": "Answer", + }, + } + ] + }, + ] + ) + + with patch.object(client, "ask_stream", return_value=events): + response = client.ask("question") + + assert response.citations == [ + { + "title": "Incremental source", + "url": "https://example.com/incremental", + "snippet": "Incremental evidence", + } + ] + + def test_legacy_final_event_still_parses(self): + """Supporting the current schema must not regress legacy FINAL events.""" + client = self._create_client() + nested_steps = [ + { + "step_type": "FINAL", + "content": {"answer": json.dumps({"answer": "Legacy answer"})}, + } + ] + + with patch.object( + client, + "ask_stream", + return_value=iter( + [ + { + "step_type": "FINAL", + "text": json.dumps(nested_steps), + "related_queries": ["next"], + } + ] + ), + ): + response = client.ask("question") + + assert response.text == "Legacy answer" + assert response.related_queries == ["next"] + + def test_transport_timeout_returns_partial_answer(self): + """A transport timeout after text arrives must return marked partial data.""" + client = self._create_client() + + def events(): + yield { + "blocks": [ + { + "intended_usage": "ask_text", + "markdown_block": { + "progress": "IN_PROGRESS", + "chunks": ["Partial answer"], + }, + } + ] + } + raise CurlTimeout("upstream stalled") + + with patch.object(client, "ask_stream", return_value=events()): + response = client.ask("question", timeout_seconds=45) + + assert response.text == "Partial answer" + assert response.partial is True + assert response.warning is not None + assert "45" in response.warning + + def test_transport_timeout_without_text_is_an_error(self): + """A timeout with no answer text must remain an explicit error.""" + client = self._create_client() + + def events(): + raise CurlTimeout("upstream stalled") + yield + + with patch.object(client, "ask_stream", return_value=events()): + with pytest.raises(CurlTimeout, match="no answer text"): + client.ask("question", timeout_seconds=45) + + def test_total_deadline_returns_partial_and_closes_stream(self): + """The client-enforced total deadline must not wait for another chunk.""" + client = self._create_client() + stream_closed = [] + + def events(): + try: + yield { + "blocks": [ + { + "intended_usage": "ask_text", + "markdown_block": { + "progress": "IN_PROGRESS", + "chunks": ["Partial by deadline"], + }, + } + ] + } + raise AssertionError("stream was consumed past the total deadline") + finally: + stream_closed.append(True) + + times = iter([0.0, 46.0]) + with patch.object(client, "ask_stream", return_value=events()): + with patch( + "src.core.perplexity_client.time.monotonic", + side_effect=lambda: next(times, 46.0), + ): + response = client.ask("question", timeout_seconds=45) + + assert response.text == "Partial by deadline" + assert response.partial is True + assert stream_closed == [True] + + def test_ask_stream_forwards_timeout_and_closes_http_response(self): + """The streaming HTTP response must close after iteration completes.""" + client = self._create_client() + http_response = MagicMock() + http_response.status_code = 200 + http_response.iter_content.return_value = [b'data: {"blocks": []}\n'] + + with patch( + "src.core.perplexity_client.cffi_requests.post", + return_value=http_response, + ) as post: + events = list(client.ask_stream("question", timeout_seconds=45)) + + assert events == [{"blocks": []}] + assert post.call_args.kwargs["timeout"] == 45 + assert post.call_args.kwargs["curl_options"] == {CurlOpt.TIMEOUT_MS: 45_000} + http_response.close.assert_called_once_with() + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/test_sse_parser.py b/tests/test_sse_parser.py index 5c57d2f..42cecb4 100644 --- a/tests/test_sse_parser.py +++ b/tests/test_sse_parser.py @@ -1,14 +1,15 @@ """Tests for Perplexity SSE parser.""" import pytest -from src.services.sse_parser import PerplexitySSEParser + from src.models.perplexity_models import ( - PerplexitySSEEvent, - PerplexityBlock, DiffBlock, JSONPatch, MarkdownBlock, + PerplexityBlock, + PerplexitySSEEvent, ) +from src.services.sse_parser import PerplexitySSEParser class TestParseEventData: @@ -375,6 +376,11 @@ def test_parse_block_defaults_diff_field_to_empty_string(self): class TestIsMarkdownBlock: """Tests for is_markdown_block method.""" + def test_plain_ask_text_returns_true(self): + """Should recognize the current plain 'ask_text' answer block.""" + result = PerplexitySSEParser.is_markdown_block("ask_text") + assert result is True + def test_ask_text_markdown_returns_true(self): """Should return True for 'ask_text_markdown'.""" result = PerplexitySSEParser.is_markdown_block("ask_text_markdown")