# 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](https://github.com/rahulanand1103/rag-citation/blob/main/docs/examples/3.example-langchain.ipynb)
* **Embeddchain example**: [embeddchain.ipynb](https://github.com/rahulanand1103/rag-citation/blob/main/docs/examples/2.example-embeddchain.ipynb)
* **Custom embedding model**: [custom_embedding_model.ipynb](https://github.com/rahulanand1103/rag-citation/blob/main/docs/examples/4.example-custom_embedding_model.ipynb)

### Installation

```bash
pip install rag-citation
```

For LLM-based citations (includes litellm + pydantic):
```bash
pip install rag-citation[llm]
```

Download a spaCy model (required for the non-LLM method):
```bash
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)

```python
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

```python
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`
```json
[
  {
    "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:

```python
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

```bash
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](LICENSE).

## Acknowledgements

- [SpaCy](https://spacy.io/)
- [SentenceTransformers](https://www.sbert.net/)
- [LiteLLM](https://github.com/BerriAI/litellm)
- [Huggingface/avsolatorio](https://huggingface.co/avsolatorio)
