Want to use RAPTOR for legal research. How to add legislation citation? about raptor HOT 2 CLOSED

from raptor import RetrievalAugmentation, RetrievalAugmentationConfig, BaseSummarizationModel
from openai import OpenAI
from tenacity import retry, stop_after_attempt, wait_random_exponential

class LegalGPT3TurboSummarizationModel(BaseSummarizationModel):
    def __init__(self, model="gpt-3.5-turbo"):
        self.model = model

    @retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(6))
    def summarize(self, context, max_tokens=500, stop_sequence=None):
        try:
            client = OpenAI()
            response = client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": "You are a helpful legal research assistant."},
                    {
                        "role": "user",
                        "content": f"Write a summary of the following legal text, making sure to include all relevant section numbers and legislation names in the summary: {context}",
                    },
                ],
                max_tokens=max_tokens,
            )
            return response.choices[0].message.content
        except Exception as e:
            print(e)
            return e

RAC = RetrievalAugmentationConfig(summarization_model=LegalGPT3TurboSummarizationModel())
RA = RetrievalAugmentation(config=RAC)
RA.add_documents(text)