From 0de97b2b302a92706ccd1283a97ece901a9fc546 Mon Sep 17 00:00:00 2001 From: Dayenne Souza Date: Wed, 13 Nov 2024 13:10:41 +0000 Subject: [PATCH] Remove dummy nodes from qtd rewrite query --- .../query_text_data/query_rewriter.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/intelligence_toolkit/query_text_data/query_rewriter.py b/intelligence_toolkit/query_text_data/query_rewriter.py index 70c6010f..1be0e49f 100644 --- a/intelligence_toolkit/query_text_data/query_rewriter.py +++ b/intelligence_toolkit/query_text_data/query_rewriter.py @@ -2,6 +2,7 @@ # Licensed under the MIT license. See LICENSE file in the project. import networkx as nx + import intelligence_toolkit.AI.utils as utils import intelligence_toolkit.query_text_data.prompts as prompts @@ -20,12 +21,13 @@ async def rewrite_query(ai_configuration, query, concept_graph, top_concepts): Returns: str: The rewritten query. """ - concepts = sorted(concept_graph.degree(), key=lambda x: x[1], reverse=True)[ - :top_concepts - ] + concepts = sorted(concept_graph.degree(), key=lambda x: x[1], reverse=True) + if "dummynode" in concepts: + concepts.remove("dummynode") + + concepts = concepts[:top_concepts] concepts_str = ", ".join([concept for concept, _ in concepts]) messages = utils.prepare_messages( prompts.query_anchoring_prompt, {"query": query, "concepts": concepts_str} ) - response = await utils.generate_text_async(ai_configuration, messages) - return response \ No newline at end of file + return await utils.generate_text_async(ai_configuration, messages) \ No newline at end of file