Merge pull request #14 from ZhuangbilityY/main

add sse transport

Author: XGenerationLab

README.md

@@ -14,7 +14,7 @@
 🌐 <a href="https://github.com/XGenerationLab/XiYan-SQL" >XiYan-SQL</a> |
 📖 <a href="https://arxiv.org/abs/2411.08599"> Arxiv</a> | 
 📄 <a href="https://paperswithcode.com/paper/xiyan-sql-a-multi-generator-ensemble" >PapersWithCode</a>
-💻 <a href="https://huggingface.co/collections/XGenerationLab/xiyansql-models-67c9844307b49f87436808fc">HuggingFace</a> |
+🤗 <a href="https://huggingface.co/collections/XGenerationLab/xiyansql-models-67c9844307b49f87436808fc">HuggingFace</a> |
 🤖 <a href="https://modelscope.cn/collections/XiYanSQL-Models-4483337b614241" >ModelScope</a> |
 🌕 <a href="https://bailian.console.aliyun.com/xiyan">析言GBI</a> 
 <br />
@@ -100,18 +100,13 @@ The following figure illustrates the performance of the XiYan MCP server as meas
 Python 3.11+ is required. 
 You can install the server through pip, and it will install the latest version:
 
-```bash
+```shell
 pip install xiyan-mcp-server
 ```
 
-After that you can directly run the server by:
-```bash
-python -m xiyan_mcp_server
-```
-But it does not provide any functions until you complete following config.
-You will get a yml file. After that you can run the server by:
-```yaml
-env YML=path/to/yml python -m xiyan_mcp_server
+If you want to install the development version from source, you can install from source code on github:
+```shell
+pip install git+https://github.com/XGenerationLab/xiyan_mcp_server.git
 ```
 
 ### Installing from Smithery.ai
@@ -125,11 +120,12 @@ You need a YAML config file to configure the server.
 A default config file is provided in config_demo.yml which looks like this:
 
 ```yaml
+mcp:
+  transport: "stdio"
 model:
   name: "XGenerationLab/XiYanSQL-QwenCoder-32B-2412"
   key: ""
   url: "https://api-inference.modelscope.cn/v1/"
-
 database:
   host: "localhost"
   port: 3306
@@ -138,6 +134,27 @@ database:
   database: ""
 ```
 
+### MCP Configuration
+You can set the transport protocol to ``stdio`` or ``sse``.
+#### STDIO
+For stdio protocol, you can set just like this:
+```yaml
+mcp:
+  transport: "stdio"
+```
+#### SSE
+For sse protocol, you can set mcp config as below:
+```yaml
+mcp:
+  transport: "sse"
+  port: 8000
+  log_level: "INFO"
+```
+The default port is `8000`. You can change the port if needed. 
+The default log level is `ERROR`. We recommend to set log level to `INFO` for more detailed information.
+
+Other configurations like `debug`, `host`, `sse_path`, `message_path` can be customized as well, but normally you don't need to modify them.
+
 ### LLM Configuration
 ``Name`` is the name of the model to use, ``key`` is the API key of the model, ``url`` is the API url of the model. We support following models.
 
@@ -203,53 +220,14 @@ model:
   name: "xiyansql-qwencoder-32b"
   key: "KEY"
   url: "https://xiyan-stream.biz.aliyun.com/service/api/xiyan-sql"
