From b4b654828ce4e7d13f8d42c40acf0b8db9912403 Mon Sep 17 00:00:00 2001 From: Kaveen Kumarasinghe Date: Sun, 5 Feb 2023 14:28:48 -0500 Subject: [PATCH] Full usage metrics reporting, nodes param for search, default 2 for search --- cogs/commands.py | 12 ++++++++++-- cogs/search_service_cog.py | 8 +++++--- gpt3discord.py | 2 +- models/index_model.py | 18 ++++++++++++------ models/search_model.py | 16 ++++++++++++---- 5 files changed, 40 insertions(+), 16 deletions(-) diff --git a/cogs/commands.py b/cogs/commands.py index 159cfd1..89b7a4c 100644 --- a/cogs/commands.py +++ b/cogs/commands.py @@ -840,6 +840,14 @@ class Commands(discord.Cog, name="Commands"): max_value=8, min_value=1, ) + @discord.option( + name="nodes", + description="The higher the number, the more accurate the results, but more expensive", + required=False, + input_type=discord.SlashCommandOptionType.integer, + max_value=10, + min_value=1, + ) @discord.guild_only() - async def search(self, ctx: discord.ApplicationContext, query: str, scope: int): - await self.search_cog.search_command(ctx, query, scope) + async def search(self, ctx: discord.ApplicationContext, query: str, scope: int, nodes:int): + await self.search_cog.search_command(ctx, query, scope, nodes) diff --git a/cogs/search_service_cog.py b/cogs/search_service_cog.py index b1b0958..f84b84a 100644 --- a/cogs/search_service_cog.py +++ b/cogs/search_service_cog.py @@ -20,13 +20,15 @@ class SearchService(discord.Cog, name="SearchService"): self, bot, gpt_model, + usage_service, ): super().__init__() self.bot = bot - self.model = Search(gpt_model) + self.usage_service = usage_service + self.model = Search(gpt_model, usage_service) # Make a mapping of all the country codes and their full country names: - async def search_command(self, ctx, query, search_scope): + async def search_command(self, ctx, query, search_scope, nodes): """Command handler for the translation command""" user_api_key = None if USER_INPUT_API_KEYS: @@ -45,7 +47,7 @@ class SearchService(discord.Cog, name="SearchService"): await ctx.defer() - response = await self.model.search(query, user_api_key, search_scope) + response = await self.model.search(query, user_api_key, search_scope,nodes) await ctx.respond( f"**Query:**\n\n{query.strip()}\n\n**Query response:**\n\n{response.response.strip()}" diff --git a/gpt3discord.py b/gpt3discord.py index 7c8efcf..d0fe7c0 100644 --- a/gpt3discord.py +++ b/gpt3discord.py @@ -171,7 +171,7 @@ async def main(): EnvService.get_google_search_api_key() and EnvService.get_google_search_engine_id() ): - bot.add_cog(SearchService(bot, model)) + bot.add_cog(SearchService(bot, model, usage_service)) print("The Search service is enabled.") bot.add_cog( diff --git a/models/index_model.py b/models/index_model.py index 50068b5..0d9265e 100644 --- a/models/index_model.py +++ b/models/index_model.py @@ -27,7 +27,7 @@ from gpt_index import ( LLMPredictor, QueryConfig, PromptHelper, - IndexStructType, + IndexStructType, OpenAIEmbedding, ) from gpt_index.readers.web import DEFAULT_WEBSITE_EXTRACTOR @@ -38,7 +38,7 @@ from services.environment_service import EnvService, app_root_path SHORT_TO_LONG_CACHE = {} -def get_and_query(user_id, index_storage, query, response_mode, nodes, llm_predictor): +def get_and_query(user_id, index_storage, query, response_mode, nodes, llm_predictor, embed_model): index: [GPTSimpleVectorIndex, ComposableGraph] = index_storage[ user_id ].get_index_or_throw() @@ -49,6 +49,7 @@ def get_and_query(user_id, index_storage, query, response_mode, nodes, llm_predi verbose=True, child_branch_factor=2, llm_predictor=llm_predictor, + embed_model=embed_model, prompt_helper=prompthelper, ) else: @@ -57,6 +58,7 @@ def get_and_query(user_id, index_storage, query, response_mode, nodes, llm_predi response_mode=response_mode, verbose=True, llm_predictor=llm_predictor, + embed_model=embed_model, similarity_top_k=nodes, prompt_helper=prompthelper, ) @@ -331,9 +333,9 @@ class Index_handler: if isinstance(_index.docstore.get_document(doc_id), Document) ] llm_predictor = LLMPredictor(llm=OpenAI(model_name="text-davinci-003")) - tree_index = GPTTreeIndex(documents=documents, llm_predictor=llm_predictor) - print("The last token usage was ", llm_predictor.last_token_usage) - await self.usage_service.update_usage(llm_predictor.last_token_usage) + embedding_model = OpenAIEmbedding() + tree_index = GPTTreeIndex(documents=documents, llm_predictor=llm_predictor, embed_model=embedding_model) + await self.usage_service.update_usage(llm_predictor.last_token_usage+embedding_model.last_token_usage) # Now we have a list of tree indexes, we can compose them if not name: @@ -354,8 +356,10 @@ class Index_handler: if isinstance(_index.docstore.get_document(doc_id), Document) ] + embedding_model = OpenAIEmbedding() # Add everything into a simple vector index - simple_index = GPTSimpleVectorIndex(documents=documents) + simple_index = GPTSimpleVectorIndex(documents=documents, embed_model=embedding_model) + await self.usage_service.update_usage(embedding_model.last_token_usage) if not name: name = f"composed_index_{date.today().month}_{date.today().day}.json" @@ -410,6 +414,7 @@ class Index_handler: try: llm_predictor = LLMPredictor(llm=OpenAI(model_name="text-davinci-003")) + embedding_model = OpenAIEmbedding() response = await self.loop.run_in_executor( None, partial( @@ -420,6 +425,7 @@ class Index_handler: response_mode, nodes, llm_predictor, + embedding_model ), ) print("The last token usage was ", llm_predictor.last_token_usage) diff --git a/models/search_model.py b/models/search_model.py index 5498386..f61e8e0 100644 --- a/models/search_model.py +++ b/models/search_model.py @@ -11,17 +11,19 @@ from gpt_index import ( QuestionAnswerPrompt, GPTSimpleVectorIndex, BeautifulSoupWebReader, - Document, + Document, PromptHelper, LLMPredictor, OpenAIEmbedding, ) from gpt_index.readers.web import DEFAULT_WEBSITE_EXTRACTOR +from langchain import OpenAI from services.environment_service import EnvService from services.usage_service import UsageService class Search: - def __init__(self, gpt_model): + def __init__(self, gpt_model, usage_service): self.model = gpt_model + self.usage_service = usage_service self.google_search_api_key = EnvService.get_google_search_api_key() self.google_search_engine_id = EnvService.get_google_search_engine_id() self.loop = asyncio.get_running_loop() @@ -55,7 +57,8 @@ class Search: else: return "An error occurred while searching." - async def search(self, query, user_api_key, search_scope): + async def search(self, query, user_api_key, search_scope, nodes): + DEFAULT_SEARCH_NODES = 2 if not user_api_key: os.environ["OPENAI_API_KEY"] = self.openai_key else: @@ -86,9 +89,14 @@ class Search: except Exception as e: traceback.print_exc() + prompthelper = PromptHelper(4096, 1024, 20) + index = GPTSimpleVectorIndex(documents) + llm_predictor = LLMPredictor(llm=OpenAI(model_name="text-davinci-003")) + embedding_model = OpenAIEmbedding() # Now we can search the index for a query: - response = index.query(query, text_qa_template=self.qaprompt) + response = index.query(query,embed_model=embedding_model,llm_predictor=llm_predictor,prompt_helper=prompthelper, similarity_top_k=nodes or DEFAULT_SEARCH_NODES, text_qa_template=self.qaprompt) + await self.usage_service.update_usage(llm_predictor.last_token_usage, embedding_model.last_token_usage) return response