Generating Citations with a Custom Model Deployment

[bcicc]

What is the issue?

I am unable to generate citations with a custom model deployment by simply yielding a CITATION_GENERATION event, nor by including citations in the STREAM_END event in the model's invoke_chat_stream method.

Here is my code in invoke_chat_stream:

yield {
      "event_type": StreamEvent.CITATION_GENERATION,
      "citations":[{
          "start": 0,
          "end": 57, #len(chunk_text),
          "text": "Carlos Alcaraz won the Wimbledon 2024 men's singles title",
          "document_ids": ["wikipedia.org"]
      }]
  }

I can see in the frontend that the citation-generation event comes through the chat-stream endpoint with the appropriate data, but no UI component is rendered. I think I may be misunderstanding something about how citations are implemented. Help would be appreciated.

Additional information

No response

[tianjing-li]

@bcicc I'll check this out the upcoming week, I don't recall the citations format by heart.

I'm probably going to step through the code with a debugger and check the actual expected format, if you're in a rush my suggestion is to step through it on your side as well.

Just from a glance, your citation looks fine however 🤔

[tianjing-li]

@bcicc The general format is correct, your error is most likely the hardcoded document_ids field you're providing. If you look in handle_stream_citation_generation, you'll see the following lines:

document_ids = event_citation.get("document_ids")
        for document_id in document_ids:
            document = document_ids_to_document.get(document_id, None)

            if document is not None:
                citation.documents.append(document)

The document id needs to exist in the document_ids_to_document dictionary provided to this function, which is generated by handle_stream_search_results. You can always hardcode that dictionary where the citations are handled to test the scenario