Merge pull request #40 from CogStack/doc-n-reorg

Improve READMEs and refactor the OAuth flow

Author: baixiac

README.md

@@ -2,6 +2,12 @@
 
 Cogstack ModelServe (CMS) is a model-serving and model-governance system created for a range of CogStack NLP tasks. Targeting language models with NER and entity linking capabilities, CMS provides a one-stop shop for serving and fine-tuning models, training lifecycle management, as well as monitoring and end-to-end observability.
 
+[![build](https://img.shields.io/github/actions/workflow/status/CogStack/CogStack-ModelServe/main.yaml)](https://github.com/CogStack/CogStack-ModelServe/actions)
+[![release](https://img.shields.io/github/v/release/CogStack/CogStack-ModelServe)](https://github.com/CogStack/CogStack-ModelServe/releases)
+[![docker](https://img.shields.io/docker/v/cogstacksystems/cogstack-modelserve?sort=semver&label=docker)](https://hub.docker.com/r/cogstacksystems/cogstack-modelserve)
+[![python](https://img.shields.io/badge/python-3.10%20%7C%203.11%20%7C%203.12-blue)](https://www.python.org/)
+[![license](https://img.shields.io/badge/license-Apache%20License%202.0-blue)](https://github.com/CogStack/CogStack-ModelServe/blob/main/LICENSE)
+
 ## Install Dependencies
 A virtual environment is highly recommended prior to installation. To install the dependencies, run:
 ```commandline
@@ -254,6 +260,15 @@ Note that to enable quantization and training features, you need to install the
 ```commandline
 pip install '.[llm]'
 ```
+
+### CMS MCP server
+You can run an MCP server to expose local or remote CMS capabilities to your preferred MCP client.
+To that end, install the following extra dependencies:
+```commandline
+pip install '.[mcp]'
+```
+For detailed configuration, please refer to the [CMS MCP Server docs](./app/mcp/README.md).
+
 #### Chat with served models
 You can also "chat" with the running model using the `/stream/ws` endpoint. For example:
 ```html

app/api/utils.py

@@ -323,8 +323,8 @@ async def init_vllm_engine(app: FastAPI,
         )
         from vllm import SamplingParams, TokensPrompt
     except ImportError:
-        logger.error("Cannot import the vLLM engine. Please install it with `pip install cms[llm]`.")
-        raise ExtraDependencyRequiredException("Cannot import the vLLM engine. Please install it with `pip install cms[llm]`.")
+        logger.error("Cannot import the vLLM engine. Please install it with `pip install '.[llm]'`.")
+        raise ExtraDependencyRequiredException("Cannot import the vLLM engine. Please install it with `pip install '.[llm]'`.")
 
     parser = FlexibleArgumentParser()
     parser = make_arg_parser(parser)

app/cli/cli.py

@@ -599,7 +599,7 @@ def run_mcp_server(
     transport: str = typer.Option("http", help="The transport type (either 'stdio', 'sse' or 'http')"),
     cms_base_url: str = typer.Option("http://127.0.0.1:8000", help="The base URL of the CMS API"),
     cms_api_key: str = typer.Option("Bearer", help="The API key for authenticating with the CMS API"),
-    mcp_api_keys: str = typer.Option("", help="Comma-separated API keys for authenticating MCP clients"),
+    cms_mcp_api_keys: str = typer.Option("", help="Comma-separated API keys for authenticating CMS MCP clients"),
     cms_mcp_oauth_enabled: Optional[bool] = typer.Option(None, help="Whether to enable OAuth2 authentication for MCP clients"),
     github_client_id: str = typer.Option("", help="The GitHub OAuth2 client ID"),
     github_client_secret: str = typer.Option("", help="The GitHub OAuth2 client secret"),
@@ -618,6 +618,13 @@ def run_mcp_server(
         port (int): The port of the MCP server. Defaults to 8080.
         transport (str): The transport type for the MCP server. Can be "stdio" or "http". Defaults to "stdio".
         cms_base_url (str): The base URL of the CMS API endpoint. Defaults to "http://localhost:8000".
+        cms_api_key (str): The API key for authenticating with the CMS API. Defaults to "Bearer".
+        cms_mcp_api_keys (str): Comma-separated API keys for authenticating CMS MCP clients. Defaults to "".
+        cms_mcp_oauth_enabled (Optional[bool]): Whether to enable OAuth2 authentication for MCP clients. Defaults to None.
+        github_client_id (str): The GitHub OAuth2 client ID, required if cms_mcp_oauth_enabled is True. Defaults to "".
+        github_client_secret (str): The GitHub OAuth2 client secret, required if cms_mcp_oauth_enabled is True. Defaults to an "".
+        google_client_id (str): The Google OAuth2 client ID, required if cms_mcp_oauth_enabled is True. Defaults to an "".
+        google_client_secret (str): The Google OAuth2 client secret, required if cms_mcp_oauth_enabled is True. Defaults to an "".
         debug (Optional[bool]): Run in debug mode if set to True.
     """
 
@@ -629,7 +636,7 @@ def run_mcp_server(
     os.environ["CMS_MCP_SERVER_PORT"] = str(port)
     os.environ["CMS_MCP_TRANSPORT"] = transport.lower()
     os.environ["CMS_API_KEY"] = cms_api_key
-    os.environ["MCP_API_KEYS"] = mcp_api_keys
+    os.environ["CMS_MCP_API_KEYS"] = cms_mcp_api_keys
     os.environ["CMS_MCP_OAUTH_ENABLED"] = "true" if cms_mcp_oauth_enabled else "false"
     os.environ["GITHUB_CLIENT_ID"] = github_client_id
     os.environ["GITHUB_CLIENT_SECRET"] = github_client_secret
@@ -645,9 +652,9 @@ def run_mcp_server(
         logger.info(f"Connected to CMS API at {cms_base_url}")
         main()
     except ImportError as e:
-        logger.error(f"Cannot import MCP. Please install it with `pip install cms[mcp]`: {e}")
+        logger.error(f"Cannot import MCP. Please install it with `pip install '.[mcp]'`: {e}")
         typer.echo(f"ERROR: Cannot import MCP: {e}")
-        typer.echo("Please install it with `pip install cms[mcp]`.")
+        typer.echo("Please install it with `pip install '.[mcp]'`.")
         raise typer.Exit(code=1)
     except KeyboardInterrupt:
         logger.info("MCP server stopped by the user")

app/mcp/README.md

@@ -12,7 +12,7 @@ pip install '.[mcp]'
 ### 2. Set Environment Variables
 ```bash
 export CMS_BASE_URL="http://127.0.0.1:8000"  # CogStack ModelServe API base URL
-export MCP_API_KEYS="key1,key2,...keyN"      # Optional: The API key(s) for authentication
+export CMS_MCP_API_KEYS="key1,key2,...keyN"  # Optional: The API key(s) for authentication
 ```
 
 ### 3. Run the Server
@@ -23,10 +23,84 @@ cms mcp run
 
 ```bash
 # HTTP transport
-export CMS_MCP_TRANSPORT=http
 cms mcp run --transport http
 ```
-Once the above succeeds, the MCP server will be running on http://127.0.0.1:8080/mcp
+Once the above succeeds, the MCP server will be running at http://127.0.0.1:8080/mcp
+
+```bash
+# SSE transport
+cms mcp run --transport sse
+```
+Once the above succeeds, the MCP server will be running at http://127.0.0.1:8080/sse
+
+## Claude Desktop Configuration
+
+To use this MCP server with Claude Desktop, add the following configuration to your `claude_desktop_config.json` file:
+
+```json
+{
+  "mcpServers": {
+    "cms-mcp-server": {
+      "command": "npx",
+      "args": [
+        "mcp-remote",
+        "http://127.0.0.1:8080/sse"
+      ]
+    }
+  }
+}
+```
+
+With API-key-based authentication:
+```bash
+cms mcp run --transport sse --cms-mcp-api-keys "key1,key2,...keyN"
+```
+```json
+{
+  "mcpServers": {
+    "cms-mcp-server": {
+      "command": "npx",
+      "args": [
+        "mcp-remote",
+        "http://127.0.0.1:8080/sse",
+        "--header",
+        "X-API-Key:${API_KEY_HEADER}"
+      ],
+      "env": {
+        "API_KEY_HEADER": "ONE_OF_THE_API_KEYS"
+      }
+    }
+  }
+}
+```
+
+With OAuth2-based authentication:
+```bash
+cms mcp run --transport sse
+  --cms-mcp-oauth-enabled \
+  --github-client-id <GITHUB_CLIENT_ID> \
+  --github-client-secret <GITHUB_CLIENT_SECRET> \
+  --google-client-id <GOOGLE_CLIENT_ID> \
+  --google-client-secret <GOOGLE_CLIENT_SECRET>
+```
+```json
+{
+  "mcpServers": {
+    "cms-mcp-server": {
+      "command": "npx",
+      "args": [
+        "mcp-remote",
+        "http://127.0.0.1:8080/sse",
+        "--header",
+        "X-API-Key:${AUTH_HEADER}"
+      ],
+      "env": {
+        "AUTH_HEADER": "Bearer <ACCESS_TOKEN>"
+      }
+    }
+  }
+}
+```
 
 ## Available Tools
 
@@ -40,27 +114,27 @@ Once the above succeeds, the MCP server will be running on http://127.0.0.1:8080
 
 ## Configuration
 
-| Environment Variable | Default | Description |
-|---------------------|---------|-------------|
+| Environment Variable | Default                 | Description |
+|---------------------|-------------------------|-------------|
 | `CMS_BASE_URL` | `http://127.0.0.1:8000` | ModelServe API base URL |
 | `CMS_MCP_SERVER_HOST` | `127.0.0.1` | MCP server host |
 | `CMS_MCP_SERVER_PORT` | `8080` | MCP server port |
 | `CMS_MCP_TRANSPORT` | `stdio` | Transport type (`stdio`, `http` or `sse`) |
 | `CMS_ACCESS_TOKEN` | Empty | Bearer token for ModelServe API |
 | `CMS_API_KEY` | `Bearer` | API key for ModelServe API |
-| `MCP_API_KEYS` | None | Comma-separated API keys for authentication |
-| `CMS_MCP_OAUTH_ENABLED` | `true` | Enable OAuth authentication |
-| `CMS_MCP_BASE_URL` | `http://<host>:<port>` | Base URL for OAuth callback |
-| `CMS_MCP_DEV` | `0` | Run in development mode (creates server instance) |
+| `CMS_MCP_API_KEYS` | None | Comma-separated API keys for authentication |
+| `CMS_MCP_OAUTH_ENABLED` | `false` | Enable OAuth authentication |
+| `CMS_MCP_BASE_URL` | `http://<host>:<port>`  | Base URL for OAuth callback |
+| `CMS_MCP_DEV` | `0` | Run in development mode |
 
 
 ## Authentication
 
 The server supports two authentication methods:
 
 ### 1. API Key Authentication
-When `MCP_API_KEYS` is set, clients must authenticate using:
-- **Header**: `x-api-key: your-key`
+When `CMS_MCP_API_KEYS` is set, clients must authenticate using:
+- **Header**: `X-API-Key: your-key`
 
 ### 2. OAuth Authentication (SSE Transport)
 When `CMS_MCP_OAUTH_ENABLED=true`, the server provides a built-in OAuth 2.0 login flow for SSE transport.
@@ -80,7 +154,7 @@ When `CMS_MCP_OAUTH_ENABLED=true`, the server provides a built-in OAuth 2.0 logi
 | `GOOGLE_CLIENT_ID` | Google OAuth client ID |
 | `GOOGLE_CLIENT_SECRET` | Google OAuth client secret |
 
-**Note:** OAuth credentials can also be set via environment variables or `.env` file. If not configured, the server will log a warning but continue running.
+**Note:** If OAuth credentials are not configured, the server will log a warning but continue running.
 
 **Session Authentication:**
 After OAuth login, a session cookie (`cms_mcp_session`) is set. Subsequent MCP requests should include this cookie for authentication.