Apply RUFF linting

This commit is contained in:
Nielson Janné 2025-03-14 23:21:34 +01:00
parent b07eca8f9b
commit e99d26ed96

View File

@ -18,25 +18,41 @@ logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
parser = argparse.ArgumentParser(description="A Sogeti Nederland Generic RAG demo.") parser = argparse.ArgumentParser(description="A Sogeti Nederland Generic RAG demo.")
parser.add_argument("-b", "--back-end", type=BackendType, choices=list(BackendType), default=BackendType.azure, parser.add_argument(
help="(Cloud) back-end to use. In the case of local, a locally installed ollama will be used.") "-b",
parser.add_argument("-p", "--pdf-data", type=Path, required=True, nargs="+", "--back-end",
help="One or multiple paths to folders or files to use for retrieval. " type=BackendType,
"If a path is a folder, all files in the folder will be used. " choices=list(BackendType),
"If a path is a file, only that file will be used. " default=BackendType.azure,
"If the path is relative it will be relative to the current working directory.") help="(Cloud) back-end to use. In the case of local, a locally installed ollama will be used.",
parser.add_argument("--pdf-chunk_size", type=int, default=1000, )
help="The size of the chunks to split the text into.") parser.add_argument(
parser.add_argument("--pdf-chunk_overlap", type=int, default=200, "-p",
help="The overlap between the chunks.") "--pdf-data",
parser.add_argument("--pdf-add-start-index", action="store_true", type=Path,
help="Add the start index to the metadata of the chunks.") required=True,
parser.add_argument("-w", "--web-data", type=str, nargs="*", default=[], nargs="+",
help="One or multiple URLs to use for retrieval.") help="One or multiple paths to folders or files to use for retrieval. "
parser.add_argument("--web-chunk-size", type=int, default=200, "If a path is a folder, all files in the folder will be used. "
help="The size of the chunks to split the text into.") "If a path is a file, only that file will be used. "
parser.add_argument("-c", "--chroma-db-location", type=Path, default=Path(".chroma_db"), "If the path is relative it will be relative to the current working directory.",
help="file path to store or load a Chroma DB from/to.") )
parser.add_argument("--pdf-chunk_size", type=int, default=1000, help="The size of the chunks to split the text into.")
parser.add_argument("--pdf-chunk_overlap", type=int, default=200, help="The overlap between the chunks.")
parser.add_argument(
"--pdf-add-start-index", action="store_true", help="Add the start index to the metadata of the chunks."
)
parser.add_argument(
"-w", "--web-data", type=str, nargs="*", default=[], help="One or multiple URLs to use for retrieval."
)
parser.add_argument("--web-chunk-size", type=int, default=200, help="The size of the chunks to split the text into.")
parser.add_argument(
"-c",
"--chroma-db-location",
type=Path,
default=Path(".chroma_db"),
help="File path to store or load a Chroma DB from/to.",
)
parser.add_argument("-r", "--reset-chrome-db", action="store_true", help="Reset the Chroma DB.") parser.add_argument("-r", "--reset-chrome-db", action="store_true", help="Reset the Chroma DB.")
args = parser.parse_args() args = parser.parse_args()
@ -60,15 +76,16 @@ def generate(state: State):
docs_content = "\n\n".join(doc.page_content for doc in state["context"]) docs_content = "\n\n".join(doc.page_content for doc in state["context"])
messages = prompt.invoke({"question": state["question"], "context": docs_content}) messages = prompt.invoke({"question": state["question"], "context": docs_content})
response = llm.invoke(messages) response = llm.invoke(messages)
return {"answer": response.content} return {"answer": response.content}
@cl.on_chat_start @cl.on_chat_start
async def on_chat_start(): async def on_chat_start():
vector_store = Chroma(collection_name="generic_rag", vector_store = Chroma(
embedding_function=get_embedding_model(args.back_end), collection_name="generic_rag",
persist_directory=str(args.chroma_db_location)) embedding_function=get_embedding_model(args.back_end),
persist_directory=str(args.chroma_db_location),
)
cl.user_session.set("vector_store", vector_store) cl.user_session.set("vector_store", vector_store)
cl.user_session.set("emb_model", get_embedding_model(args.back_end)) cl.user_session.set("emb_model", get_embedding_model(args.back_end))
@ -103,8 +120,9 @@ async def set_starters():
try: try:
starters.append(cl.Starter(label=starter["label"], message=starter["message"])) starters.append(cl.Starter(label=starter["label"], message=starter["message"]))
except KeyError: except KeyError:
logging.warning("CHAINLIT_STARTERS environment is not a list with " logging.warning(
"dictionaries containing 'label' and 'message' keys.") "CHAINLIT_STARTERS environment is not a list with dictionaries containing 'label' and 'message' keys."
)
return starters return starters