diff --git a/litellm/llms/perplexity/chat/transformation.py b/litellm/llms/perplexity/chat/transformation.py index 955fdff0818..27e6415ff8b 100644 --- a/litellm/llms/perplexity/chat/transformation.py +++ b/litellm/llms/perplexity/chat/transformation.py @@ -13,6 +13,8 @@ from litellm.types.utils import Usage, PromptTokensDetailsWrapper from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.utils import ModelResponse +from litellm.types.llms.openai import ChatCompletionAnnotation +from litellm.types.llms.openai import ChatCompletionAnnotationURLCitation class PerplexityChatConfig(OpenAIGPTConfig): @@ -102,7 +104,10 @@ class PerplexityChatConfig(OpenAIGPTConfig): # Extract and enhance usage with Perplexity-specific fields try: raw_response_json = raw_response.json() - self._enhance_usage_with_perplexity_fields(model_response, raw_response_json) + self._enhance_usage_with_perplexity_fields( + model_response, raw_response_json + ) + self._add_citations_as_annotations(model_response, raw_response_json) except Exception as e: verbose_logger.debug(f"Error extracting Perplexity-specific usage fields: {e}") @@ -131,7 +136,9 @@ class PerplexityChatConfig(OpenAIGPTConfig): if citations: # Count total characters in citations as a proxy for citation tokens # This is an estimation - in practice, you might want to use proper tokenization - total_citation_chars = sum(len(str(citation)) for citation in citations if citation) + total_citation_chars = sum( + len(str(citation)) for citation in citations if citation + ) # Rough estimation: ~4 characters per token (OpenAI's general rule) if total_citation_chars > 0: citation_tokens = max(1, total_citation_chars // 4) @@ -150,7 +157,9 @@ class PerplexityChatConfig(OpenAIGPTConfig): num_search_queries = raw_response_json.get("search_queries") # Create or update prompt_tokens_details to include web search requests and citation tokens - if citation_tokens > 0 or (num_search_queries is not None and num_search_queries > 0): + if citation_tokens > 0 or ( + num_search_queries is not None and num_search_queries > 0 + ): if usage.prompt_tokens_details is None: usage.prompt_tokens_details = PromptTokensDetailsWrapper() @@ -161,3 +170,82 @@ class PerplexityChatConfig(OpenAIGPTConfig): # Store search queries count in the standard web_search_requests field if num_search_queries is not None and num_search_queries > 0: usage.prompt_tokens_details.web_search_requests = num_search_queries + + def _add_citations_as_annotations( + self, model_response: ModelResponse, raw_response_json: dict + ) -> None: + """ + Extract citations and search_results from Perplexity API response + and add them as ChatCompletionAnnotation objects to the message. + """ + if not model_response.choices: + return + + # Get the first choice (assuming single response) + choice = model_response.choices[0] + if not hasattr(choice, "message") or choice.message is None: + return + + message = choice.message + annotations = [] + + # Extract citations from the response + citations = raw_response_json.get("citations", []) + search_results = raw_response_json.get("search_results", []) + + # Create a mapping of URLs to search result titles + url_to_title = {} + for result in search_results: + if isinstance(result, dict) and "url" in result and "title" in result: + url_to_title[result["url"]] = result["title"] + + # Get the message content to find citation positions + content = getattr(message, "content", "") + if not content: + return + + # Find all citation markers like [1], [2], [3], [4] in the text + import re + + citation_pattern = r"\[(\d+)\]" + citation_matches = list(re.finditer(citation_pattern, content)) + + # Create a mapping of citation numbers to URLs + citation_number_to_url = {} + for i, citation in enumerate(citations): + if isinstance(citation, str): + citation_number_to_url[i + 1] = citation # 1-indexed + + # Create annotations for each citation match found in the text + for match in citation_matches: + citation_number = int(match.group(1)) + if citation_number in citation_number_to_url: + url = citation_number_to_url[citation_number] + title = url_to_title.get(url, "") + + # Create the URL citation annotation with actual text positions + url_citation: ChatCompletionAnnotationURLCitation = { + "url": url, + "title": title, + "start_index": match.start(), + "end_index": match.end(), + } + + annotation: ChatCompletionAnnotation = { + "type": "url_citation", + "url_citation": url_citation, + } + + annotations.append(annotation) + + # Add annotations to the message if we have any + if annotations: + if not hasattr(message, "annotations") or message.annotations is None: + message.annotations = [] + message.annotations.extend(annotations) + + # Also add the raw citations and search_results as attributes for backward compatibility + if citations: + setattr(model_response, "citations", citations) + if search_results: + setattr(model_response, "search_results", search_results) \ No newline at end of file diff --git a/tests/test_litellm/llms/perplexity/chat/test_perplexity_chat_transformation.py b/tests/test_litellm/llms/perplexity/chat/test_perplexity_chat_transformation.py index 6f64f46b4a7..784e6f6fe63 100644 --- a/tests/test_litellm/llms/perplexity/chat/test_perplexity_chat_transformation.py +++ b/tests/test_litellm/llms/perplexity/chat/test_perplexity_chat_transformation.py @@ -406,4 +406,305 @@ class TestPerplexityChatTransformation: assert model_response.usage.prompt_tokens_details is not None web_search_requests = model_response.usage.prompt_tokens_details.web_search_requests - assert web_search_requests == 4 \ No newline at end of file + assert web_search_requests == 4 + + # Tests for citation annotations functionality + def test_add_citations_as_annotations_basic(self): + """Test basic citation annotation creation.""" + config = PerplexityChatConfig() + + # Create a ModelResponse with content + from litellm.types.utils import Choices, Message + message = Message(content="This response has citations[1][2] in the text.", role="assistant") + choice = Choices(finish_reason="stop", index=0, message=message) + model_response = ModelResponse() + model_response.choices = [choice] + + # Mock raw response with citations and search results + raw_response_json = { + "citations": [ + "https://example.com/page1", + "https://example.com/page2" + ], + "search_results": [ + {"title": "Example Page 1", "url": "https://example.com/page1"}, + {"title": "Example Page 2", "url": "https://example.com/page2"} + ] + } + + # Add citations as annotations + config._add_citations_as_annotations(model_response, raw_response_json) + + # Check that annotations were created + annotations = getattr(message, 'annotations', None) + assert annotations is not None + assert len(annotations) == 2 + + # Check first annotation + annotation1 = annotations[0] + assert annotation1['type'] == 'url_citation' + url_citation1 = annotation1['url_citation'] + assert url_citation1['url'] == "https://example.com/page1" + assert url_citation1['title'] == "Example Page 1" + # Check that start_index and end_index are valid positions + assert url_citation1['start_index'] >= 0 + assert url_citation1['end_index'] > url_citation1['start_index'] + # Verify the positions correspond to [1] in the text + assert message.content[url_citation1['start_index']:url_citation1['end_index']] == "[1]" + + # Check second annotation + annotation2 = annotations[1] + assert annotation2['type'] == 'url_citation' + url_citation2 = annotation2['url_citation'] + assert url_citation2['url'] == "https://example.com/page2" + assert url_citation2['title'] == "Example Page 2" + # Check that start_index and end_index are valid positions + assert url_citation2['start_index'] >= 0 + assert url_citation2['end_index'] > url_citation2['start_index'] + # Verify the positions correspond to [2] in the text + assert message.content[url_citation2['start_index']:url_citation2['end_index']] == "[2]" + + # Check backward compatibility + assert hasattr(model_response, 'citations') + assert hasattr(model_response, 'search_results') + assert model_response.citations == raw_response_json['citations'] + assert model_response.search_results == raw_response_json['search_results'] + + def test_add_citations_as_annotations_empty_citations(self): + """Test handling of empty citations array.""" + config = PerplexityChatConfig() + + # Create a ModelResponse with content + from litellm.types.utils import Choices, Message + message = Message(content="This response has citations[1][2] but no citations array.", role="assistant") + choice = Choices(finish_reason="stop", index=0, message=message) + model_response = ModelResponse() + model_response.choices = [choice] + + # Mock raw response with empty citations + raw_response_json = { + "citations": [], + "search_results": [] + } + + # Add citations as annotations + config._add_citations_as_annotations(model_response, raw_response_json) + + # Check that no annotations were created + annotations = getattr(message, 'annotations', None) + assert annotations is None or len(annotations) == 0 + + def test_add_citations_as_annotations_no_citation_patterns(self): + """Test handling when text has no citation patterns.""" + config = PerplexityChatConfig() + + # Create a ModelResponse with content without citation patterns + from litellm.types.utils import Choices, Message + message = Message(content="This response has no citation markers in the text.", role="assistant") + choice = Choices(finish_reason="stop", index=0, message=message) + model_response = ModelResponse() + model_response.choices = [choice] + + # Mock raw response with citations + raw_response_json = { + "citations": [ + "https://example.com/page1", + "https://example.com/page2" + ], + "search_results": [ + {"title": "Example Page 1", "url": "https://example.com/page1"}, + {"title": "Example Page 2", "url": "https://example.com/page2"} + ] + } + + # Add citations as annotations + config._add_citations_as_annotations(model_response, raw_response_json) + + # Check that no annotations were created + annotations = getattr(message, 'annotations', None) + assert annotations is None or len(annotations) == 0 + + def test_add_citations_as_annotations_mismatched_numbers(self): + """Test handling of citation numbers that don't match available citations.""" + config = PerplexityChatConfig() + + # Create a ModelResponse with content + from litellm.types.utils import Choices, Message + message = Message(content="This response has citations[1][5] but only 3 citations available.", role="assistant") + choice = Choices(finish_reason="stop", index=0, message=message) + model_response = ModelResponse() + model_response.choices = [choice] + + # Mock raw response with only 3 citations + raw_response_json = { + "citations": [ + "https://example.com/page1", + "https://example.com/page2", + "https://example.com/page3" + ], + "search_results": [ + {"title": "Example Page 1", "url": "https://example.com/page1"}, + {"title": "Example Page 2", "url": "https://example.com/page2"}, + {"title": "Example Page 3", "url": "https://example.com/page3"} + ] + } + + # Add citations as annotations + config._add_citations_as_annotations(model_response, raw_response_json) + + # Check that only one annotation was created (for [1]) + annotations = getattr(message, 'annotations', None) + assert annotations is not None + assert len(annotations) == 1 + + # Check the annotation + annotation = annotations[0] + assert annotation['type'] == 'url_citation' + url_citation = annotation['url_citation'] + assert url_citation['url'] == "https://example.com/page1" + assert url_citation['title'] == "Example Page 1" + + def test_add_citations_as_annotations_missing_titles(self): + """Test handling when search results don't have titles.""" + config = PerplexityChatConfig() + + # Create a ModelResponse with content + from litellm.types.utils import Choices, Message + message = Message(content="This response has citations[1][2] with search results but no titles.", role="assistant") + choice = Choices(finish_reason="stop", index=0, message=message) + model_response = ModelResponse() + model_response.choices = [choice] + + # Mock raw response with missing titles + raw_response_json = { + "citations": [ + "https://example.com/page1", + "https://example.com/page2" + ], + "search_results": [ + {"url": "https://example.com/page1"}, # No title + {"title": "Example Page 2", "url": "https://example.com/page2"} + ] + } + + # Add citations as annotations + config._add_citations_as_annotations(model_response, raw_response_json) + + # Check that annotations were created + annotations = getattr(message, 'annotations', None) + assert annotations is not None + assert len(annotations) == 2 + + # Check first annotation (no title) + annotation1 = annotations[0] + url_citation1 = annotation1['url_citation'] + assert url_citation1['title'] == "" # Empty title for missing title + + # Check second annotation (has title) + annotation2 = annotations[1] + url_citation2 = annotation2['url_citation'] + assert url_citation2['title'] == "Example Page 2" + + def test_add_citations_as_annotations_non_numeric_patterns(self): + """Test handling of non-numeric citation patterns.""" + config = PerplexityChatConfig() + + # Create a ModelResponse with content containing non-numeric patterns + from litellm.types.utils import Choices, Message + message = Message(content="This response has patterns: [a] [b] [1] [c] [2].", role="assistant") + choice = Choices(finish_reason="stop", index=0, message=message) + model_response = ModelResponse() + model_response.choices = [choice] + + # Mock raw response with citations + raw_response_json = { + "citations": [ + "https://example.com/page1", + "https://example.com/page2" + ], + "search_results": [ + {"title": "Example Page 1", "url": "https://example.com/page1"}, + {"title": "Example Page 2", "url": "https://example.com/page2"} + ] + } + + # Add citations as annotations + config._add_citations_as_annotations(model_response, raw_response_json) + + # Check that only numeric patterns were processed + annotations = getattr(message, 'annotations', None) + assert annotations is not None + assert len(annotations) == 2 # Only [1] and [2] should be processed + + # Check that the annotations correspond to [1] and [2] + urls = [ann['url_citation']['url'] for ann in annotations] + assert "https://example.com/page1" in urls + assert "https://example.com/page2" in urls + + def test_add_citations_as_annotations_empty_content(self): + """Test handling of empty content.""" + config = PerplexityChatConfig() + + # Create a ModelResponse with empty content + from litellm.types.utils import Choices, Message + message = Message(content="", role="assistant") + choice = Choices(finish_reason="stop", index=0, message=message) + model_response = ModelResponse() + model_response.choices = [choice] + + # Mock raw response with citations + raw_response_json = { + "citations": ["https://example.com/page1"], + "search_results": [{"title": "Example Page 1", "url": "https://example.com/page1"}] + } + + # Add citations as annotations + config._add_citations_as_annotations(model_response, raw_response_json) + + # Check that no annotations were created + annotations = getattr(message, 'annotations', None) + assert annotations is None or len(annotations) == 0 + + def test_add_citations_as_annotations_no_choices(self): + """Test handling when model_response has no choices.""" + config = PerplexityChatConfig() + + # Create a ModelResponse without choices + model_response = ModelResponse() + model_response.choices = [] # Explicitly set empty choices + + # Mock raw response with citations + raw_response_json = { + "citations": ["https://example.com/page1"], + "search_results": [{"title": "Example Page 1", "url": "https://example.com/page1"}] + } + + # Should not raise an error + config._add_citations_as_annotations(model_response, raw_response_json) + + # No annotations should be created since choices is empty + assert len(model_response.choices) == 0 + + def test_add_citations_as_annotations_no_message(self): + """Test handling when choice has no message.""" + config = PerplexityChatConfig() + + # Create a ModelResponse with choice but no message + from litellm.types.utils import Choices + choice = Choices(finish_reason="stop", index=0, message=None) + model_response = ModelResponse() + model_response.choices = [choice] + + # Mock raw response with citations + raw_response_json = { + "citations": ["https://example.com/page1"], + "search_results": [{"title": "Example Page 1", "url": "https://example.com/page1"}] + } + + # Should not raise an error + config._add_citations_as_annotations(model_response, raw_response_json) + + # Check that no annotations were created (message content is None) + assert choice.message.content is None + # No annotations should be created since content is None + assert not hasattr(choice.message, 'annotations') or choice.message.annotations is None \ No newline at end of file