token_limit_retriever limit token number retrieved from retriever

[humblemat810]

Checked

• I searched existing ideas and did not find a similar one
• I added a very descriptive title
• I've clearly described the feature request and motivation for it

Feature request

Given that language models usually limit the number of token in and out altogether, also because models charge by number of token, users will want to limit number of tokens that will goes into the chatmodel/llm.

Motivation

Given that language models usually limit the number of token in and out altogether, also because models charge by number of token, users will want to limit number of tokens that will goes into the chatmodel/llm.

Proposal (If applicable)

I have already implemented a token limit retriever that limit the token number when retrieving data. And would like to add into community contribution.

it is already working in my personal project

Will add more documentation and examples with test if this idea is welcomed

========= in token_limit_retriever.py==========

__all__ = ["TokenLimitRetriver"]
from langchain.schema.retriever import BaseRetriever
from langchain_core.retrievers import RetrieverLike
from langchain.schema.document import Document

from typing import Any, Dict, List, Optional 
from langchain.callbacks.manager import CallbackManagerForRetrieverRun, AsyncCallbackManagerForRetrieverRun

from langchain.pydantic_v1 import BaseModel, root_validator
from typing import Callable, Optional

class TokenLimitRetrieverMixin(BaseModel):
    token_limit : int
    token_cutoff_strategy: str
    remaining_token_fillup_callback : Optional[Callable[[str,int], str]] = None
    # remaining_token_fillup_callback is a callback that take the final document and number of token remaining
    tokeniser_callback : Callable 
    # a tokenliser callable that take in a 
    retriever : RetrieverLike
    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        assert ((callable(values['remaining_token_fillup_callback'])
                   ^ (values['token_cutoff_strategy'] == 'complete_document'))), \
                   'cannot set token_cutoff_strategy to be "complete_document" and set cufoff function at the same time'+\
                   f"{values['token_cutoff_strategy']} and {values['remaining_token_fillup_callback']}"
        assert values["token_cutoff_strategy"] in ['complete_document', 'partial_document']
        return values
    
class TokenLimitRetriver(BaseRetriever, TokenLimitRetrieverMixin):
    """
    Wrap another retriever and limit the token.
    tokenizer callback is require to count the token
    
    remaining_token_fillup_callback is called when there are token quote left for part of the final
    document. A new document containing partial of the final document will be appended to end of 
    the retrieved list of documents.

    
    Args:
        token_cutoff_strategy (Literal['complete_doccument', 'partial_document']):  specify strategy to cutoff, whether whole document or part of the final document
        retriever (RetrieverLike): retriever to be wrapped

    Example:


    """

    def _get_relevant_documents(
            self, query: str, *, run_manager: CallbackManagerForRetrieverRun, **kwargs
    ) -> List[Document]:
        
        
        docs = self.retriever._get_relevant_documents(query, run_manager = run_manager,  **kwargs)
        cummulative_token_len = 0
        len_doc_token = 0
        i = 0
        len_doc_token = len(self.tokeniser_callback(docs[i].page_content))
        while cummulative_token_len + len_doc_token <= self.token_limit:
            
            cummulative_token_len += len_doc_token
            
            # check next
            i+=1
            if i == len(docs):
                break
            len_doc_token = len(self.tokeniser_callback(docs[i].page_content))

        tokens_remaining = self.token_limit - cummulative_token_len
        to_return = docs[:i]
        if self.token_cutoff_strategy == 'partial_document':
            if tokens_remaining > 0 and i < len(docs):
            # chop final document to fit the remaining token
                to_return.append(
                            Document(
                                self.remaining_token_fillup_callback(
                                    docs[i].page_content,tokens_remaining)
                                    )
                        )

        return to_return

    async def _aget_relevant_documents(
            self,
            query: str,
            *,
            run_manager: AsyncCallbackManagerForRetrieverRun,
            **kwargs: Any,
    ) -> List[Document]:
        docs = await self.retriever._aget_relevant_documents(query, run_manager = run_manager,  **kwargs)
        cummulative_token_len = 0
        len_doc_token = 0
        i = 0
        len_doc_token = len(self.tokeniser_callback(docs[i].page_content))
        while cummulative_token_len + len_doc_token <= self.token_limit:
            cummulative_token_len += len_doc_token
            i+=1
            if i == len(docs):
                break
            len_doc_token = len(self.tokeniser_callback(docs[i].page_content))

        tokens_remaining = self.token_limit - cummulative_token_len
        to_return = docs[:i]
        if self.token_cutoff_strategy == 'partial_document':
            if tokens_remaining > 0 and i < len(docs):
            # chop final document to fit the remaining token
                to_return.append(
                            Document(
                                self.remaining_token_fillup_callback(
                                    docs[i].page_content,tokens_remaining)
                                    )
                        )

        return to_return

