I want to add metadata with llamaindex

[bhavya-giri]

from phoenix.otel import register
from openinference.instrumentation.llama_index import LlamaIndexInstrumentor

tracer_provider = register(
project_name="default", endpoint="url"
)

LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)

need help adding metadata

[dosubot[bot]]

To add metadata when using LlamaIndexInstrumentor with OpenInference, you can use the using_metadata context manager to include metadata in your OpenTelemetry context. Here's an example of how you can do this:

from openinference.instrumentation import using_metadata
from openinference.instrumentation.llama_index import LlamaIndexInstrumentor
from phoenix.otel import register

# Register the tracer provider
tracer_provider = register(
    project_name="default", endpoint="url"
)

# Instrument LlamaIndex
LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)

# Define your metadata
metadata = {
    "key-1": "value_1",
    "key-2": "value_2",
    # Add more key-value pairs as needed
}

# Use the context manager to add metadata
with using_metadata(metadata):
    # Your code here will have the metadata added to the spans
    # For example, querying or other operations
    ...

This will ensure that the specified metadata is included in the spans generated within the context block ^[1].

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Other}

[bhavya-giri]

@dosu, how do i do it here?

import timeit
from loguru import logger
import os
import openai

from core.ai_config import AIConfig
from core.config import config

from core.llms.groq import GroqFactory

from core.llms.openai import OpenAIFactory
from core.llms.anthropic import AnthropicFactory

from core.llms.gemini import GeminiFactory

from .TrackPerformance import log_execution_time, log_performance
from .utils import InferenceUtils
from llama_index.core import VectorStoreIndex
from llama_index.core.query_engine import CitationQueryEngine
from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.core.query_engine import RetrieverQueryEngine

from llama_index.core.postprocessor import LLMRerank

from llama_index.postprocessor.colbert_rerank import ColbertRerank

from llama_index.core.postprocessor import SentenceTransformerRerank

import nest_asyncio
import cProfile
import pstats
import io
import concurrent.futures
from schemas.Inference import ResponseQueryRetriver
from phoenix.otel import register
from openinference.instrumentation.llama_index import LlamaIndexInstrumentor
from openinference.instrumentation import using_metadata

tracer_provider = register(
project_name="default", endpoint="http://10.10.11.6:6006/v1/traces"
)

LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)

nest_asyncio.apply()

os.environ["TOKENIZERS_PARALLELISM"] = "false"

utils = InferenceUtils()

ai_config = AIConfig()
AI_MODEL = {
"gpt-4o-mini": "gpt-4o-mini",
"gpt-4o": "gpt-4o",
"claude-sonnet": "claude-3-5-sonnet-20240620",
}

rerank = SentenceTransformerRerank(model=config.RERANK_MODEL, top_n=config.RERANK_TOP_N)

@log_execution_time
def construct_index(project_id, aiModel, log_prefix: str = None):

ai_configuration = ai_config.load_config_by_matter_id(project_id)
embed_model = InferenceUtils.set_embedding_model_config(ai_configuration)

milvus_collection_name = f"matter_{project_id}"
logger.debug(f"{log_prefix} Milvus collection name: {milvus_collection_name}")

logger.debug("{log_prefix} Retrieving vector store...")
vector_store = InferenceUtils.get_milvus_store(
    milvus_collection_name, ai_configuration
)
logger.debug(f"{log_prefix} Vector store retrieved!")

logger.debug("{log_prefix} Creating storage context...")
storage_context = InferenceUtils.get_storage_context(
    vector_store, collection_name=milvus_collection_name
)
logger.debug(f"{log_prefix} Storage context created!")

logger.debug("{log_prefix} Creating VectorStoreIndex...")

index = VectorStoreIndex.from_vector_store(
    vector_store=vector_store, embed_model=embed_model
)

index.from_documents(
    documents=[],
    storage_context=storage_context,
    show_progress=True,
)

return index

def get_llm(aiModel, api_key=None):

# if aiModel == "mistral":
#     llm_factory = GroqFactory(
#         model="mixtral-8x7b-32768",
#         api_key="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",
#     )
if aiModel == "claude-3-5-sonnet-20240620":
    llm_factory = AnthropicFactory(model=aiModel, api_key=api_key)
# elif aiModel == "gemini":
#     llm_factory = GeminiFactory(
#         model="models/gemini-1.5-pro-latest",
#     )
else:
    llm_factory = OpenAIFactory(model=aiModel, api_key=api_key)

llm = llm_factory.get_llm()
return llm

def construct_engine(
retriver,
index,
llm_uuid,
aiModel,
similarity_top_k,
log_prefix: str = None,
project_id: str = None,
):

api_key = ai_config.get_api_key_llm_id(llm_uuid)

