Skip to content

Commit 1105cff

Browse files
mn-mikkemn-mikke
authored
Merge pull request #10 from h2oai/mn/endpoint-filtering
Add endpoint filtering
2 parents 7a54e6f + 6e0a000 commit 1105cff

File tree

3 files changed

+93
-10
lines changed

3 files changed

+93
-10
lines changed

‎README.md‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,15 @@ make install
3030
## Usage
3131

3232
### Environment Variables
33-
- **API_KEY** - (required) H2OGPTe access key. See [documentation](https://docs.h2o.ai/enterprise-h2ogpte/guide/apis#create-an-api-key) on how to get the key.
33+
- **H2OGPTE_API_KEY** - (required) H2OGPTe access key. See [documentation](https://docs.h2o.ai/enterprise-h2ogpte/guide/apis#create-an-api-key) on how to get the key.
3434
- **H2OGPTE_SERVER_URL** - The url of H2OGPTe server. Default value is [https://h2ogpte.genai.h2o.ai](https://h2ogpte.genai.h2o.ai).
35-
- **ALL_ENDPOINTS_AS_TOOLS** - A boolean flag, specifing whether all REST API endpoints should be represented as MCP tools. If disabled, GET endpoints will be represented as resources. Default value is `true`.
35+
- **H2OGPTE_ALL_ENDPOINTS_AS_TOOLS** - A boolean flag, specifing whether all REST API endpoints should be represented as MCP tools. If disabled, GET endpoints will be represented as resources. Default value is `true`.
36+
- **H2OGPTE_ENDPOINT_SET** - A set of REST API endpoints that will be used for MCP tools or resources. The default value is `all_without_async_ingest`. Possible values:
37+
- `all` - All REST API endpoints on the H2OGPTe server
38+
- `all_without_async_ingest` - All endpoints without asynchronous ingestion endpoints. These endpoints start and returns a job. E.g.: `create_ingest_upload_job`
39+
- `basic` - A mininal set of endpoints for chatting with collections and ingesting new documents.
40+
- `custom` - A set of endpoints defined by the user. If chossen, the `H2OGPTE_CUSTOM_ENDPOINT_SET_FILE` variable must be set.
41+
- **H2OGPTE_CUSTOM_ENDPOINT_SET_FILE** - A path to file with the list of REST API endpoints. Each endpoint name must be an a separate line. The name of the endpoint is the `operationId` attribute in REST API spec file (e.g.: [https://h2ogpte.genai.h2o.ai/api-spec.yaml](https://h2ogpte.genai.h2o.ai/api-spec.yaml))
3642

3743
### Example Configuration
3844
An example MCP server configuration for MCP clients. E.g.: Cursor, Claude Desktop
@@ -43,9 +49,9 @@ An example MCP server configuration for MCP clients. E.g.: Cursor, Claude Deskto
4349
"h2ogpte-mcp-server": {
4450
"command": "h2ogpte-mcp-server",
4551
"env": {
46-
"API_KEY": "sk-...",
52+
"H2OGPTE_API_KEY": "sk-...",
4753
"H2OGPTE_SERVER_URL": "https://h2ogpte.genai.h2o.ai",
48-
"ALL_ENDPOINTS_AS_TOOLS": true
54+
"H2OGPTE_ALL_ENDPOINTS_AS_TOOLS": "true"
4955
}
5056
}
5157
}

‎src/h2ogpte_mcp_server/server.py‎

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,14 @@
22
import yaml
33
from fastmcp import FastMCP
44
from fastmcp.utilities.openapi import OpenAPIParser
5-
from .settings import settings
5+
from .settings import settings, basic_endpoints
66
from .tools import register_custom_tools
7+
from .settings import EndpointSet
8+
from typing import List
79

810
async def start_server():
9-
mux_service_url = settings.h2ogpte_server_url
11+
print(f"Starting H2OGPTe MCP API server with endpoint set '{settings.endpoint_set.value}'.")
12+
mux_service_url = settings.server_url
1013

1114
# Load your OpenAPI spec
1215
client = httpx.AsyncClient(base_url=f"{mux_service_url}")
@@ -29,7 +32,21 @@ async def start_server():
2932
)
3033

3134
await register_custom_tools(mcp)
32-
await remove_create_job_tools(mcp)
35+
36+
if settings.endpoint_set == EndpointSet.ALL_WITHOUT_ASYNC_INGEST:
37+
await remove_create_job_tools(mcp)
38+
elif settings.endpoint_set == EndpointSet.BASIC:
39+
await reduce_tools_and_resources(mcp, basic_endpoints)
40+
elif settings.endpoint_set == EndpointSet.CUSTOM:
41+
if not settings.custom_endpoint_set_file:
42+
raise ValueError("Custom endpoint set file is not set. Please set the CUSTOM_ENDPOINT_SET_FILE environment variable.")
43+
with open(settings.custom_endpoint_set_file, "r") as f:
44+
custom_endpoints = [endpoint.strip() for endpoint in f.readlines()]
45+
print(f"Custom endpoints: {custom_endpoints}")
46+
await reduce_tools_and_resources(mcp, custom_endpoints)
47+
elif settings.endpoint_set == EndpointSet.ALL:
48+
pass
49+
3350

3451
await mcp.run_async()
3552

@@ -44,3 +61,17 @@ async def remove_create_job_tools(mcp: FastMCP):
4461
if tool.startswith("create_") and tool.endswith("_job"):
4562
print(f"Skipping tool {tool}")
4663
mcp.remove_tool(tool)
64+
65+
66+
async def reduce_tools_and_resources(mcp: FastMCP, endpoints: List[str]):
67+
tools = await mcp.get_tools()
68+
for tool in tools.keys():
69+
if tool not in endpoints:
70+
print(f"Skipping tool {tool}")
71+
mcp.remove_tool(tool)
72+
73+
resources = await mcp.get_resources()
74+
for resource in resources.keys():
75+
if resource not in endpoints:
76+
print(f"Skipping resource {resource}")
77+
mcp.remove_resource(resource)

‎src/h2ogpte_mcp_server/settings.py‎

Lines changed: 49 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,57 @@
11
from pydantic import Field
22
from pydantic_settings import BaseSettings
3+
from typing import Optional
4+
from enum import Enum
35

6+
class EndpointSet(Enum):
7+
BASIC = "basic"
8+
ALL = "all"
9+
ALL_WITHOUT_ASYNC_INGEST = "all_without_async_ingest"
10+
CUSTOM = "custom"
411

512
class Settings(BaseSettings):
6-
h2ogpte_server_url: str = Field("https://h2ogpte.genai.h2o.ai")
13+
server_url: str = Field("https://h2ogpte.genai.h2o.ai")
714
api_key: str = Field()
815
all_endpoints_as_tools: bool = Field(True)
16+
endpoint_set: EndpointSet = Field(EndpointSet.ALL_WITHOUT_ASYNC_INGEST, case_sensitive=False)
17+
custom_endpoint_set_file: Optional[str] = Field(None)
918

10-
11-
settings = Settings()
19+
settings = Settings(_env_prefix="H2OGPTE_")
20+
basic_endpoints = [
21+
"create_collection",
22+
"list_collections",
23+
"update_collection",
24+
"update_collection_prompt_template",
25+
"delete_collection_prompt_template",
26+
"reset_collection_prompt_settings",
27+
"get_collection_settings",
28+
"update_collection_settings",
29+
"get_collection_chat_settings",
30+
"update_collection_chat_settings",
31+
"list_documents_for_collection",
32+
"insert_document_into_collection",
33+
"delete_document_from_collection",
34+
"list_chat_sessions_for_collection",
35+
"list_documents",
36+
"update_document",
37+
"create_prompt_template",
38+
"list_prompt_templates",
39+
"get_default_prompt_template",
40+
"update_prompt_template",
41+
"upload_file",
42+
"ingest_upload",
43+
"ingest_from_plain_text",
44+
"ingest_from_website",
45+
"ingest_from_s3",
46+
"ingest_from_gcs",
47+
"ingest_from_azure_blob_storage",
48+
"create_chat_session",
49+
"list_chat_sessions",
50+
"update_chat_session",
51+
"update_chat_session_prompt_template",
52+
"delete_chat_session_prompt_template",
53+
"get_completion",
54+
"list_models",
55+
"list_embedding_models",
56+
"get_default_embedding_model",
57+
]

0 commit comments

Comments
 (0)