Spaces:
Runtime error
Runtime error
File size: 1,031 Bytes
1e53020 6a57640 1e53020 6a57640 1e53020 6a57640 1e53020 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
from typing import List
from langchain.chains.combine_documents.stuff import StuffDocumentsChain
from langchain.chat_models import ChatOpenAI
from langchain.chat_models.base import BaseChatModel
from langchain.docstore.document import Document
from knowledge_gpt.core.debug import FakeChatModel
def pop_docs_upto_limit(query: str, chain: StuffDocumentsChain, docs: List[Document], max_len: int) -> List[Document]:
"""Pops documents from a list until the final prompt length is less
than the max length."""
token_count: int = chain.prompt_length(docs, question=query) # type: ignore
while token_count > max_len and len(docs) > 0:
docs.pop()
token_count = chain.prompt_length(docs, question=query) # type: ignore
return docs
def get_llm(model: str, **kwargs) -> BaseChatModel:
if model == "debug":
return FakeChatModel()
if "gpt" in model:
return ChatOpenAI(model=model, **kwargs) # type: ignore
raise NotImplementedError(f"Model {model} not supported!")
|