Skip to content

Commit 4fc0cdd

Browse files
committed
Replace fastapi_mcp with the MCP Python SDK (FastMCP) to fix RecursionError on recursive Pydantic models, expose explicit @mcp.tool() handlers that call FeatureStore directly, and add registry discovery, data access, and materialization tools with flat LLM-friendly schemas.
Signed-off-by: Chaitany patel <patelchaitany93@gmail.com> Rerun-CI: true
1 parent 7ebcf03 commit 4fc0cdd

File tree

9 files changed

+983
-291
lines changed

9 files changed

+983
-291
lines changed

examples/agent_feature_store/agent.py

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -167,15 +167,12 @@ async def _discover_feast_tools() -> dict[str, str]:
167167

168168

169169
def _parse_online_features(data: dict) -> dict[str, Any]:
170-
"""Parse the response from get-online-features into a flat dict."""
171-
results = data.get("results", [])
172-
feature_names = data.get("metadata", {}).get("feature_names", [])
173-
features: dict[str, Any] = {}
174-
for i, name in enumerate(feature_names):
175-
values = results[i].get("values", [])
176-
val = values[0] if values else None
170+
"""Parse the MCP to_dict() response into a flat dict of first-row values."""
171+
features = {}
172+
for key, values in data.items():
173+
val = values[0] if isinstance(values, list) and values else values
177174
if val is not None:
178-
features[name] = val
175+
features[key] = val
179176
return features
180177

181178

@@ -214,15 +211,14 @@ async def tool_search_knowledge_base(query: str) -> list[dict[str, Any]]:
214211
"api_version": 2,
215212
},
216213
)
217-
results = data.get("results", [])
218-
feature_names = data.get("metadata", {}).get("feature_names", [])
219-
220-
num_docs = len(results[0]["values"]) if results else 0
214+
first_col = next(iter(data.values()), [])
215+
num_docs = len(first_col) if isinstance(first_col, list) else 0
221216
docs = []
222217
for doc_idx in range(num_docs):
223218
doc = {}
224-
for feat_idx, name in enumerate(feature_names):
225-
doc[name] = results[feat_idx]["values"][doc_idx]
219+
for key, values in data.items():
220+
if isinstance(values, list) and doc_idx < len(values):
221+
doc[key] = values[doc_idx]
226222
if doc.get("title"):
227223
docs.append(doc)
228224
return docs

sdk/python/feast/feature_server.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -353,8 +353,14 @@ async def lifespan(app: FastAPI):
353353

354354
await store.initialize()
355355
async_refresh()
356+
357+
mcp_sm = getattr(app.state, "mcp_session_manager", None)
356358
try:
357-
yield
359+
if mcp_sm:
360+
async with mcp_sm.run():
361+
yield
362+
else:
363+
yield
358364
finally:
359365
stop_refresh()
360366
if offline_batcher is not None:
Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
# MCP (Model Context Protocol) server implementations for Feast
22

33
from .mcp_config import McpFeatureServerConfig
4-
from .mcp_server import add_mcp_support_to_app
4+
from .mcp_server import add_mcp_support_to_app, create_mcp_server
55

66
__all__ = [
77
"McpFeatureServerConfig",
88
"add_mcp_support_to_app",
9+
"create_mcp_server",
910
]

sdk/python/feast/infra/mcp_servers/mcp_config.py

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,19 +8,16 @@
88
class McpFeatureServerConfig(BaseFeatureServerConfig):
99
"""MCP (Model Context Protocol) Feature Server configuration."""
1010

11-
# Feature server type selector
1211
type: Literal["mcp"] = "mcp"
1312

14-
# Enable MCP server support - defaults to False as requested
1513
mcp_enabled: StrictBool = False
1614

17-
# MCP server name for identification
1815
mcp_server_name: StrictStr = "feast-mcp-server"
1916

20-
# MCP server version
2117
mcp_server_version: StrictStr = "1.0.0"
2218

2319
mcp_transport: Literal["sse", "http"] = "sse"
2420

25-
# The endpoint definition for transformation_service (inherited from base)
21+
mcp_base_path: StrictStr = "/mcp"
22+
2623
transformation_service_endpoint: StrictStr = "localhost:6566"

0 commit comments

Comments
 (0)