llama_stack_mcp_sse.py

from llama_stack_client import Agent, AgentEventLogger, RAGDocument, LlamaStackClient
from llama_stack.apis.common.content_types import URL

client = LlamaStackClient(base_url="http://localhost:8321")

models = client.models.list()

# Select the first LLM and first embedding models
model_id = next(m for m in models if m.model_type == "llm").identifier
embedding_model_id = (
    em := next(m for m in models if m.model_type == "embedding")
).identifier
embedding_dimension = em.metadata["embedding_dimension"]

client.toolgroups.register(
        toolgroup_id="mcp::test",
        provider_id="model-context-protocol",
        mcp_endpoint=URL(uri="http://localhost:8080/sse")
)

agent = Agent(
    client,
    model=model_id,
    instructions="You are a helpful assistant",
    tools=[
        "mcp::test",
    ],
)

prompt = input("prompt>")

response = agent.create_turn(
    messages=[{"role": "user", "content": prompt}],
    session_id=agent.create_session("mcp-demo"),
    stream=True,
)

for log in AgentEventLogger().log(response):
    log.print()