kubeflow / kubeflow/mcp-server

feat: local AI agent with pluggable provider architecture and token-efficient tool discovery

Open
#15 2 comments 0 reactions 1 assignee Claimed by @abhijeet-dhumal View on GitHub
area/agents area/core enhancement
Dominant language
Python
Stars
44
Forks
54
Avg merge
2d 18h
Merged PRs (30d)
29

Description

### Description

## Summary

Add a local AI agent layer for kubeflow-mcp that drives Kubeflow training workflows using Ollama-hosted (or other) LLMs, with a pluggable provider architecture and token-efficient tool modes to avoid context overflow.

## Motivation

- kubeflow-mcp serve is the right surface for IDE clients and production orchestrators — but for local development, iterating over a full MCP client/server round-trip adds friction: process lifecycle, transport setup, auth config, and latency on every tool call.
- A local agent plane solves this differently: same trainer callables, same tool discovery logic, but called in-process through a thin LLM loop. The result is a faster inner loop for developers who want to test training workflows conversationally without standing up a server.
- The token problem is real regardless of transport: loading all tool schemas upfront consumes 5K+ tokens, exhausting context on smaller models. Progressive and semantic modes bound this to under 100 tokens until the agent actually needs a specific tool's schema.
- The pluggable provider architecture means the community isn't locked into one LLM stack — Ollama today, LiteLLM gateway tomorrow, without touching the CLI or tool wiring.

### Use Case

- A developer is iterating on a new training configuration locally. They have Ollama running with qwen3:8b and a Kind or remote k8s cluster with the Kubeflow Trainer operator v2.2.0 installed. They don't want to configure an MCP client in their IDE just to test a workflow - they want a conversational loop they can drive from the terminal.
- The same developer later switches to --mode progressive when testing with a smaller model that has a tighter context window — the tool discovery overhead drops from ~200 to ~85 tokens without changing any code or workflow.
- A platform team member uses --provider litellm --model gpt-4o-mini to route through their existing LiteLLM gateway for audit/cost tracking (tool parity pending).

### Proposed Solution

The agent runs as a separate execution plane - it does not route through kubeflow-mcp serve. Instead it calls trainer callables directly in-process via the shared core.dynamic_tools registry. This is intentional for local/dev use: lower latency, no auth overhead, no MCP hop. The serve plane remains the vision-primary path for IDE clients and future orchestrators.

```
kubeflow-mcp agent -> AgentProvider -> LlamaIndex FunctionAgent -> trainer callables
kubeflow-mcp serve -> FastMCP -> _audit_wrap -> trainer callables
```
Both planes can share core/dynamic_tools.py for progressive/semantic meta-tools.

## Acceptance Criteria

- [ ] Agent connects to a running Ollama instance and executes Kubeflow training workflows
- [ ] Progressive mode reduces token usage vs full-schema loading
- [ ] Semantic mode ranks tools by embedding similarity
- [ ] CLI agent subcommand works end-to-end with qwen3:8b (or any Ollama model)
- [ ] LiteLLM placeholder for future iteration for prod deployment schema

### Component

TrainerClient

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.