llm = get_llm(aiModel, api_key)
colbert_reranker = ColbertRerank(
    top_n=10,
    model="colbert-ir/colbertv2.0",
    tokenizer="colbert-ir/colbertv2.0",
    keep_retrieval_score=True,
)
if retriver == "topk":
    logger.debug("{log_prefix} Creating CitationQueryEngine for topK mode...")
    query_engine = CitationQueryEngine.from_args(
        index,
        node_postprocessors=[colbert_reranker],
        similarity_top_k=similarity_top_k,
        llm=llm,
    )
    logger.debug("{log_prefix} CitationQueryEngine created!")
    return query_engine
elif retriver == "tree_summarize":
    logger.debug(
        "{log_prefix} Creating CitationQueryEngine for tree_summarize mode..."
    )
    query_engine_tree = CitationQueryEngine.from_args(
        index,
        response_mode="tree_summarize",
        citation_chunk_size=config.CHUNK_SIZE,
        similarity_top_k=similarity_top_k,
        llm=llm,
    )
    logger.debug("{log_prefix} CitationQueryEngine created!")
    return query_engine_tree
elif retriver == "fusion":
    logger.debug("{log_prefix} Creating retriever for fusion mode...")
    query_engine_tree = CitationQueryEngine.from_args(
        index,
        response_mode="tree_summarize",
        citation_chunk_size=config.CHUNK_SIZE,
        similarity_top_k=similarity_top_k,
        llm=llm,
    )
    vector_retriever = index.as_retriever(similarity_top_k=config.TOP_K)
    retriever = QueryFusionRetriever(
        [query_engine_tree, vector_retriever],
        similarity_top_k=similarity_top_k,
        num_queries=10,
        use_async=True,
        verbose=True,
        mode="reciprocal_rerank",
        llm=llm,
        # query_gen_prompt="...",  # we could override the query generation prompt here
    )

    logger.debug("{log_prefix} Querying with fusion mode...")

    fusion_retriever = RetrieverQueryEngine(retriever)
    return fusion_retriever

class InferenceRetriver:
"""
This class is used to do inference on the index.
"""

def __init__(self):
    logger.debug("Initializing Inference")

def __execute_query(self, query_engine, user_input):
    return query_engine.query(f"{user_input}")

def query_index(
    self,
    user_input: str,
    project_id: str,
    llm_uuid: str,
    aiModel: str,
    retriver: str,
    uuid: str,
    sid: str,
    matter_context: str = "",
    similarity_top_k: int = 10,
):
    # Query the data
    # storage_context = StorageContext.from_defaults(persist_dir="./storage")
    # index = load_index_from_storage(storage_context)
    metadata = {
        "project_id": project_id,
        "llm_uuid": llm_uuid,
        "model": aiModel,
        "retriever_type": retriver,
        "request_id": uuid,
        "session_id": sid
    }
    log_prefix = f"sid: {sid} ::: UUID: {uuid} :::"

    try:
        aiModel = AI_MODEL[aiModel]

        if matter_context:
            user_input = f"{user_input} {matter_context}"

        logger.debug(f"{log_prefix}")
        logger.debug(
            f"""{log_prefix} AIModel: {aiModel} @@ TopK: {similarity_top_k} @@ Project Id: {project_id}"""
        )

        index = construct_index(project_id, aiModel, log_prefix=log_prefix)
        logger.debug("Index created")

        query_engine = construct_engine(
            retriver,
            index,
            llm_uuid,
            aiModel,
            similarity_top_k,
            log_prefix=log_prefix,
            project_id=project_id,
        )
        
        using_metadata(metadata)
        def __execute_tree_summarize():
            logger.debug(f"{log_prefix} Querying with tree_summarize mode...")
            start_time = timeit.default_timer()
            tree_summary_response = self.__execute_query(
                query_engine=query_engine, user_input=user_input
            )
            end_time = timeit.default_timer()
            log_performance(
                f"{log_prefix} Execution time of :: Tree_summarize :: ",
                end_time - start_time,
            )
            return tree_summary_response
        
        using_metadata(metadata)
        def __execute_top_k():
            logger.debug(f"{log_prefix} Querying with topK mode...")
            start_time = timeit.default_timer()
            response_topk = self.__execute_query(
                query_engine=query_engine, user_input=user_input
            )
            end_time = timeit.default_timer()
            log_performance(
                f"{log_prefix} Execution time of :: TopK :: ", end_time - start_time
            )
            return response_topk
        
        using_metadata(metadata)
        def __execute_fusion():
            logger.debug(f"{log_prefix} Response retriever for fusion mode...")
            start_time = timeit.default_timer()
            # formatted_user_input = format_fusion_prompt(user_input)
            fusion_response = self.__execute_query(
                query_engine=query_engine, user_input=f"Summarize {user_input}"
            )
            end_time = timeit.default_timer()
            log_performance(
                f"{log_prefix} Execution time of :: Fusion Query :: ",
                end_time - start_time,
            )
            return fusion_response

        def __execute_gemini():
            logger.debug("{log_prefix} Querying with gemini topK mode...")
            start_time = timeit.default_timer()
            response_gemini = self.__execute_query(
                query_engine=query_engine, user_input=user_input
            )
            end_time = timeit.default_timer()
            log_performance(
                f"{log_prefix} Execution time of :: Gemini TopK :: ",
                end_time - start_time,
            )
            return response_gemini

        RETRIVER = {
            "topk": {"display_name": "AI Agent 1", "function": __execute_top_k},
            "tree_summarize": {
                "display_name": "AI Agent 2",
                "function": __execute_tree_summarize,
            },
            "fusion": {"display_name": "AI Agent 3", "function": __execute_fusion},
            "gemini": {"display_name": "AI Agent 4", "function": __execute_gemini},
        }

        try:
            response = RETRIVER[retriver]["function"]()

        except Exception as e:
            logger.error(
                f"{log_prefix} Question: {user_input} ::: An error occurred: {e}"
            )

        return {retriver: response}
    except Exception as e:
        logger.error(f"Error in query_index UUID: {uuid} token: {sid} - {e}")

        return None

