Merge pull request #14 from h2oai/mn/custom-openapi-spec

Add option to specify custom endpoint spec file

Author: mn-mikke

README.md

@@ -39,6 +39,7 @@ make install
   - `basic` - A mininal set of endpoints for chatting with collections and ingesting new documents.
   - `custom` - A set of endpoints defined by the user. If chossen, the `H2OGPTE_CUSTOM_ENDPOINT_SET_FILE` variable must be set.
 - **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)) 
+- **H2OGPTE_CUSTOM_ENDPOINT_SPEC_FILE** - A path to OpenAPI spec file in YAML format describing REST API of the H2OGPTe server. If not specified, the file is obtained from the H2OGPTe server itself. This environement variable should be used only for debugging purposes.
 
 ### Example Configuration
 An example MCP server configuration for MCP clients. E.g.: Cursor, Claude Desktop

src/h2ogpte_mcp_server/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.1.4-dev"
+__version__ = "0.1.4"
 
 import asyncio
 from .server import start_server

src/h2ogpte_mcp_server/server.py

@@ -12,10 +12,7 @@ async def start_server():
     mux_service_url = settings.server_url
 
     # Load your OpenAPI spec
-    client = httpx.AsyncClient(base_url=f"{mux_service_url}")
-    response = await client.get("/api-spec.yaml")
-    yaml_spec = response.content
-    openapi_spec = yaml.load(yaml_spec, Loader=yaml.CLoader)
+    openapi_spec = await load_openapi_spec(mux_service_url)
 
     # Create an HTTP client for your API
     headers = {"Authorization": f"Bearer {settings.api_key}"}
@@ -50,6 +47,17 @@ async def start_server():
 
     await mcp.run_async()
 
+async def load_openapi_spec(mux_service_url):
+    if settings.custom_openapi_spec_file:
+        with open(settings.custom_openapi_spec_file, "r") as f:
+            custom_openapi_spec = yaml.load(f, Loader=yaml.CLoader)
+        return custom_openapi_spec
+    else:
+        client = httpx.AsyncClient(base_url=f"{mux_service_url}")
+        response = await client.get("/api-spec.yaml")
+        yaml_spec = response.content
+        openapi_spec = yaml.load(yaml_spec, Loader=yaml.CLoader)
+        return openapi_spec
 
 def _patched_convert_to_parameter_location(self, param_in: "ParameterLocation") -> str:
     return param_in.value

src/h2ogpte_mcp_server/settings.py

@@ -15,6 +15,7 @@ class Settings(BaseSettings):
     all_endpoints_as_tools: bool = Field(True)
     endpoint_set: EndpointSet = Field(EndpointSet.ALL_WITHOUT_ASYNC_INGEST, case_sensitive=False)
     custom_endpoint_set_file: Optional[str] = Field(None)
+    custom_openapi_spec_file: Optional[str] = Field(None)
 
 settings = Settings(_env_prefix="H2OGPTE_")
 basic_endpoints = [