================ in test_token_limit_retriever.py===========

the tests are not yet written in favour of any testing framework yet.

#%%
from typing import List

from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever
from typing import Dict

from langchain.pydantic_v1 import root_validator
class ToyRetriever(BaseRetriever):
    """A toy retriever that contains the top k documents that contain the user query.

    This retriever only implements the sync method _get_relevant_documents.

    If the retriever were to involve file access or network access, it could benefit
    from a native async implementation of `_aget_relevant_documents`.

    As usual, with Runnables, there's a default async implementation that's provided
    that delegates to the sync implementation running on another thread.
    """

    documents: List[Document] = [
        Document('this is a nice day'),
        Document('Peter looks great today'),
        Document('Will AI take over the world?'),
        Document('Langchain is a useful community contributed library.')]
    
    """List of documents to retrieve from."""
    k: int
    """Number of top results to return"""
    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
       
        return values
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        """Sync implementations for retriever."""
        matching_documents = []
        for document in self.documents:
            if len(matching_documents) > self.k:
                return matching_documents

            if query.lower() in document.page_content.lower():
                matching_documents.append(document)
        return matching_documents

test_retriever_to_be_wrapped = ToyRetriever(k = 4)


#%%
from token_limit_retriever import TokenLimitRetriver

import tiktoken
def test1():
    def tokenize(text):
        enc = tiktoken.encoding_for_model("gpt-3.5-turbo")
        tokens = enc.encode(text)
        return tokens

    def decode_token_upto(text: str, n_token_left: int):
        enc = tiktoken.encoding_for_model("gpt-3.5-turbo")
        return enc.decode(enc.encode(text)[:n_token_left])
    simple_token_limit_chain =TokenLimitRetriver(token_limit=20, remaining_token_fillup_callback=decode_token_upto,
                       tokeniser_callback=tokenize, retriever=test_retriever_to_be_wrapped,
                       token_cutoff_strategy = 'partial_document'
                       )
    res = simple_token_limit_chain.invoke('a')
    assert(len(res) == 4)
    for i in res:
        print(len(tokenize(i.page_content)))
        from pydantic.v1.error_wrappers import ValidationError
    # test for invalid optoins
    try:
        simple_token_limit_chain =TokenLimitRetriver(token_limit=20, remaining_token_fillup_callback=decode_token_upto,
                        tokeniser_callback=tokenize, retriever=test_retriever_to_be_wrapped,
                        token_cutoff_strategy = 'complete_document'
                        )
    
    
        res = simple_token_limit_chain.invoke('a')
    except ValidationError:
        pass
    # test for complete document options
    simple_token_limit_chain =TokenLimitRetriver(token_limit=20,
                       tokeniser_callback=tokenize, retriever=test_retriever_to_be_wrapped,
                       token_cutoff_strategy = 'complete_document'
                       )
    res = simple_token_limit_chain.invoke('a')
    assert(len(res) == 3)

test1()
#%%

def test2():
    def tokenize(text):
        return text.split(' ')

    def decode_token_upto(text: str, n_token_left: int):
        return ' '.join(text.split(' ')[:n_token_left])
    

    simple_token_limit_chain =TokenLimitRetriver(token_limit=20,
                       tokeniser_callback=tokenize, retriever=test_retriever_to_be_wrapped,
                       token_cutoff_strategy = 'complete_document'
                       )
    res = simple_token_limit_chain.invoke('a')
    assert (len(res) == 3)

    simple_token_limit_chain =TokenLimitRetriver(token_limit=20, remaining_token_fillup_callback=decode_token_upto,
                       tokeniser_callback=tokenize, retriever=test_retriever_to_be_wrapped,
                       token_cutoff_strategy = 'partial_document'
                       )
    res = simple_token_limit_chain.invoke('a')
    assert(len(res) == 4)
test2()
# %%