JSONalayze

JSONalyzeQueryEngine

Bases: BaseQueryEngine

JSON List Shape Data Analysis Query Engine.

Converts natural language statasical queries to SQL within in-mem SQLite queries.

list_of_dict(List[Dict[str, Any]]): List of dictionaries to query. jsonalyze_prompt (BasePromptTemplate): The JSONalyze prompt to use. use_async (bool): Whether to use async. analyzer (Callable): The analyzer that executes the query. sql_parser (BaseSQLParser): The SQL parser that ensures valid SQL being parsed from llm output. synthesize_response (bool): Whether to synthesize a response. response_synthesis_prompt (BasePromptTemplate): The response synthesis prompt to use. table_name (str): The table name to use. verbose (bool): Whether to print verbose output.

Source code in llama-index-core/llama_index/core/query_engine/jsonalyze_query_engine.py

class JSONalyzeQueryEngine(BaseQueryEngine):
    """JSON List Shape Data Analysis Query Engine.

    Converts natural language statasical queries to SQL within in-mem SQLite queries.

    list_of_dict(List[Dict[str, Any]]): List of dictionaries to query.
    jsonalyze_prompt (BasePromptTemplate): The JSONalyze prompt to use.
    use_async (bool): Whether to use async.
    analyzer (Callable): The analyzer that executes the query.
    sql_parser (BaseSQLParser): The SQL parser that ensures valid SQL being parsed
        from llm output.
    synthesize_response (bool): Whether to synthesize a response.
    response_synthesis_prompt (BasePromptTemplate): The response synthesis prompt
        to use.
    table_name (str): The table name to use.
    verbose (bool): Whether to print verbose output.
    """

    def __init__(
        self,
        list_of_dict: List[Dict[str, Any]],
        llm: Optional[LLM] = None,
        jsonalyze_prompt: Optional[BasePromptTemplate] = None,
        use_async: bool = False,
        analyzer: Optional[Callable] = None,
        sql_parser: Optional[BaseSQLParser] = None,
        synthesize_response: bool = True,
        response_synthesis_prompt: Optional[BasePromptTemplate] = None,
        table_name: str = DEFAULT_TABLE_NAME,
        verbose: bool = False,
        **kwargs: Any,
    ) -> None:
        """Initialize params."""
        self._list_of_dict = list_of_dict
        self._llm = llm or Settings.llm
        self._jsonalyze_prompt = jsonalyze_prompt or DEFAULT_JSONALYZE_PROMPT
        self._use_async = use_async
        self._analyzer = load_jsonalyzer(use_async, analyzer)
        self._sql_parser = sql_parser or DefaultSQLParser()
        self._synthesize_response = synthesize_response
        self._response_synthesis_prompt = (
            response_synthesis_prompt or DEFAULT_RESPONSE_SYNTHESIS_PROMPT
        )
        self._table_name = table_name
        self._verbose = verbose

        super().__init__(callback_manager=Settings.callback_manager)

    def _get_prompts(self) -> Dict[str, Any]:
        """Get prompts."""
        return {
            "jsonalyze_prompt": self._jsonalyze_prompt,
            "response_synthesis_prompt": self._response_synthesis_prompt,
        }

    def _update_prompts(self, prompts: PromptDictType) -> None:
        """Update prompts."""
        if "jsonalyze_prompt" in prompts:
            self._jsonalyze_prompt = prompts["jsonalyze_prompt"]
        if "response_synthesis_prompt" in prompts:
            self._response_synthesis_prompt = prompts["response_synthesis_prompt"]

    def _get_prompt_modules(self) -> PromptMixinType:
        """Get prompt sub-modules."""
        return {}

    def _query(self, query_bundle: QueryBundle) -> Response:
        """Answer an analytical query on the JSON List."""
        query = query_bundle.query_str
        if self._verbose:
            print_text(f"Query: {query}\n", color="green")

        # Perform the analysis
        sql_query, table_schema, results = self._analyzer(
            self._list_of_dict,
            query_bundle,
            self._llm,
            table_name=self._table_name,
            prompt=self._jsonalyze_prompt,
            sql_parser=self._sql_parser,
        )
        if self._verbose:
            print_text(f"SQL Query: {sql_query}\n", color="blue")
            print_text(f"Table Schema: {table_schema}\n", color="cyan")
            print_text(f"SQL Response: {results}\n", color="yellow")

        if self._synthesize_response:
            response_str = self._llm.predict(
                self._response_synthesis_prompt,
                sql_query=sql_query,
                table_schema=table_schema,
                sql_response=results,
                query_str=query_bundle.query_str,
            )
            if self._verbose:
                print_text(f"Response: {response_str}", color="magenta")
        else:
            response_str = str(results)
        response_metadata = {"sql_query": sql_query, "table_schema": str(table_schema)}

        return Response(response=response_str, metadata=response_metadata)

    async def _aquery(self, query_bundle: QueryBundle) -> Response:
        """Answer an analytical query on the JSON List."""
        query = query_bundle.query_str
        if self._verbose:
            print_text(f"Query: {query}", color="green")

        # Perform the analysis
        sql_query, table_schema, results = self._analyzer(
            self._list_of_dict,
            query,
            self._llm,
            table_name=self._table_name,
            prompt=self._jsonalyze_prompt,
        )
        if self._verbose:
            print_text(f"SQL Query: {sql_query}\n", color="blue")
            print_text(f"Table Schema: {table_schema}\n", color="cyan")
            print_text(f"SQL Response: {results}\n", color="yellow")

        if self._synthesize_response:
            response_str = await self._llm.apredict(
                self._response_synthesis_prompt,
                sql_query=sql_query,
                table_schema=table_schema,
                sql_response=results,
                query_str=query_bundle.query_str,
            )
            if self._verbose:
                print_text(f"Response: {response_str}", color="magenta")
        else:
            response_str = json.dumps(
                {
                    "sql_query": sql_query,
                    "table_schema": table_schema,
                    "sql_response": results,
                }
            )
        response_metadata = {"sql_query": sql_query, "table_schema": str(table_schema)}

        return Response(response=response_str, metadata=response_metadata)