from llama_index.core import Document, Settings, VectorStoreIndex
from llama_index.core.embeddings import MockEmbedding
from llama_index.core.llms.mock import MockLLM
from llama_index.observability.otel import LlamaIndexOpenTelemetry
from opentelemetry.sdk.trace.export import ConsoleSpanExporter

Settings.embed_model = MockEmbedding(embed_dim=8)
Settings.llm = MockLLM(max_tokens=16)

instrumentor = LlamaIndexOpenTelemetry(
    service_name_or_resource="llamaindex-demo",
    span_exporter=ConsoleSpanExporter(),
    span_processor="simple",
)
instrumentor.start_registering()

index = VectorStoreIndex.from_documents(
    [Document(text="OpenTelemetry exports LlamaIndex retrieval and query spans.")]
)
query_engine = index.as_query_engine(similarity_top_k=1)
response = query_engine.query("What does OpenTelemetry export?")

print("Answer:", response)