(
query,
knowledge_graph_inst: BaseGraphStorage,
entities_vdb: BaseVectorStorage,
community_reports: BaseKVStorage[CommunitySchema],
text_chunks_db: BaseKVStorage[TextChunkSchema],
query_param: QueryParam,
tokenizer_wrapper,
)
| 842 | |
| 843 | |
| 844 | async def _build_local_query_context( |
| 845 | query, |
| 846 | knowledge_graph_inst: BaseGraphStorage, |
| 847 | entities_vdb: BaseVectorStorage, |
| 848 | community_reports: BaseKVStorage[CommunitySchema], |
| 849 | text_chunks_db: BaseKVStorage[TextChunkSchema], |
| 850 | query_param: QueryParam, |
| 851 | tokenizer_wrapper, |
| 852 | ): |
| 853 | results = await entities_vdb.query(query, top_k=query_param.top_k) |
| 854 | if not len(results): |
| 855 | return None |
| 856 | node_datas = await knowledge_graph_inst.get_nodes_batch([r["entity_name"] for r in results]) |
| 857 | if not all([n is not None for n in node_datas]): |
| 858 | logger.warning("Some nodes are missing, maybe the storage is damaged") |
| 859 | node_degrees = await knowledge_graph_inst.node_degrees_batch([r["entity_name"] for r in results]) |
| 860 | node_datas = [ |
| 861 | {**n, "entity_name": k["entity_name"], "rank": d} |
| 862 | for k, n, d in zip(results, node_datas, node_degrees) |
| 863 | if n is not None |
| 864 | ] |
| 865 | use_communities = await _find_most_related_community_from_entities( |
| 866 | node_datas, query_param, community_reports, tokenizer_wrapper |
| 867 | ) |
| 868 | use_text_units = await _find_most_related_text_unit_from_entities( |
| 869 | node_datas, query_param, text_chunks_db, knowledge_graph_inst, tokenizer_wrapper |
| 870 | ) |
| 871 | use_relations = await _find_most_related_edges_from_entities( |
| 872 | node_datas, query_param, knowledge_graph_inst, tokenizer_wrapper |
| 873 | ) |
| 874 | logger.info( |
| 875 | f"Using {len(node_datas)} entites, {len(use_communities)} communities, {len(use_relations)} relations, {len(use_text_units)} text units" |
| 876 | ) |
| 877 | entites_section_list = [["id", "entity", "type", "description", "rank"]] |
| 878 | for i, n in enumerate(node_datas): |
| 879 | entites_section_list.append( |
| 880 | [ |
| 881 | i, |
| 882 | n["entity_name"], |
| 883 | n.get("entity_type", "UNKNOWN"), |
| 884 | n.get("description", "UNKNOWN"), |
| 885 | n["rank"], |
| 886 | ] |
| 887 | ) |
| 888 | entities_context = list_of_list_to_csv(entites_section_list) |
| 889 | |
| 890 | relations_section_list = [ |
| 891 | ["id", "source", "target", "description", "weight", "rank"] |
| 892 | ] |
| 893 | for i, e in enumerate(use_relations): |
| 894 | relations_section_list.append( |
| 895 | [ |
| 896 | i, |
| 897 | e["src_tgt"][0], |
| 898 | e["src_tgt"][1], |
| 899 | e["description"], |
| 900 | e["weight"], |
| 901 | e["rank"], |
no test coverage detected