[dosubot[bot]]

To add metadata using the using_metadata context manager in your query_index method, you should wrap the relevant code blocks where you want the metadata to be applied. Here's how you can modify your code:

def query_index(
    self,
    user_input: str,
    project_id: str,
    llm_uuid: str,
    aiModel: str,
    retriver: str,
    uuid: str,
    sid: str,
    matter_context: str = "",
    similarity_top_k: int = 10,
):
    metadata = {
        "project_id": project_id,
        "llm_uuid": llm_uuid,
        "model": aiModel,
        "retriever_type": retriver,
        "request_id": uuid,
        "session_id": sid
    }
    log_prefix = f"sid: {sid} ::: UUID: {uuid} :::"

    try:
        aiModel = AI_MODEL[aiModel]

        if matter_context:
            user_input = f"{user_input} {matter_context}"

        logger.debug(f"{log_prefix}")
        logger.debug(
            f"""{log_prefix} AIModel: {aiModel} @@ TopK: {similarity_top_k} @@ Project Id: {project_id}"""
        )

        index = construct_index(project_id, aiModel, log_prefix=log_prefix)
        logger.debug("Index created")

        query_engine = construct_engine(
            retriver,
            index,
            llm_uuid,
            aiModel,
            similarity_top_k,
            log_prefix=log_prefix,
            project_id=project_id,
        )

        with using_metadata(metadata):
            def __execute_tree_summarize():
                logger.debug(f"{log_prefix} Querying with tree_summarize mode...")
                start_time = timeit.default_timer()
                tree_summary_response = self.__execute_query(
                    query_engine=query_engine, user_input=user_input
                )
                end_time = timeit.default_timer()
                log_performance(
                    f"{log_prefix} Execution time of :: Tree_summarize :: ",
                    end_time - start_time,
                )
                return tree_summary_response

            def __execute_top_k():
                logger.debug(f"{log_prefix} Querying with topK mode...")
                start_time = timeit.default_timer()
                response_topk = self.__execute_query(
                    query_engine=query_engine, user_input=user_input
                )
                end_time = timeit.default_timer()
                log_performance(
                    f"{log_prefix} Execution time of :: TopK :: ", end_time - start_time
                )
                return response_topk

            def __execute_fusion():
                logger.debug(f"{log_prefix} Response retriever for fusion mode...")
                start_time = timeit.default_timer()
                fusion_response = self.__execute_query(
                    query_engine=query_engine, user_input=f"Summarize {user_input}"
                )
                end_time = timeit.default_timer()
                log_performance(
                    f"{log_prefix} Execution time of :: Fusion Query :: ",
                    end_time - start_time,
                )
                return fusion_response

            def __execute_gemini():
                logger.debug("{log_prefix} Querying with gemini topK mode...")
                start_time = timeit.default_timer()
                response_gemini = self.__execute_query(
                    query_engine=query_engine, user_input=user_input
                )
                end_time = timeit.default_timer()
                log_performance(
                    f"{log_prefix} Execution time of :: Gemini TopK :: ",
                    end_time - start_time,
                )
                return response_gemini

            RETRIVER = {
                "topk": {"display_name": "AI Agent 1", "function": __execute_top_k},
                "tree_summarize": {
                    "display_name": "AI Agent 2",
                    "function": __execute_tree_summarize,
                },
                "fusion": {"display_name": "AI Agent 3", "function": __execute_fusion},
                "gemini": {"display_name": "AI Agent 4", "function": __execute_gemini},
            }

            try:
                response = RETRIVER[retriver]["function"]()

            except Exception as e:
                logger.error(
                    f"{log_prefix} Question: {user_input} ::: An error occurred: {e}"
                )

            return {retriver: response}
    except Exception as e:
        logger.error(f"Error in query_index UUID: {uuid} token: {sid} - {e}")

        return None

By wrapping the function definitions and their calls within the with using_metadata(metadata): block, the metadata will be included in the spans generated during their execution ^[1].

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Other}

[axiomofjoy]

Hey @bhavya-giri, please check out this portion of the docs here.