Perplexity AI Citations

[vikigenius]

Description

Perplexity AI has a citations field in their response that usually works like this

response_stream = client.chat.completions.create(
    model="sonar-pro",
    messages=messages,
    stream=True,
)

citations = []
for response in response_stream:
    print(response.choices[0].delta.content, end="")
    citations.extend(response.citations)

You now have access to a list of citations that you can use reliably.

Is there any way in pydantic_ai to access this?

The perplexity model is not very reliable at producing inline citations. So it would be useful to be able to do this and access their citations.

[Kludex]

Hi @vikigenius,

Can you give me an MRE to run Perplexity AI?

[vikigenius]

If you mean using the standard OpenAI API

import os
from openai import OpenAI

YOUR_API_KEY = os.getenv("PERPLEXITY_API_KEY")

messages = [
    {
        "role": "system",
        "content": (
            "You are an artificial intelligence assistant that helps users answer questions "
        ),
    },
    {   
        "role": "user",
        "content": (
            "What is the geopolitical impact of the Germany election, especially given Trump's behaviour in America."
        ),
    },
]

client = OpenAI(api_key=YOUR_API_KEY, base_url="https://api.perplexity.ai/")

# chat completion with streaming
response_stream = client.chat.completions.create(
    model="sonar-pro",
    messages=messages,
    stream=True,
)

citations = []
for response in response_stream:
    print(response.choices[0].delta.content, end="")
    citations.extend(response.citations)
print(list(set(citations)))

[vikigenius]

I did manage to regex substitute citations into markdown links by creating a custom model like this.

# -*- coding: utf-8 -*-
from __future__ import annotations
import re
from collections.abc import AsyncIterator
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import cast, override
from openai import AsyncStream
from openai.types import chat
from openai.types.chat import ChatCompletionChunk
from pydantic_ai import UnexpectedModelBehavior, usage
from pydantic_ai.messages import ModelResponse, ModelResponsePart, ModelResponseStreamEvent, TextPart, ToolCallPart
from pydantic_ai.models.openai import OpenAIModel, OpenAIStreamedResponse
from pydantic_ai._utils import Unset, PeekableAsyncStream



_citation_pattern: re.Pattern[str] = re.compile(r'\[\^?(\d+)\]')


def _replace_chunk_with_markdown(chunk: str, citations: list[str]):
    """Replaces chunk with cited chunk."""
    def _replace_with_markdown(re_match: re.Match[str]):
        """Function to replace `[^num]` with `[Link num](URL)`"""
        num = int(re_match.group(1))  # Extract the number
        if 1 <= num <= len(citations):  # Ensure it's a valid index
            return f"[[{num}]]({citations[num-1]})"  # Convert to markdown format

        return re_match.group(0)  # Return original if out of range
    return _citation_pattern.sub(_replace_with_markdown, chunk)


@dataclass(init=False)
class PerplexityModel(OpenAIModel):
    """A Perplexity model that uses the OpenAI API.

    Internally, this uses the [OpenAI Python client](https://github.com/openai/openai-python) to interact with the API.

    Apart from `__init__`, all methods are private or match those of the base class.
    """

    @override
    def _process_response(self, response: chat.ChatCompletion) -> ModelResponse:
        """Process a non-streamed response, and prepare a message to return."""
        timestamp = datetime.fromtimestamp(response.created, tz=timezone.utc)
        choice = response.choices[0]
        items: list[ModelResponsePart] = []
        if choice.message.content is not None:
            citations = getattr(response, 'citations', [])
            content =  _replace_chunk_with_markdown(choice.message.content, citations)
            items.append(TextPart(content))
        if choice.message.tool_calls is not None:
            for c in choice.message.tool_calls:
                items.append(ToolCallPart(c.function.name, c.function.arguments, c.id))
        return ModelResponse(items, model_name=response.model, timestamp=timestamp)

    @override
    async def _process_streamed_response(self, response: AsyncStream[ChatCompletionChunk]) -> PerplexityStreamedResponse:
        """Process a streamed response, and prepare a streaming response to return."""
        peekable_response = PeekableAsyncStream(response)
        first_chunk = await peekable_response.peek()
        if isinstance(first_chunk, Unset):
            raise UnexpectedModelBehavior('Streamed response ended without content or tool calls')

        return PerplexityStreamedResponse(
            _model_name=self._model_name,
            _response=peekable_response,
            _timestamp=datetime.fromtimestamp(first_chunk.created, tz=timezone.utc),
        )

@dataclass
class PerplexityStreamedResponse(OpenAIStreamedResponse):
    """Implementation of `StreamedResponse` for OpenAI models."""
    citations: list[str] = field(default_factory=list)

    @override
    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
        async for chunk in self._response:
            self._usage += _map_usage(chunk)

            try:
                choice = chunk.choices[0]
                self.citations = cast(list[str], getattr(chunk, 'citations', []))
            except IndexError:
                continue

            if choice.delta.content is not None:
                # Handle the text part of the response
                content =  _replace_chunk_with_markdown(choice.delta.content, self.citations)
                yield self._parts_manager.handle_text_delta(vendor_part_id='content', content=content)

            for dtc in choice.delta.tool_calls or []:
                maybe_event = self._parts_manager.handle_tool_call_delta(
                    vendor_part_id=dtc.index,
                    tool_name=dtc.function and dtc.function.name,
                    args=dtc.function and dtc.function.arguments,
                    tool_call_id=dtc.id,
                )
                if maybe_event is not None:
                    yield maybe_event


def _map_usage(response: chat.ChatCompletion | ChatCompletionChunk) -> usage.Usage:
    response_usage = response.usage
    if response_usage is None:
        return usage.Usage()
    else:
        details: dict[str, int] = {}
        if response_usage.completion_tokens_details is not None:
            details.update(response_usage.completion_tokens_details.model_dump(exclude_none=True))
        if response_usage.prompt_tokens_details is not None:
            details.update(response_usage.prompt_tokens_details.model_dump(exclude_none=True))
        return usage.Usage(
            request_tokens=response_usage.prompt_tokens,
            response_tokens=response_usage.completion_tokens,
            total_tokens=response_usage.total_tokens,
            details=details,
        )

[jerry-reevo]

@vikigenius Out of curiosity, do you ever see a response chunk split the citation? e.g.,

chunk 1: Statement[1
chunk 2: ][2]. Other text

[vikigenius]

@vikigenius Out of curiosity, do you ever see a response chunk split the citation? e.g.,

chunk 1: Statement[1 chunk 2: ][2]. Other text

Nope, never seen it happen so far

[stuartmaxwell]

I was also looking for this, but I will note that this is an undocumented feature - Perplexity's API docs don't even mention it, even though you can see the citations in the API response.

As a workaround, when using Perplexity with Pydantic_ai, I include some instructions similar to this: "Please format your response as Markdown and include citations in the response using the Github style of including footnotes."