vulndify-mcp-server

--- vulndify_mcp_server-0.3.0/src/vulndify_mcp_server/server.py (excerpt) ---
"""Minimal MCP server exposing a demo hello tool."""

from mcp.server.fastmcp import FastMCP
from pydantic import BaseModel
import os
import sys
import signal
import requests
import json

SERVER_NAME = "vulndify-demo"

mcp = FastMCP(name=SERVER_NAME)


class QueryResponse(BaseModel):
    conversation_id: str
    message: str


def signal_handler(sig, frame):
    print("Thanks for using mcp server...")
    sys.exit(0)


def hello() -> str:
    """Return the fixed response used by the demo tool."""
    return "hello"


@mcp.tool(name="hello")
def hello_tool() -> str:
    """Return a fixed string for demonstration."""
    return hello()



@mcp.tool()
def query(target, message, conversation_id='')->QueryResponse:
    """
    发送消息给目标应用，并获得回复消息
    
    Parameters:
    - target(str): 目标应用的APP ID
    - message(str): 发送的消息内容
    - conversation_id(str): 会话ID，默认为空

    Return(tuple): 
    - 会话ID
    - 回复内容
    """
    url = "http://10.37.1.177/v1/chat-messages"

    # 请求头
    headers = {
        "Authorization": f"Bearer {target}",
        "Content-Type": "application/json"
    }

    # 请求体数据
    payload = {
        "inputs": {},
        "query": message,
        "response_mode": "streaming",
        "conversation_id": conversation_id,
        "user": 'test',
        "files": []
    }

    try:
        # 发送POST请求
        response = requests.post(
            url,
            headers=headers,
            data=json.dumps(payload),
            stream=True  # 对于流式响应需要设置stream=True