431 lines
18 KiB
Python
Executable File
431 lines
18 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
MCP (Model Context Protocol) сервер для интеграции локальной нейросети oMLX в PyCharm / Junie.
|
|
|
|
Данный скрипт реализует стандартный протокол MCP по транспорту stdio с использованием
|
|
исключительно стандартной библиотеки Python 3.12+ (urllib, json, sys, os).
|
|
Он транслирует команды от внешнего агента (Junie) в запросы к локальному REST API oMLX.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import urllib.error
|
|
import urllib.request
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
|
|
def load_env_file() -> None:
|
|
"""
|
|
Загружает переменные окружения из файла .env в корне проекта,
|
|
если они еще не установлены в системном окружении.
|
|
Не перезаписывает уже существующие переменные окружения.
|
|
"""
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
project_root = os.path.dirname(script_dir)
|
|
candidate_paths = [
|
|
os.path.join(project_root, ".env"),
|
|
os.path.join(os.getcwd(), ".env"),
|
|
]
|
|
|
|
for env_path in candidate_paths:
|
|
if os.path.isfile(env_path):
|
|
try:
|
|
with open(env_path, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
clean_line = line.strip()
|
|
if not clean_line or clean_line.startswith("#") or "=" not in clean_line:
|
|
continue
|
|
key, val = clean_line.split("=", 1)
|
|
key = key.strip()
|
|
val = val.strip().strip('"').strip("'")
|
|
if key and key not in os.environ:
|
|
os.environ[key] = val
|
|
break
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# Автоматическая загрузка конфигурации из .env при старте скрипта
|
|
load_env_file()
|
|
|
|
# Конфигурация по умолчанию (динамически берется из переменных окружения / .env)
|
|
DEFAULT_OMLX_BASE_URL: str = os.environ.get("OMLX_BASE_URL", "http://127.0.0.1:8008/v1")
|
|
DEFAULT_OMLX_API_KEY: str = os.environ.get("OMLX_API_KEY", "")
|
|
DEFAULT_OMLX_MODEL: str = os.environ.get("OMLX_MODEL_NAME", "default")
|
|
REQUEST_TIMEOUT_SECONDS: int = int(os.environ.get("OMLX_TIMEOUT", "120"))
|
|
|
|
|
|
def make_omlx_request(endpoint: str, data: Optional[Dict[str, Any]] = None, method: str = "GET") -> Dict[str, Any]:
|
|
"""
|
|
Выполняет HTTP-запрос к локальному REST API oMLX.
|
|
|
|
:param endpoint: Относительный путь эндпоинта (например, '/chat/completions' или '/models').
|
|
:param data: Тело запроса в формате словаря (для POST-запросов).
|
|
:param method: HTTP-метод ('GET' или 'POST').
|
|
:return: Распарсенный JSON-ответ от сервера oMLX.
|
|
:raises RuntimeError: При ошибках подключения или получении некорректного ответа.
|
|
"""
|
|
base_url = DEFAULT_OMLX_BASE_URL.rstrip("/")
|
|
url = f"{base_url}{endpoint}"
|
|
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {DEFAULT_OMLX_API_KEY}",
|
|
"User-Agent": "MCP-oMLX-Bridge/1.0",
|
|
}
|
|
|
|
payload_bytes = json.dumps(data).encode("utf-8") if data is not None else None
|
|
request = urllib.request.Request(url=url, data=payload_bytes, headers=headers, method=method)
|
|
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=REQUEST_TIMEOUT_SECONDS) as response:
|
|
raw_body = response.read().decode("utf-8")
|
|
return json.loads(raw_body)
|
|
except urllib.error.HTTPError as error:
|
|
err_msg = error.read().decode("utf-8", errors="replace")
|
|
raise RuntimeError(f"Ошибка oMLX API HTTP {error.code}: {err_msg}") from error
|
|
except urllib.error.URLError as error:
|
|
raise RuntimeError(
|
|
f"Не удалось подключиться к oMLX по адресу {url}. Проверьте, запущен ли oMLX. Ошибка: {error.reason}"
|
|
) from error
|
|
except Exception as error:
|
|
raise RuntimeError(f"Непредвиденная ошибка при запросе к oMLX: {str(error)}") from error
|
|
|
|
|
|
def get_available_models() -> List[str]:
|
|
"""
|
|
Получает список идентификаторов моделей, доступных на сервере oMLX.
|
|
|
|
:return: Список строк с названиями моделей.
|
|
"""
|
|
try:
|
|
result = make_omlx_request("/models", method="GET")
|
|
data = result.get("data", [])
|
|
return [m.get("id", "") for m in data if isinstance(m, dict) and m.get("id")]
|
|
except Exception:
|
|
return []
|
|
|
|
|
|
def resolve_model_name(requested_model: str = "") -> str:
|
|
"""
|
|
Определяет точное имя модели для использования в запросе.
|
|
Если имя не указано или равно 'default', пытается автоматически выбрать
|
|
первую доступную модель из списка активных на локальном сервере oMLX.
|
|
|
|
:param requested_model: Запрошенное пользователем или конфигурацией имя модели.
|
|
:return: Выбранное имя модели.
|
|
"""
|
|
clean_model = requested_model.strip()
|
|
if clean_model and clean_model.lower() != "default":
|
|
return clean_model
|
|
|
|
if DEFAULT_OMLX_MODEL and DEFAULT_OMLX_MODEL.lower() != "default":
|
|
return DEFAULT_OMLX_MODEL
|
|
|
|
available_models = get_available_models()
|
|
if available_models:
|
|
return available_models[0]
|
|
|
|
return "default"
|
|
|
|
|
|
def tool_ask_local_model(
|
|
prompt: str,
|
|
system_prompt: str = "",
|
|
model: str = "",
|
|
temperature: float = 0.2,
|
|
max_tokens: Optional[int] = None,
|
|
) -> str:
|
|
"""
|
|
Отправляет запрос к локальной модели oMLX для генерации текста или решения задачи.
|
|
|
|
:param prompt: Текст запроса (пользовательский промпт).
|
|
:param system_prompt: Опциональная системная инструкция.
|
|
:param model: Имя модели (если пусто, выбирается активная модель oMLX).
|
|
:param temperature: Температура генерации (от 0.0 до 1.0).
|
|
:param max_tokens: Ограничение на максимальное количество токенов в ответе.
|
|
:return: Сгенерированный локальной моделью текст ответа.
|
|
"""
|
|
target_model = resolve_model_name(model)
|
|
|
|
messages: List[Dict[str, str]] = []
|
|
if system_prompt:
|
|
messages.append({"role": "system", "content": system_prompt})
|
|
messages.append({"role": "user", "content": prompt})
|
|
|
|
payload: Dict[str, Any] = {
|
|
"model": target_model,
|
|
"messages": messages,
|
|
"temperature": temperature,
|
|
}
|
|
if max_tokens is not None:
|
|
payload["max_tokens"] = max_tokens
|
|
|
|
result = make_omlx_request("/chat/completions", data=payload, method="POST")
|
|
|
|
choices = result.get("choices", [])
|
|
if not choices:
|
|
return "Ответ от oMLX пуст или не содержит вариантов (choices)."
|
|
|
|
message_content = choices[0].get("message", {}).get("content", "")
|
|
return str(message_content)
|
|
|
|
|
|
def tool_check_local_model_status() -> str:
|
|
"""
|
|
Проверяет статус подключения к oMLX и возвращает список доступных локальных моделей.
|
|
|
|
:return: Строка с отчетом о доступности сервера и списке загруженных моделей.
|
|
"""
|
|
try:
|
|
result = make_omlx_request("/models", method="GET")
|
|
data = result.get("data", [])
|
|
model_ids = [m.get("id", "unknown") for m in data if isinstance(m, dict)]
|
|
models_str = ", ".join(model_ids) if model_ids else "Модели не обнаружены в ответе"
|
|
return f"Сервер oMLX доступен по адресу {DEFAULT_OMLX_BASE_URL}. Доступные модели: {models_str}."
|
|
except Exception as error:
|
|
return f"Проверка статуса завершилась ошибкой: {str(error)}"
|
|
|
|
|
|
def tool_review_code_locally(code_snippet: str, instructions: str = "") -> str:
|
|
"""
|
|
Выполняет локальный анализ и ревью фрагмента кода силами oMLX.
|
|
|
|
:param code_snippet: Исходный код для анализа.
|
|
:param instructions: Дополнительные требования к проверке (например, поиск багов, оптимизация).
|
|
:return: Замечания и рекомендации локальной модели.
|
|
"""
|
|
sys_prompt = (
|
|
"Ты — опытный локальный ассистент по ревью кода. "
|
|
"Проанализируй предоставленный фрагмент кода, найди потенциальные ошибки, "
|
|
"узкие места и предложи конкретные улучшения без лишних вводных слов."
|
|
)
|
|
user_prompt = f"Код для ревью:\n```\n{code_snippet}\n```"
|
|
if instructions:
|
|
user_prompt += f"\n\nСпециальные указания:\n{instructions}"
|
|
|
|
return tool_ask_local_model(prompt=user_prompt, system_prompt=sys_prompt, temperature=0.1)
|
|
|
|
|
|
# Спецификация поддерживаемых MCP-инструментов
|
|
MCP_TOOLS: List[Dict[str, Any]] = [
|
|
{
|
|
"name": "ask_local_model",
|
|
"description": (
|
|
"Отправляет запрос локальной модели oMLX (на Apple Silicon) и возвращает ответ. "
|
|
"Используется для делегирования локальных вычислений, обработки больших текстов, "
|
|
"генерации вариантов или сохранения приватности."
|
|
),
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"prompt": {
|
|
"type": "string",
|
|
"description": "Основной текст запроса или задачи для локальной модели.",
|
|
},
|
|
"system_prompt": {
|
|
"type": "string",
|
|
"description": "Опциональная системная инструкция для локальной модели.",
|
|
},
|
|
"model": {
|
|
"type": "string",
|
|
"description": "Имя конкретной модели oMLX. Если не указано, берется активная модель.",
|
|
},
|
|
"temperature": {
|
|
"type": "number",
|
|
"description": "Температура генерации (по умолчанию 0.2).",
|
|
},
|
|
"max_tokens": {
|
|
"type": "integer",
|
|
"description": "Максимальное количество токенов в ответе.",
|
|
},
|
|
},
|
|
"required": ["prompt"],
|
|
},
|
|
},
|
|
{
|
|
"name": "check_local_model_status",
|
|
"description": "Проверяет доступность локального сервера oMLX и возвращает список моделей.",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {},
|
|
},
|
|
},
|
|
{
|
|
"name": "review_code_locally",
|
|
"description": "Запускает локальное ревью фрагмента кода на модели oMLX.",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"code_snippet": {
|
|
"type": "string",
|
|
"description": "Фрагмент кода для анализа.",
|
|
},
|
|
"instructions": {
|
|
"type": "string",
|
|
"description": "Особые пожелания или фокус внимания при анализе.",
|
|
},
|
|
"model": {
|
|
"type": "string",
|
|
"description": "Имя конкретной модели oMLX.",
|
|
},
|
|
},
|
|
"required": ["code_snippet"],
|
|
},
|
|
},
|
|
]
|
|
|
|
|
|
def send_response(response: Dict[str, Any]) -> None:
|
|
"""
|
|
Отправляет JSON-RPC ответ в стандартный поток вывода stdout.
|
|
|
|
:param response: Словарь с ответом по спецификации JSON-RPC 2.0.
|
|
"""
|
|
body = json.dumps(response, ensure_ascii=False)
|
|
sys.stdout.write(body + "\n")
|
|
sys.stdout.flush()
|
|
|
|
|
|
def handle_request(message: Dict[str, Any]) -> None:
|
|
"""
|
|
Обрабатывает входящее JSON-RPC сообщение протокола MCP.
|
|
|
|
:param message: Распарсенное входящее сообщение.
|
|
"""
|
|
msg_id = message.get("id")
|
|
method = message.get("method")
|
|
params = message.get("params", {})
|
|
|
|
# Уведомления без id (например, notifications/initialized) не требуют ответа
|
|
if msg_id is None:
|
|
return
|
|
|
|
if method == "initialize":
|
|
# Ответ на инициализацию сервера MCP
|
|
send_response({
|
|
"jsonrpc": "2.0",
|
|
"id": msg_id,
|
|
"result": {
|
|
"protocolVersion": "2024-11-05",
|
|
"capabilities": {
|
|
"tools": {}
|
|
},
|
|
"serverInfo": {
|
|
"name": "omlx-local-bridge",
|
|
"version": "1.0.0"
|
|
}
|
|
}
|
|
})
|
|
return
|
|
|
|
if method == "ping":
|
|
send_response({"jsonrpc": "2.0", "id": msg_id, "result": {}})
|
|
return
|
|
|
|
if method == "tools/list":
|
|
send_response({
|
|
"jsonrpc": "2.0",
|
|
"id": msg_id,
|
|
"result": {
|
|
"tools": MCP_TOOLS
|
|
}
|
|
})
|
|
return
|
|
|
|
if method == "tools/call":
|
|
tool_name = params.get("name")
|
|
arguments = params.get("arguments", {})
|
|
|
|
try:
|
|
if tool_name == "ask_local_model":
|
|
max_tokens_val = arguments.get("max_tokens")
|
|
output = tool_ask_local_model(
|
|
prompt=arguments.get("prompt", ""),
|
|
system_prompt=arguments.get("system_prompt", ""),
|
|
model=arguments.get("model", ""),
|
|
temperature=float(arguments.get("temperature", 0.2)),
|
|
max_tokens=int(max_tokens_val) if max_tokens_val is not None else None,
|
|
)
|
|
elif tool_name == "check_local_model_status":
|
|
output = tool_check_local_model_status()
|
|
elif tool_name == "review_code_locally":
|
|
output = tool_review_code_locally(
|
|
code_snippet=arguments.get("code_snippet", ""),
|
|
instructions=arguments.get("instructions", ""),
|
|
)
|
|
else:
|
|
send_response({
|
|
"jsonrpc": "2.0",
|
|
"id": msg_id,
|
|
"error": {
|
|
"code": -32601,
|
|
"message": f"Неизвестный инструмент: {tool_name}"
|
|
}
|
|
})
|
|
return
|
|
|
|
send_response({
|
|
"jsonrpc": "2.0",
|
|
"id": msg_id,
|
|
"result": {
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": output
|
|
}
|
|
],
|
|
"isError": False
|
|
}
|
|
})
|
|
except Exception as error:
|
|
send_response({
|
|
"jsonrpc": "2.0",
|
|
"id": msg_id,
|
|
"result": {
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": f"Ошибка выполнения инструмента {tool_name}: {str(error)}"
|
|
}
|
|
],
|
|
"isError": True
|
|
}
|
|
})
|
|
return
|
|
|
|
# Неподдерживаемый метод
|
|
send_response({
|
|
"jsonrpc": "2.0",
|
|
"id": msg_id,
|
|
"error": {
|
|
"code": -32601,
|
|
"message": f"Метод не поддерживается: {method}"
|
|
}
|
|
})
|
|
|
|
|
|
def main() -> None:
|
|
"""Главный цикл чтения сообщений из stdin."""
|
|
for line in sys.stdin:
|
|
clean_line = line.strip()
|
|
if not clean_line:
|
|
continue
|
|
try:
|
|
message = json.loads(clean_line)
|
|
handle_request(message)
|
|
except json.JSONDecodeError:
|
|
# Игнорируем некорректные JSON-строки
|
|
continue
|
|
except Exception as error:
|
|
# Логируем фатальные сбои в stderr
|
|
sys.stderr.write(f"Ошибка MCP сервера: {str(error)}\n")
|
|
sys.stderr.flush()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|