From 963e2fecfb12fa05cb7127a9c9833327325b027b Mon Sep 17 00:00:00 2001 From: Saryev Rustam Date: Sat, 27 May 2023 17:29:46 +0300 Subject: [PATCH] feat: add Halo spinner when creating vector store - Add Halo spinner when creating vector store to provide visual feedback to the user. - Also update success message to reflect the number of documents in the vector store. --- pyproject.toml | 2 +- talk_codebase/llm.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index c08e9ff..f699cc9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "talk-codebase" -version = "0.1.16" +version = "0.1.17" description = "talk-codebase is a powerful tool for querying and analyzing codebases." authors = ["Saryev Rustam "] readme = "README.md" diff --git a/talk_codebase/llm.py b/talk_codebase/llm.py index 34fbeae..61569bd 100644 --- a/talk_codebase/llm.py +++ b/talk_codebase/llm.py @@ -2,6 +2,7 @@ import os import questionary import tiktoken +from halo import Halo from langchain import FAISS from langchain.callbacks.manager import CallbackManager from langchain.chains import ConversationalRetrievalChain @@ -42,8 +43,10 @@ def create_vector_store(root_dir, openai_api_key): if not approve: exit(0) + spinners = Halo(text='Creating vector store', spinner='dots').start() embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key) db = FAISS.from_documents(texts, embeddings) + spinners.succeed(f"Created vector store with {len(docs)} documents") return db