scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

add Perplexity citation annotations support (#13225)

Browse source
This commit is contained in:
Sameer Kankute 2025-08-02 21:17:35 +05:30 • committed by GitHub
parent eb173f9155
commit 1e33dc50a0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 393 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

94
litellm/llms/perplexity/chat/transformation.py
View file

@ -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)

303
tests/test_litellm/llms/perplexity/chat/test_perplexity_chat_transformation.py
View file

@ -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
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