Coverage for intelligence_toolkit/query_text_data/query_rewriter.py: 0%

1# Copyright (c) 2024 Microsoft Corporation. All rights reserved. 

2# Licensed under the MIT license. See LICENSE file in the project. 

3 

4import networkx as nx 

5 

6import intelligence_toolkit.AI.utils as utils 

7import intelligence_toolkit.query_text_data.prompts as prompts 

8 

9 

10async def rewrite_query(ai_configuration, query, concept_graph, top_concepts): 

11 """ 

12 Rewrite a user query to better match the concepts contained in the dataset. 

13 

14 The output query should retain all the key phrases from the input query, but may expand on them with additional concepts and phrasing to better match relevant concepts in the dataset. If no concepts are relevant, the output query should be the same as the input query. 

15 

16 Args: 

17 query (str): The user query to rewrite. 

18 concept_graph (nx.Graph): The concept graph representing the dataset. 

19 top_concepts (int): The number of top concepts to consider. 

20 

21 Returns: 

22 str: The rewritten query. 

23 """ 

24 concepts = sorted(concept_graph.degree(), key=lambda x: x[1], reverse=True) 

25 if "dummynode" in concepts: 

26 concepts.remove("dummynode") 

27 

28 concepts = concepts[:top_concepts] 

29 concepts_str = ", ".join([concept for concept, _ in concepts]) 

30 messages = utils.prepare_messages( 

31 prompts.query_anchoring_prompt, {"query": query, "concepts": concepts_str} 

32 ) 

33 return await utils.generate_text_async(ai_configuration, messages, stream=False)