def _gremlin_generate_query()

in hugegraph-llm/src/hugegraph_llm/operators/hugegraph_op/graph_rag_query.py [0:0]

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

    def _gremlin_generate_query(self, context: Dict[str, Any]) -> Dict[str, Any]:
        query = context["query"]
        vertices = context.get("match_vids")
        query_embedding = context.get("query_embedding")

        self._gremlin_generator.clear()
        self._gremlin_generator.example_index_query(num_examples=self._num_gremlin_generate_example)
        gremlin_response = self._gremlin_generator.gremlin_generate_synthesize(
            context["simple_schema"], vertices=vertices, gremlin_prompt=self._gremlin_prompt
        ).run(query=query, query_embedding=query_embedding)
        if self._num_gremlin_generate_example > 0:
            gremlin = gremlin_response["result"]
        else:
            gremlin = gremlin_response["raw_result"]
        log.info("Generated gremlin: %s", gremlin)
        context["gremlin"] = gremlin
        try:
            result = self._client.gremlin().exec(gremlin=gremlin)["data"]
            if result == [None]:
                result = []
            context["graph_result"] = [json.dumps(item, ensure_ascii=False) for item in result]
            if context["graph_result"]:
                context["graph_result_flag"] = 1
                context["graph_context_head"] = (
                    f"The following are graph query result " f"from gremlin query `{gremlin}`.\n"
                )
        except Exception as e:  # pylint: disable=broad-except
            log.error(e)
            context["graph_result"] = ""
        return context