Query AI Indices from LangChain

import asyncio
import os

from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain.agents import create_agent
from langchain_openai import ChatOpenAI

CONTEXT_ENGINE_TOOLS = {
    "platform_context_engine_list_ai_indices",
    "platform_context_engine_describe_ai_index",
    "platform_context_engine_query_ai_indices",
}

SYSTEM_PROMPT = """\
You have access to Elastic Context Engine knowledge through three tools.

To answer a question from that knowledge:
1. Call platform_context_engine_list_ai_indices to see what is available.
2. Call platform_context_engine_describe_ai_index on the index you pick.
3. Call platform_context_engine_query_ai_indices with ES|QL built from that description.

Never guess an index ID, an ES|QL target, or a field name — the describe output gives you
all three. The server applies the space and permissions condition, so never add one to a query.
Answer from the rows you get back and cite the Knowledge Indicator titles.
"""


async def main() -> None:
    kibana = os.environ["KIBANA_URL"].rstrip("/")
    space = os.environ.get("KIBANA_SPACE")
    base = f"{kibana}/s/{space}" if space else kibana

    client = MultiServerMCPClient(
        {
            "kibana": {
                "transport": "streamable_http",
                "url": f"{base}/api/agent_builder/mcp",
                "headers": {"Authorization": f"ApiKey {os.environ['KIBANA_API_KEY']}"},
            }
        }
    )

    all_tools = await client.get_tools()
    tools = [t for t in all_tools if t.name in CONTEXT_ENGINE_TOOLS]

    llm = ChatOpenAI(
        model="anthropic/claude-sonnet-4-6",
        openai_api_key=os.environ["OPENROUTER_API_KEY"],
        openai_api_base="https://openrouter.ai/api/v1",
    )
    agent = create_agent(llm, tools)

    result = await agent.ainvoke(
        {
            "messages": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": "What is our refund policy for annual plans?"},
            ]
        }
    )
    print(result["messages"][-1].content)


if __name__ == "__main__":
    asyncio.run(main())
		

import os

import httpx
from langchain.agents import create_agent
from langchain.tools import tool
from langchain_openai import ChatOpenAI

SYSTEM_PROMPT = """\
You have access to Elastic Context Engine knowledge through three tools.

To answer a question from that knowledge:
1. Call list_ai_indices to see what is available.
2. Call describe_ai_index on the index you pick.
3. Call query_ai_indices with ES|QL built from that description.

Never guess an index ID, an ES|QL target, or a field name — the describe output gives you
all three. The server applies the space and permissions condition, so never add one to a query.
Answer from the rows you get back and cite the Knowledge Indicator titles.
"""

_kibana = os.environ["KIBANA_URL"].rstrip("/")
_space = os.environ.get("KIBANA_SPACE")
_base = f"{_kibana}/s/{_space}" if _space else _kibana

client = httpx.Client(
    base_url=f"{_base}/api/context_engine",
    headers={
        "Authorization": f"ApiKey {os.environ['KIBANA_API_KEY']}",
        "elastic-api-version": "2023-10-31",
        "kbn-xsrf": "true",
        "Content-Type": "application/json",
    },
    timeout=60,
)


@tool
def list_ai_indices() -> list[dict]:
    """List the AI Indices available in this space, with the ES|QL target to query each
    one against. Call this first, before describing or querying anything."""
    response = client.get("/ai_index")
    response.raise_for_status()
    return [
        {
            "id": entry["id"],
            "esql_target": entry["dest"]["value"],
            "description": entry.get("description"),
        }
        for entry in response.json()["ai_indices"]
    ]


@tool
def describe_ai_index(ai_index_id: str) -> str:
    """Describe one AI Index: the ES|QL target to put after FROM, every field it exposes
    and which are semantic, the knowledge indicator types and tags it holds, and example
    queries. Always call this before writing a query against an index."""
    response = client.get(f"/ai_index/{ai_index_id}/_describe")
    response.raise_for_status()
    return response.json()["response"]


@tool
def query_ai_indices(query: str, params: dict | None = None, limit: int = 20) -> dict:
    """Run an ES|QL query against one or more AI Indices and return {columns, values}.
    Build the query from the output of describe_ai_index: use its FROM target and its
    field names verbatim. Pass user input as named parameters in params rather than
    writing it into the query string. Do not add any space, tenant, or permissions
    condition to the query."""
    body = {"query": query, "limit": limit}
    if params:
        body["params"] = params
    response = client.post("/ai_index/_query", json=body)
    response.raise_for_status()
    return response.json()


def main() -> None:
    llm = ChatOpenAI(
        model="anthropic/claude-sonnet-4-6",
        openai_api_key=os.environ["OPENROUTER_API_KEY"],
        openai_api_base="https://openrouter.ai/api/v1",
    )
    agent = create_agent(llm, [list_ai_indices, describe_ai_index, query_ai_indices])

    result = agent.invoke(
        {
            "messages": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": "What is our refund policy for annual plans?"},
            ]
        }
    )
    print(result["messages"][-1].content)


if __name__ == "__main__":
    main()
		

import asyncio
import os
from urllib.request import urlopen

from deepagents.backends import StateBackend
from deepagents.backends.utils import create_file_data
from deepagents.middleware import FilesystemMiddleware, SkillsMiddleware
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain.agents import create_agent
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import InMemorySaver

