Vectara Ragna Plugin

[peachkeel]

I understand that adding external source storages is not currently a priority for the project, yet it was at the top of my TODO list. Therefore, please find below a plugin that enables Ragna to store and retrieve items from Vectara. This code has minimal dependencies, and I hope it can be included in a future release. I welcome any comments, questions, bug reports, or concerns.

# Requirements:
#   pip install Authlib [1.2.1]

import json
import os
import requests
import uuid

from authlib.integrations.requests_client import OAuth2Session

from ragna.core import Config, Document, Source
from ragna.source_storages._vector_database import VectorDatabaseSourceStorage

#
# NOTE: Environmental variables VECTARA_CUSTOMER_ID, VECTARA_CLIENT_ID, and
#       VECTARA_CLIENT_SECRET must all be set for this code to work.
#
class Vectara(VectorDatabaseSourceStorage):

    def __init__(
            self,
            config: Config,
            *,
            customer_id=os.environ.get("VECTARA_CUSTOMER_ID"),
            client_id=os.environ.get("VECTARA_CLIENT_ID"),
            client_secret=os.environ.get("VECTARA_CLIENT_SECRET")
        ) -> None:
        super().__init__(config)
        self._vcustid = customer_id
        self._vclntid = client_id
        self._vclntxx = client_secret
        self._jwt_token = self._get_jwt_token()
        self._session = requests.Session()

    def _get_jwt_token(self):
        auth_url = f"https://vectara-prod-{self._vcustid}.auth.us-west-2.amazoncognito.com"
        token_endpoint = f"{auth_url}/oauth2/token"
        session = OAuth2Session(self._vclntid, self._vclntxx, scope="")
        token = session.fetch_token(token_endpoint, grant_type="client_credentials")
        return token["access_token"]

    def _do_post(self, endpoint, body):
        response = self._session.post(
            "https://api.vectara.io" + endpoint,
            data=json.dumps(body),
            verify=True,
            headers={
                "customer-id": self._vcustid,
                "Authorization": f"Bearer {self._jwt_token}"
            }
        )
        if response.status_code == 401: # JWT token expired?
            self._jwt_token = self._get_jwt_token()
            return self._do_post(endpoint, body)
        else:
            response.raise_for_status()
            return response.json()

    def _rest_list_corpus(self, name):
        result = self._do_post("/v1/list-corpora",
            {
                "numResults": 1,
                "filter": name
            }
        )
        return None if not result["corpus"] else result["corpus"][0]["id"]

    def _rest_create_corpus(self, name):
        result = self._do_post("/v1/create-corpus",
            {
                "corpus": {
                    "name": name,
                    "description": "Represents a chat in the Ragna app."
                }
            }
        )
        return result["corpusId"]

    ############################################################################
    # [INPUT] Document Example
    ############################################################################
    #{'handler': <ragna.core.TxtDocumentHandler object at 0x7f8796f9fa00>,
    # 'id': UUID('c523a076-72e3-422d-b1c9-f6968a2884f7'),
    # 'metadata': {'path': '/home/ubuntu/.cache/ragna/documents/c523a076-72e3-422d-b1c9-f6968a2884f7'},
    # 'name': 'Tracy L.doc.txt'}
    ############################################################################
    def store(
        self,
        documents: list[Document],
        *,
        chat_id: uuid.UUID,
        chunk_size: int = 500,    # keep enabled in UI for retrieve() control
        chunk_overlap: int = None # disable in UI since Vectara handles
    ) -> None:
        vcorpid = self._rest_list_corpus(str(chat_id))
        assert vcorpid is None, "Corpus for chat predates chat!"
        vcorpid = self._rest_create_corpus(str(chat_id))

        for document in documents:
            vdoc = (
                {
                    "customerId": self._vcustid,
                    "corpusId": vcorpid,
                    "document": {
                        "documentId": str(document.id),
                        "section": [
                            {"text": page.text}
                            for page in document.extract_pages()
                        ]
                    }
                }
            )
            self._do_post("/v1/index", vdoc)

    ############################################################################
    # [OUTPUT] Source Example
    ############################################################################
    #{'content': 'Tracy L. Geaccone I graduated from University of [...]',
    # 'document': <ragna.core.LocalDocument object at 0x7f87742cf640>,
    # 'id': '21125af7-8d8a-42a9-ac2a-2a7ee32d2f16',
    # 'location': '',
    # 'num_tokens': 110}
    ############################################################################
    def retrieve(
        self,
        documents: list[Document],
        prompt: str,
        *,
        chat_id: uuid.UUID,
        chunk_size: int = 500, # keep enabled in UI
        num_tokens: int = 4000 # keep enabled in UI
    ) -> list[Source]:
        vcorpid = self._rest_list_corpus(str(chat_id))
        assert vcorpid is not None, "Chat has no corresponding corpus!"

        est_words_per_token = 0.75
        est_words_per_sentence = 15
        est_words_per_chunk = chunk_size * est_words_per_token
        est_sentences_per_chunk = est_words_per_chunk / est_words_per_sentence

        vquery = (
            {
                "query": [
                    {
                        "query": prompt,
                        "start": 0,
                        "numResults": max(int(num_tokens * 2 / chunk_size), 5),
                        "contextConfig": {
                            "sentencesBefore": int(est_sentences_per_chunk / 2),
                            "sentencesAfter": int(est_sentences_per_chunk / 2)
                        },
                        "corpusKey": [
                            {
                                "customerId": self._vcustid,
                                "corpusId": vcorpid,
                                "lexicalInterpolationConfig": {
                                    "lambda": 0.025
                                }
                            }
                        ]
                    }
                ]
            }
        )

        result = self._do_post("/v1/query", vquery)

        document_map = {str(document.id): document for document in documents}
        sources = self._take_sources_up_to_max_tokens(
            (
                Source(
                    id=str(uuid.uuid4()),
                    document=document_map[vdoc["id"]],
                    location='',
                    content=vresp["text"],
                    num_tokens=len(self._tokenizer.encode(vresp["text"]))
                )
                for vresp, vdoc in zip(
                    result["responseSet"][0]["response"],
                    result["responseSet"][0]["document"]
                )
            ),
            max_tokens=num_tokens
        )

        return sources

[pmeier]

Thanks for taking the time to post this for everyone to use!

I can't fully review, but here are a few comments to improve if you want to:

• You can declare environment variable requirements. For example
  

  ragna/ragna/assistants/_api.py

  Lines 19 to 21 in 4962665

  	@classmethod
  	def requirements(cls) -> list[Requirement]:
  	return [EnvVarRequirement(cls._API_KEY_ENV_VAR)]

  That way, you can be sure that these environment variables are available when the component is instantiated and you will get a proper error message if not.

• requests.post has a json= parameter so you don't have to dump the body manually.

• You can put the auth token as session header so you don't have to explicitly pass it with every request

• Similarly, you can also put a base url for the session if you are only hitting a single API (which you do)

• Your _do_post method has a potential infinite loop if the cause of the 401 is not an expired token

• You should not need to keep chunk_size on the .store() method if you use it as part of .retrieve(). If you do, please file an issue. Unfortunately however, you need to keep both chunk_overlap, chunk_size, and num_tokens on any method there, because they are hardcoded in the UI for now. We are working on this.

• Not sure about your chunk_size usage in .retrieve. Since you don't control the chunking in .store it makes little sense to use it to estimate the number of documents you want to pull? In the future when we have fixed the UI, you probably want to just have a parameter to let the user select how many source to pull at maximum? Or is there any indication in Vectara how long a chunk is?