def extract_graph()

in hugegraph-llm/src/hugegraph_llm/utils/graph_index_utils.py [0:0]

• 26 lines of code
• 6 McCabe index (conditional complexity)

def extract_graph(input_file, input_text, schema, example_prompt) -> str:
    texts = read_documents(input_file, input_text)
    builder = KgBuilder(LLMs().get_chat_llm(), Embeddings().get_embedding(), get_hg_client())
    if not schema:
        return "ERROR: please input with correct schema/format."

    error_message = parse_schema(schema, builder)
    if error_message:
        return error_message
    builder.chunk_split(texts, "document", "zh").extract_info(example_prompt, "property_graph")

    try:
        context = builder.run()
        if not context["vertices"] and not context["edges"]:
            log.info("Please check the schema.(The schema may not match the Doc)")
            return json.dumps(
                {
                    "vertices": context["vertices"],
                    "edges": context["edges"],
                    "warning": "The schema may not match the Doc"
                },
                ensure_ascii=False,
                indent=2
            )
        return json.dumps({"vertices": context["vertices"], "edges": context["edges"]}, ensure_ascii=False, indent=2)
    except Exception as e:  # pylint: disable=broad-exception-caught
        log.error(e)
        raise gr.Error(str(e))