Skip to content

Commit 0ca9890

Browse files
IckerdayIckerday
committed
Save progress for now
1 parent fcae942 commit 0ca9890

11 files changed

Lines changed: 212 additions & 132 deletions

File tree

‎pyproject.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@ classifiers = [
3131

3232
dependencies = [
3333
"fastmcp>=2.12.4",
34+
"httpx>=0.28.1",
35+
"mcp>=1.15.0",
3436
"python-dotenv>=0.21.1",
3537
]
3638
optional-dependencies = { compat = ["six>=1.17.0"] }
@@ -75,5 +77,3 @@ select = [
7577
"RUF", # ruff-specific rules
7678
]
7779

78-
[tool.uv.workspace]
79-
members = ["tests/system/test_apps/mcp_enabled_app"]

‎splunklib/mcp/mcp.py‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
import httpx
2+
from mcp.types import Tool as MCPTool
3+
4+
from splunklib.mcp.tools.models import AddTool, AddToolsRequest
5+
6+
7+
async def send_mcp_registrations(
8+
endpoint_url: str,
9+
tool_registrations: list[MCPTool],
10+
server_file_path: str,
11+
):
12+
async with httpx.AsyncClient() as client:
13+
add_req = AddToolsRequest(
14+
tools=[
15+
AddTool(script_path=server_file_path, spec=tool)
16+
for tool in tool_registrations
17+
]
18+
)
19+
20+
res = await client.post(endpoint_url, json=add_req.model_dump())
21+
print(res.status_code)
22+
print(res.text)
23+
24+
25+
async def execute_tool(endpoint_url: str):
26+
async with httpx.AsyncClient() as client:
27+
res = await client.post(endpoint_url)
28+
print(res.text)

‎splunklib/mcp/tools/models.py‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
from dataclasses import field
2+
from typing import Any, Literal
3+
4+
from fastmcp.client import Client
5+
from fastmcp.server.dependencies import get_context
6+
from fastmcp.tools import Tool as FastMCPTool
7+
from fastmcp.tools.tool import ToolResult
8+
from mcp.types import Tool as MCPTool
9+
from pydantic.main import BaseModel
10+
from typing_extensions import override
11+
12+
13+
class SplunkMeta(BaseModel):
14+
permissions: list[str] = field(default=[])
15+
tool_type: str = field(default="")
16+
schema_version: str = field(default="")
17+
execution_mode: str = field(default="")
18+
execution_endpoint: str = field(default="")
19+
20+
21+
class McpInputOutputSchema(BaseModel):
22+
type: Literal["object"] = "object"
23+
properties: dict[str, Any] = field(default_factory=lambda: {}) # pyright: ignore[reportExplicitAny]
24+
required: list[str] = field(default_factory=lambda: [])
25+
26+
27+
class AddTool(BaseModel):
28+
script_path: str
29+
spec: MCPTool
30+
31+
32+
class AddToolsRequest(BaseModel):
33+
tools: list[AddTool]
34+
35+
36+
class DeleteToolsRequest(BaseModel):
37+
tools: list[str]
38+
39+
40+
class ProxiedTool(FastMCPTool):
41+
script: str
42+
43+
@override
44+
async def run(self, arguments: dict[str, Any]) -> ToolResult: # pyright: ignore[reportExplicitAny]
45+
async def progress_handler(
46+
progress: float, total: float | None, message: str | None
47+
) -> None:
48+
await get_context().report_progress(progress, total, message)
49+
50+
c = Client(transport=self.script)
51+
52+
async with c:
53+
res = await c.call_tool(
54+
self.name, arguments, progress_handler=progress_handler
55+
)
56+
57+
# TODO: we are missing some fields ....
58+
# res.is_error
59+
# res.data
60+
return ToolResult(
61+
content=res.content, structured_content=res.structured_content
62+
)

splunklib/mcp/tools.py renamed to splunklib/mcp/tools/registrations.py

Lines changed: 11 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,14 @@
11
import configparser
2-
import os
3-
from dataclasses import asdict, dataclass, field
4-
from typing import Any, Literal
2+
from dataclasses import asdict
3+
from typing import Literal
54

65
from fastmcp.client import Client
7-
from fastmcp.client.transports import PythonStdioTransport
8-
from mcp.types import Tool
9-
10-
11-
@dataclass
12-
class SplunkMeta:
13-
permissions: list[str] = field()
14-
tool_type: str = field(default="")
15-
schema_version: str = field(default="")
16-
execution_mode: str = field(default="")
17-
execution_endpoint: str = field(default="")
18-
19-
20-
@dataclass
21-
class McpInputOutputSchema:
22-
type: Literal["object"] = "object"
23-
properties: dict[str, Any] = field(default_factory=lambda: {}) # pyright: ignore[reportExplicitAny]
24-
required: list[str] = field(default_factory=lambda: [])
6+
from mcp.types import Tool as MCPTool
257

8+
from splunklib.mcp.tools.models import (
9+
McpInputOutputSchema,
10+
SplunkMeta,
11+
)
2612

2713
tool_reg_prefix = "app:mcp_tool"
2814

@@ -48,7 +34,7 @@ def match_input_schema(input: Literal["query_string"] | Literal["other"]):
4834
raise NotImplementedError("We don't know what to put here lol")
4935

5036

51-
def parse_ai_conf(file_path: str) -> list[Tool]:
37+
def parse_ai_conf(file_path: str) -> list[MCPTool]:
5238
config = configparser.ConfigParser()
5339
all_sections_len = config.read(file_path)
5440
if len(all_sections_len) == 0:
@@ -88,23 +74,12 @@ def parse_ai_conf(file_path: str) -> list[Tool]:
8874
return ini_tools
8975

9076

91-
async def get_tools(server_path: str):
77+
async def get_mcp_tools(server_path: str) -> list[MCPTool]:
78+
"""Connects to local MCP server to get tools registered with a @tool decorator"""
9279
mcp_client = Client(server_path)
9380

94-
tools = []
81+
tools: list[MCPTool] = []
9582
async with mcp_client:
9683
tools = await mcp_client.list_tools()
9784

98-
# TODO: Get registrations from ai.conf
99-
# curr_path = os.path.join(os.getcwd(), "..", "default", "app.conf")
100-
# yaml_tool_registrations: list[Tool] = parse_ai_conf(curr_path)
101-
10285
return tools
103-
104-
105-
async def register_tools_from(file_paths: list[str]) -> None:
106-
"""TODO
107-
1. `POST /tools` with MCP payload
108-
2.
109-
"""
110-
print(file_paths)

‎tests/system/test_apps/mcp_enabled_app/README.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,3 +13,7 @@ This example aspires to verify the points listed in [POC - AI with Splunk Apps](
1313
- Run `uv sync`
1414
- `source .venv/bin/activate`
1515
- Run the code blocks `bin/mcp_enabled_app.ipynb`
16+
17+
## TODO
18+
19+
- Research using server composition: <https://gofastmcp.com/servers/composition>

‎tests/system/test_apps/mcp_enabled_app/bin/execute_tool_example.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
import asyncio
2+
3+
4+
def execute_tool_example(endpoint_url: str):
5+
pass
6+
7+
8+
if __name__ == "__main__":
9+
MCP_SERVER_HOST: str = "0.0.0.0"
10+
MCP_SERVER_PORT: int = 2137
11+
12+
asyncio.run(execute_tool_example(""))

‎tests/system/test_apps/mcp_enabled_app/bin/mcp_enabled_app.ipynb‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,14 +13,6 @@
1313
" - [ ] Capture e.g. tool_name, description, inputs, outputs\n"
1414
]
1515
},
16-
{
17-
"cell_type": "code",
18-
"execution_count": null,
19-
"id": "18ab5550",
20-
"metadata": {},
21-
"outputs": [],
22-
"source": []
23-
},
2416
{
2517
"cell_type": "markdown",
2618
"id": "363201e8",

‎tests/system/test_apps/mcp_enabled_app/bin/mcp_tools.py‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
import asyncio
2+
import json
3+
import time
4+
5+
from fastmcp import Context, FastMCP
6+
7+
from splunklib import client
8+
from splunklib.mcp.tools.models import SplunkMeta
9+
from splunklib.results import JSONResultsReader
10+
11+
app_mcp_server = FastMCP("GeneratingCSC PoC Server")
12+
13+
14+
@app_mcp_server.tool(
15+
description="""
16+
The `generatingcsc` command generates a specific number of records.
17+
18+
Example:
19+
``| generatingcsc count=4``
20+
Returns a 4 records having text 'Test Event'.
21+
""",
22+
meta=SplunkMeta(
23+
permissions=["role:search_admin", "role:aws_analyst"],
24+
tool_type="search",
25+
schema_version="1.0",
26+
execution_endpoint="",
27+
execution_mode="",
28+
).model_dump(),
29+
enabled=True,
30+
)
31+
async def generating_csc(count: int, ctx: Context) -> list[str]:
32+
service = client.connect(
33+
scheme="https",
34+
host="localhost",
35+
port="8089",
36+
username="admin",
37+
password="changed!",
38+
autologin=True,
39+
)
40+
stream = service.jobs.oneshot(
41+
f"| generatingcsc count={abs(count)}", output_mode="json"
42+
)
43+
results: JSONResultsReader = JSONResultsReader(stream)
44+
for progress in range(0, 5):
45+
await ctx.report_progress((progress + 1) * 2, 100, "Addition in progress")
46+
time.sleep(0.25)
47+
48+
quuuuuux = [json.dumps(r) for r in list(results)]
49+
print(quuuuuux)
50+
return quuuuuux
51+
52+
53+
if __name__ == "__main__":
54+
MCP_SERVER_HOST: str = "0.0.0.0"
55+
MCP_SERVER_PORT: int = 2137
56+
57+
app_mcp_server.run("stdio", show_banner=False)
58+
# asyncio.run(
59+
# app_mcp_server.run_async(
60+
# show_banner=False,
61+
# # host=MCP_SERVER_HOST,
62+
# # port=MCP_SERVER_PORT
63+
# )
64+
# )

‎tests/system/test_apps/mcp_enabled_app/bin/post_install_example.py‎

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,24 @@
1-
from splunklib.mcp import tools
1+
import asyncio
2+
import os
23

4+
from splunklib.mcp.mcp import send_mcp_registrations
5+
from splunklib.mcp.tools import registrations
36

4-
async def post_install(splunk_url: str, auth_token: str) -> None:
5-
# TODO: Implement
6-
try:
7-
await tools.register_tools_from(["./tools.py", "../local/ai.conf"])
8-
except Exception as e:
9-
print(e)
10-
raise e
7+
8+
async def post_install(server_file_path: str, endpoint_url: str) -> None:
9+
tool_registrations = await registrations.get_mcp_tools(server_file_path)
10+
11+
await send_mcp_registrations(
12+
endpoint_url,
13+
tool_registrations,
14+
server_file_path,
15+
)
16+
17+
18+
if __name__ == "__main__":
19+
asyncio.run(
20+
post_install(
21+
f"{os.getcwd()}/tests/system/test_apps/mcp_enabled_app/bin/mcp_tools.py",
22+
"http://0.0.0.0:8090/tools",
23+
)
24+
)

‎tests/system/test_apps/mcp_enabled_app/bin/tools.py‎

Lines changed: 0 additions & 58 deletions
This file was deleted.

0 commit comments

Comments
 (0)