CONTEXT_ENGINE_TOOLS = {
    "platform_context_engine_list_ai_indices",
    "platform_context_engine_describe_ai_index",
    "platform_context_engine_query_ai_indices",
}

SKILL_URL = (
    "https://raw.githubusercontent.com/elastic/agent-skills"
    "/main/skills/kibana/kibana-context-engine/SKILL.md"
)


async def main() -> None:
    kibana = os.environ["KIBANA_URL"].rstrip("/")
    space = os.environ.get("KIBANA_SPACE")
    base = f"{kibana}/s/{space}" if space else kibana

    client = MultiServerMCPClient(
        {
            "kibana": {
                "transport": "streamable_http",
                "url": f"{base}/api/agent_builder/mcp",
                "headers": {"Authorization": f"ApiKey {os.environ['KIBANA_API_KEY']}"},
            }
        }
    )

    all_tools = await client.get_tools()
    tools = [t for t in all_tools if t.name in CONTEXT_ENGINE_TOOLS]

    with urlopen(SKILL_URL) as response:
        skill = response.read().decode()

    backend = StateBackend()
    skill_files = {
        "/skills/kibana-context-engine/SKILL.md": create_file_data(skill),
    }

    llm = ChatOpenAI(
        model="anthropic/claude-sonnet-4-6",
        openai_api_key=os.environ["OPENROUTER_API_KEY"],
        openai_api_base="https://openrouter.ai/api/v1",
    )
    agent = create_agent(
        llm,
        tools,
        middleware=[
            FilesystemMiddleware(backend=backend),
            SkillsMiddleware(backend=backend, sources=["/skills/"]),
        ],
        checkpointer=InMemorySaver(),
    )

    result = await agent.ainvoke(
        {
            "messages": [
                {"role": "user", "content": "What is our refund policy for annual plans?"},
            ],
            "files": skill_files,
        },
        config={"configurable": {"thread_id": "1"}},
    )
    print(result["messages"][-1].content)


if __name__ == "__main__":
    asyncio.run(main())
		

import os
from urllib.request import urlopen

import httpx
from deepagents.backends import StateBackend
from deepagents.backends.utils import create_file_data
from deepagents.middleware import FilesystemMiddleware, SkillsMiddleware
from langchain.agents import create_agent
from langchain.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import InMemorySaver

SKILL_URL = (
    "https://raw.githubusercontent.com/elastic/agent-skills"
    "/main/skills/kibana/kibana-context-engine/SKILL.md"
)

_kibana = os.environ["KIBANA_URL"].rstrip("/")
_space = os.environ.get("KIBANA_SPACE")
_base = f"{_kibana}/s/{_space}" if _space else _kibana

client = httpx.Client(
    base_url=f"{_base}/api/context_engine",
    headers={
        "Authorization": f"ApiKey {os.environ['KIBANA_API_KEY']}",
        "elastic-api-version": "2023-10-31",
        "kbn-xsrf": "true",
        "Content-Type": "application/json",
    },
    timeout=60,
)


@tool
def list_ai_indices() -> list[dict]:
    """List the AI Indices available in this space, with the ES|QL target to query each
    one against. Call this first, before describing or querying anything."""
    response = client.get("/ai_index")
    response.raise_for_status()
    return [
        {
            "id": entry["id"],
            "esql_target": entry["dest"]["value"],
            "description": entry.get("description"),
        }
        for entry in response.json()["ai_indices"]
    ]


@tool
def describe_ai_index(ai_index_id: str) -> str:
    """Describe one AI Index: the ES|QL target to put after FROM, every field it exposes
    and which are semantic, the knowledge indicator types and tags it holds, and example
    queries. Always call this before writing a query against an index."""
    response = client.get(f"/ai_index/{ai_index_id}/_describe")
    response.raise_for_status()
    return response.json()["response"]


@tool
def query_ai_indices(query: str, params: dict | None = None, limit: int = 20) -> dict:
    """Run an ES|QL query against one or more AI Indices and return {columns, values}.
    Build the query from the output of describe_ai_index: use its FROM target and its
    field names verbatim. Pass user input as named parameters in params rather than
    writing it into the query string. Do not add any space, tenant, or permissions
    condition to the query."""
    body = {"query": query, "limit": limit}
    if params:
        body["params"] = params
    response = client.post("/ai_index/_query", json=body)
    response.raise_for_status()
    return response.json()


def main() -> None:
    with urlopen(SKILL_URL) as response:
        skill = response.read().decode()

    backend = StateBackend()
    skill_files = {
        "/skills/kibana-context-engine/SKILL.md": create_file_data(skill),
    }

    llm = ChatOpenAI(
        model="anthropic/claude-sonnet-4-6",
        openai_api_key=os.environ["OPENROUTER_API_KEY"],
        openai_api_base="https://openrouter.ai/api/v1",
    )
    agent = create_agent(
        llm,
        [list_ai_indices, describe_ai_index, query_ai_indices],
        middleware=[
            FilesystemMiddleware(backend=backend),
            SkillsMiddleware(backend=backend, sources=["/skills/"]),
        ],
        checkpointer=InMemorySaver(),
    )

    result = agent.invoke(
        {
            "messages": [
                {"role": "user", "content": "What is our refund policy for annual plans?"},
            ],
            "files": skill_files,
        },
        config={"configurable": {"thread_id": "1"}},
    )
    print(result["messages"][-1].content)


if __name__ == "__main__":
    main()