RAG Citation: Enhancing RAG Pipelines with Automatic Citations

Project Overview

RAG Citation combines Retrieval-Augmented Generation (RAG) with automatic citation generation. It enhances the credibility of RAG-generated content by providing relevant citations for the information used in generating responses.

The library supports two citation methods:

• Non-LLM (default): Fast, lightweight citation using SpaCy NER and SentenceTransformers semantic similarity -- no LLM API calls needed.
• LLM-based: Uses any LLM (via LiteLLM) with structured output for citation generation.

Key Features

• Dual Approach: Choose between a fast non-LLM pipeline or an LLM-based pipeline depending on your needs.
• Semantic Search: Identifies relevant source documents based on meaning and context rather than just keyword matching.
• Named Entity Recognition: Extracts and returns relevant named entities from LLM-generated answers, such as people, organizations, money and dates.
• Flexible Integration: Can be easily integrated into any RAG pipeline.
• Custom Embedding Models: Bring your own embedding model by implementing the BaseEmbeddingModel interface.
• Hallucination Detection (Beta): Flags instances where the LLM-generated answer contains entities (DATE, MONEY, CARDINAL, ORDINAL, QUANTITY, TIME) that cannot be found in the context.

Quickstart

• Langchain example: langchain.ipynb
• Embeddchain example: embeddchain.ipynb
• Custom embedding model: custom_embedding_model.ipynb

Installation

pip install rag-citation

For LLM-based citations (includes litellm + pydantic):

pip install rag-citation[llm]

Download a spaCy model (required for the non-LLM method):

python -m spacy download en_core_web_sm   # small
python -m spacy download en_core_web_md   # medium
python -m spacy download en_core_web_lg   # large

Non-LLM Method (Default)

from rag_citation import CiteItem, Inference
import uuid

## Sample context from vectorDB or semantic search
documents = [
    "Elon MuskCEO, Tesla$221.6B$439M (0.20%)Real Time Net Worthas of 8/6/24Reflects change since 5 pm ET of prior trading day. 1 in the world todayPhoto by Martin Schoeller for ForbesAbout Elon MuskElon Musk cofounded six companies, including electric car maker Tesla, rocket producer SpaceX and tunneling startup Boring Company.He owns about 12% of Tesla excluding options, but has pledged more than half his shares as collateral for personal loans of up to $3.5 billion.In early 2024, a Delaware judge voided Musk's 2018 deal to receive options equaling an additional 9% of Tesla.",
    "people in the world; as of August 2024[update], Forbes estimates his net worth to be US$241 billion.[3] Musk was born in Pretoria to model Maye and businessman and engineer Errol Musk, and briefly attended the University of Pretoria before immigrating to Canada at age 18, acquiring citizenship through his Canadian-born mother. Two years later, he matriculated at Queen's University at Kingston in Canada. Musk later transferred to the University of Pennsylvania and received bachelor's degrees in economics and physics."
]

## Example answer generated by an LLM
answer = "Elon Musk's net worth is estimated to be US$241 billion as of August 2024."

## Helper function to create context in the correct format
def format_document(documents):
    context = []
    for document in documents:
        context.append(
            {
                "source_id": str(uuid.uuid4()),
                "document": document,
                "meta": [{"meta-data": "some-info"}],
            }
        )
    return context

context = format_document(documents)
cite_item = CiteItem(answer=answer, context=context)

## Initialize the Inference (non-LLM is the default)
inference = Inference(spacy_model="sm", embedding_model="md")

## Get citation and other information
output = inference(cite_item)

print(output.citation)        # List of citations with source documents
print(output.hallucination)   # True if hallucination detected
print(output.missing_word)    # List of entities not found in context

LLM Method

from rag_citation import CiteItem, Inference

## Same context/answer setup as above...
cite_item = CiteItem(answer=answer, context=context)

## Initialize with LLM method (uses LiteLLM -- any supported provider works)
inference = Inference(
    method="llm",
    model="gpt-4o",              # or "anthropic/claude-sonnet-4-20250514", "azure/gpt-4o", etc.
    api_key="your-api-key",      # optional if set via environment variable
)

output = inference(cite_item)
print(output.citation)

## Optionally pass conversation history for richer context
messages = [
    {"role": "user", "content": "What is Elon Musk's net worth?"},
    {"role": "assistant", "content": answer},
]
output = inference(cite_item, messages=messages)

Output Explanation

output.citation

[
  {
    "answer_sentences": "Elon Musk's net worth is estimated to be US$241 billion as of August 2024.",
    "cite_document": [
      {
        "document": "people in the world; as of August 2024[update], Forbes estimates his net worth to be US$241 billion.[3]",
        "source_id": "23d1f1f0-2afa-4749-8639-78ec685fd837",
        "entity": [
          { "word": "US$241 billion", "entity_name": "MONEY" },
          { "word": "August 2024", "entity_name": "DATE" }
        ],
        "meta": [{ "url": "https://www.forbes.com/profile/elon-musk/" }]
      }
    ]
  }
]

Key	Description
answer_sentences	Sentence from the answer that is cited.
cite_document	List of source documents supporting the sentence. Each contains:
	- document: Text from the source document.
	- source_id: Unique identifier for the source document.
	- entity: Named entities found (with word and entity_name).
	- meta: Metadata passed in the original context.

output.hallucination

False — Indicates whether the output contains hallucinated information.

output.missing_word

[] — List of entities (DATE, MONEY, etc.) present in the answer but not found in any source document.

Configuration

Non-LLM Parameters

Parameter	Description	Default
spacy_model	SpaCy model size: "sm", "md", or "lg"	"sm"
embedding_model	Embedding model size ("sm", "md", "lg") or a custom BaseEmbeddingModel instance	"sm"
therhold_value	Cosine similarity threshold for semantic matching	0.88

Embedding model mapping:

Alias	Model
"sm"	avsolatorio/GIST-small-Embedding-v0
"md"	avsolatorio/GIST-Embedding-v0
"lg"	avsolatorio/GIST-large-Embedding-v0

LLM Parameters

Parameter	Description	Default
model	LiteLLM model identifier (e.g., "gpt-4o", "anthropic/claude-sonnet-4-20250514")	Required
api_key	API key for the LLM provider	None (reads from env)
temperature	LLM temperature	0.0
max_tokens	Maximum tokens for LLM response	4096
**litellm_kwargs	Additional LiteLLM parameters (e.g., api_base, api_version)	—

Custom Embedding Model

You can use your own embedding model by implementing the BaseEmbeddingModel interface:

from rag_citation.base_model import BaseEmbeddingModel

class MyEmbeddingModel(BaseEmbeddingModel):
    def embedding(self, sentence: str):
        # Return embeddings as a tensor
        ...

inference = Inference(spacy_model="sm", embedding_model=MyEmbeddingModel())

Running Tests

cd rag-citation

# Test non-LLM method (no API key needed)
python test/non_llm.py

# Test LLM method (set your model/api_key in the file first)
python test/llm.py

Makefile Commands

Command	Description
make build	Build the .whl and .tar.gz distributions
make install	Install the built package
make install-llm	Install with LLM extras (litellm, pydantic)
make clean	Remove build artifacts

Contributing

We welcome contributions! Here's how you can help:

• Report Bugs: Submit issues on GitHub.
• Suggest Features: Open an issue with your ideas.
• Code Contributions: Fork, make changes, and submit a pull request.
• Documentation: Update and enhance our docs.

License

This project is licensed under the MIT License.

Acknowledgements

• SpaCy
• SentenceTransformers
• LiteLLM
• Huggingface/avsolatorio