andreped/ReferenceBot

1from typing import List2 3from langchain.chains.combine_documents.stuff import StuffDocumentsChain4from langchain.chat_models import ChatOpenAI5from langchain.chat_models.base import BaseChatModel6from langchain.docstore.document import Document7 8from knowledge_gpt.core.debug import FakeChatModel9 10 11def pop_docs_upto_limit(query: str, chain: StuffDocumentsChain, docs: List[Document], max_len: int) -> List[Document]:12    """Pops documents from a list until the final prompt length is less13    than the max length."""14 15    token_count: int = chain.prompt_length(docs, question=query)  # type: ignore16 17    while token_count > max_len and len(docs) > 0:18        docs.pop()19        token_count = chain.prompt_length(docs, question=query)  # type: ignore20 21    return docs22 23 24def get_llm(model: str, **kwargs) -> BaseChatModel:25    if model == "debug":26        return FakeChatModel()27 28    if "gpt" in model:29        return ChatOpenAI(model=model, **kwargs)  # type: ignore30 31    raise NotImplementedError(f"Model {model} not supported!")32