Is that possible to use llama idex 'Pandas Query engine' function in Autogen functions and add to a new agent?

[mraguth]

Here my goal is to create an agent to do 'Pandas query' If I use just llama_index, it works fine. Now the question here is how I convert this llama_index code as a function and assign it to an agent. please help.

My code....

from llama_index.llms.azure_openai import AzureOpenAI
from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
import logging
import sys

logging.basicConfig(
    stream=sys.stdout, level=logging.INFO
)  # logging.DEBUG for more verbose output
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))


api_key = "xxxxxx"
azure_endpoint = "xxxxx
api_version = "xxxx

llm = AzureOpenAI(
    model="gpt-35-turbo-16k",
    deployment_name="gpt-35-turbo-16k",
    api_key=api_key,
    azure_endpoint=azure_endpoint,
    api_version=api_version,
)

# You need to deploy your own embedding model as well as your own chat completion model
embed_model = AzureOpenAIEmbedding(
    model="text-embedding-ada-002",
    deployment_name="gpt-35-turbo-16k",
    api_key=api_key,
    azure_endpoint=azure_endpoint,
    api_version=api_version,
)

from llama_index.core import Settings

Settings.llm = llm
Settings.embed_model = embed_model

import logging
import sys
from IPython.display import Markdown, display

import pandas as pd
from llama_index.core.query_engine import PandasQueryEngine


logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))


##################### Config End #############################################


df = pd.DataFrame(
    {
        "city": ["Toronto", "Tokyo", "Berlin"],
        "population": [2930000, 13960000, 3645000],
    }
)

query_engine = PandasQueryEngine(df=df, verbose=True, synthesize_response=True)

query = "Write 3 bullet points of hightlights based on provided dataset and provide point by point"

response = query_engine.query(query)

display(Markdown(f"<b>{response}</b>"))