-database:
 ```
 
 Note: this model service is just for trial, if you need to use it in production, please contact us.
 
+##### (3) Local version
 Alternatively, you can also deploy the model [XiYanSQL-qwencoder-32B](https://github.com/XGenerationLab/XiYanSQL-QwenCoder) on your own server.
+See [Local Model](src/xiyan_mcp_server/local_model/README.md) for more details.
 
-#### Local Model
-Note: the local model is slow (about 12 seconds per query on my macbook).
-If you need a stable and fast service, we still recommend to use the modelscope version.
-
-To run xiyan_mcp_server in local mode, you need 
-1) a PC/Mac with at least 16GB RAM
-2) 6GB disk space
-
-Step 1: Install additional Python packages
-```bash
-pip install flask modelscope torch==2.2.2 accelerate>=0.26.0 numpy=2.2.3
-```
-
-Step 2: (optional) manually download the model
-We recommend [xiyansql-qwencoder-3b](https://www.modelscope.cn/models/XGenerationLab/XiYanSQL-QwenCoder-3B-2502/).
-You can manually download the model by
-```bash
-modelscope download --model XGenerationLab/XiYanSQL-QwenCoder-3B-2502
-```
-It will take you 6GB disk space.
-
-Step 3: download the script and run server. src/xiyan_mcp_server/local_xiyan_server.py
-
-
-
-```bash
-python local_xiyan_server.py
-```
-The server will be running on http://localhost:5090/
-
-Step 4: prepare config and run xiyan_mcp_server
-the config.yml should be like:
-```yml
-model:
-  name: "xiyansql-qwencoder-3b"
-  key: "KEY"
-  url: "http://127.0.0.1:5090"
-```
-
-Till now the local mode is ready.
 
 ### Database Configuration
 ``host``, ``port``, ``user``, ``password``, ``database`` are the connection information of the database.
@@ -284,7 +262,21 @@ database:
 
 Note that ``dialect`` should be ``postgresql`` for postgresql.
 ## Launch
-### Claude Desktop
+
+### Server Launch
+
+If you want to launch server with `sse`, you have to run the following command in a terminal:
+```shell
+YML=path/to/yml python -m xiyan_mcp_server
+```
+Then you should see the information on http://localhost:8000/sse in your browser. (Defaultly, change if your mcp server runs on other host/port)
+
+Otherwise, if you use `stdio` transport protocol, you usually declare the mcp server command in specific mcp application instead of launching it in a terminal.
+However, you can still debug with this command if needed.
+
+### Client Setting
+
+#### Claude Desktop
 Add this in your Claude Desktop config file, ref <a href="https://github.com/XGenerationLab/xiyan_mcp_server/blob/main/imgs/claude_desktop.jpg">Claude Desktop config example</a>
 ```json
 {
@@ -303,26 +295,59 @@ Add this in your Claude Desktop config file, ref <a href="https://github.com/XGe
 }
 ```
 **Please note that the Python command here requires the complete path to the Python executable (`/xxx/python`); otherwise, the Python interpreter cannot be found. You can determine this path by using the command `which python`. The same applies to other applications as well.**
-### Cline
-Prepare the config like [Claude Desktop](#claude-desktop)
 
-### Goose
-Add following command in the config, ref <a href="https://github.com/XGenerationLab/xiyan_mcp_server/blob/main/imgs/goose.jpg">Goose config example</a>
+Claude Desktop currently does not support the SSE transport protocol.
 
-```yaml
+#### Cline
+Prepare the config like [Claude Desktop](#claude-desktop)
+
+#### Goose
+If you use `stdio`, add following command in the config, ref <a href="https://github.com/XGenerationLab/xiyan_mcp_server/blob/main/imgs/goose.jpg">Goose config example</a>
+```shell
 env YML=path/to/yml /xxx/python -m xiyan_mcp_server
 ```
-### Cursor
-Use the same command like [Goose](#goose).
+Otherwise, if you use `sse`, change Type to `SSE` and set the endpoint to `http://127.0.0.1:8000/sse`
+#### Cursor
+Use the similar command as follows.
+
+For `stdio`:
+```json
+{
+  "mcpServers": {
+    "xiyan-mcp-server": {
+      "command": "/xxx/python",
+      "args": [
+        "-m",
+        "xiyan_mcp_server"
+      ],
+      "env": {
+        "YML": "path/to/yml"
+      }
+    }
+  }
+}
+```
+For `sse`:
+```json
+{
+  "mcpServers": {
+    "xiyan_mcp_server_1": {
+      "url": "http://localhost:8000/sse"
+    }
+  }
+}
+```
 
 
-### Witsy
+#### Witsy
 Add following in command:
-```yaml
+```shell
 /xxx/python -m xiyan_mcp_server
 ```
 Add an env: key is YML and value is the path to your yml file.
 Ref <a href="https://github.com/XGenerationLab/xiyan_mcp_server/blob/main/imgs/witsy.jpg">Witsy config example</a>
+
+
 ## It Does Not Work!
 Contact us:
 <a href="https://github.com/XGenerationLab/xiyan_mcp_server/blob/main/imgs/dinggroup_out.png">Ding Group钉钉群</a>｜ 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "xiyan_mcp_server"
-version = "0.1.4"
+version = "0.1.5.dev0"
 description = "A Model Context Protocol (MCP) server that utilizes XiyanSQL with databases. This server enables AI assistants to list tables, read data, and execute natural language queries"
 readme = "README.md"
 requires-python = ">=3.11"

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name='xiyan_mcp_server',  # 包的名字
-    version='0.1.2',  # 版本号
+    version='0.1.5.dev0',  # 版本号
     packages=find_packages(),  # 自动找到项目中的包
     install_requires=[  # 必要的包依赖
         # 'numpy',  # 示例：若有依赖包，将其列在此

src/xiyan_mcp_server/__init__.py

@@ -1,8 +0,0 @@
-from . import server
-
-def main():
-    """Main entry point for the package."""
-    server.main()  # 修改为调用server.py中的main函数
-
-# Expose important items at package level
-__all__ = ['main', 'server']

src/xiyan_mcp_server/__main__.py

@@ -1,4 +1,10 @@
-from . import main
+
+from .server import mcp, global_config, mcp_config
+
+def main():
+    mcp.run(transport=mcp_config.get('transport', 'stdio'))
+
+
 
 if __name__ == "__main__":
     main()

src/xiyan_mcp_server/local_model/README.md

@@ -0,0 +1,42 @@
+## Local Model Configuration
+Note: the local model is slow (about 12 seconds per query on my macbook).
+If you need a stable and fast service, we still recommend to use the modelscope version.
+
+To run xiyan_mcp_server in local mode, you need 
+1) a PC/Mac/Machine with at least 16GB RAM
+2) 6GB disk space
+
+The above setting is for model of size 3B. You can adjust the settings to run a 32B model on a server.
+
+### Step 1: Install additional Python packages
+```bash
+pip install flask modelscope torch==2.2.2 accelerate>=0.26.0 numpy=2.2.3
+```
+
+### Step 2: (optional) manually download the model
+We recommend [xiyansql-qwencoder-3b](https://www.modelscope.cn/models/XGenerationLab/XiYanSQL-QwenCoder-3B-2502/).
+You can manually download the model by
+```bash
+modelscope download --model XGenerationLab/XiYanSQL-QwenCoder-3B-2502
+```
+It will take you 6GB disk space.
+
+### Step 3: download the script and run server. 
+
+Script is located at `src/xiyan_mcp_server/local_model/local_xiyan_server.py`
+
+```bash
+python local_xiyan_server.py
+```
+The server will be running on http://localhost:5090/
+
+### Step 4: prepare config and run xiyan_mcp_server
+the config.yml should be like:
+```yml
+model:
+  name: "xiyansql-qwencoder-3b"
+  key: "KEY"
+  url: "http://127.0.0.1:5090"
+```
+
+Till now the local model is ready.

src/xiyan_mcp_server/local_xiyan_server.py renamed to src/xiyan_mcp_server/local_model/local_xiyan_server.py

@@ -1,18 +1,21 @@
+import argparse
 import logging
 import os
+from typing import Literal
 import yaml  # 添加yaml库导入
 
 from mysql.connector import connect, Error
 from mcp.server import  FastMCP
 from mcp.types import TextContent
+
 from .utils.db_config import DBConfig
 from .database_env import DataBaseEnv
 from .utils.db_source import HITLSQLDatabase
 from .utils.db_util import init_db_conn
 from .utils.file_util import extract_sql_from_qwen
 from .utils.llm_util import call_openai_sdk
 
-mcp = FastMCP("xiyan")
+
 
 
 # Configure logging
@@ -22,8 +25,10 @@
 )
 logger = logging.getLogger("xiyan_mcp_server")
 
+
 def get_yml_config():
     config_path = os.getenv("YML", os.path.join(os.path.dirname(__file__), "config_demo.yml"))
+    logger.info(f"Loading configuration from {config_path}")
     try:
         with open(config_path, 'r') as file:
             config = yaml.safe_load(file)
@@ -44,15 +49,18 @@ def get_xiyan_config(db_config):
 
 
 global_config = get_yml_config()
-#print(global_config)
+mcp_config = global_config.get('mcp', {})
 model_config = global_config['model']
 global_db_config = global_config.get('database')
-database_name = global_db_config.get('database','')
 global_xiyan_db_config = get_xiyan_config(global_db_config)
 dialect = global_db_config.get('dialect','mysql')
-#print("dialect is !!!!"+dialect)
 
-@mcp.resource(dialect+'://'+database_name)
+
+
+mcp = FastMCP("xiyan", **mcp_config)
+
+
+@mcp.resource(dialect+'://'+global_db_config.get('database',''))
 async def read_resource() -> str:
 
     db_engine = init_db_conn(global_xiyan_db_config)
@@ -76,7 +84,7 @@ async def read_resource(table_name) -> str:
         raise RuntimeError(f"Database error: {str(e)}")
 
 
-def sql_gen_and_execute(db_env, query: str):
+def sql_gen_and_execute(db_env: DataBaseEnv, query: str):
     """
     Transfers the input natural language question to sql query (known as Text-to-sql) and executes it on the database.
      Args:
@@ -174,17 +182,20 @@ def get_data(query: str)-> list[TextContent]:
     """Fetch the data from database through a natural language query
 
     Args:
-        query: The query in natual language
+        query: The query in natural language
     """
 
     res=call_xiyan(query)
     return [TextContent(type="text", text=res)]
 
-def main():
-    mcp.run()
 
 
-if __name__ == "__main__":
+def main():
+    parser = argparse.ArgumentParser(description="Run MCP server.")
+    parser.add_argument('transport', nargs='?', default='stdio', choices=['stdio', 'sse'],
+                        help='Transport type (stdio or sse)')
+    args = parser.parse_args()
+    mcp.run(transport=args.transport)
 
+if __name__ == "__main__":
     main()
-