import asyncio

from llama_index.core import Document, Settings, VectorStoreIndex
from llama_index.core.agent.workflow import FunctionAgent, ToolCallResult
from llama_index.core.base.llms.types import ChatMessage, MessageRole, ToolCallBlock
from llama_index.core.embeddings import MockEmbedding
from llama_index.core.llms import MockLLM
from llama_index.core.llms.mock import MockFunctionCallingLLM
from llama_index.core.tools import QueryEngineTool


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

documents = [
    Document(
        text=(
            "Orion Search release notes: the platform team owns the July rollout. "
            "The maintenance window starts at 09:00 UTC."
        )
    )
]
index = VectorStoreIndex.from_documents(documents)
query_engine = index.as_query_engine(similarity_top_k=1, response_mode="context_only")