add: интеграция локальной нейросети oMLX с добавлением конфигурации и сервера MCP
This commit is contained in:
Executable
+430
@@ -0,0 +1,430 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user