62 lines
2.2 KiB
Python
62 lines
2.2 KiB
Python
|
import asyncio
|
||
|
|
from typing import Any, Dict, List, Optional
|
||
|
|
from django.conf import settings
|
||
|
|
from mcp_agent.mcp_client import MCPClient
|
||
|
|
from .models import AgentModel
|
||
|
|
|
||
|
|
|
||
|
|
async def _call_mcp(tool: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
|
||
|
|
"""Internal async helper to call the MCP HTTP bridge via MCPClient."""
|
||
|
|
server_url = getattr(settings, "MCP_AGENT_URL")
|
||
|
|
client = MCPClient(server_url)
|
||
|
|
try:
|
||
|
|
resp = await client.send(tool, arguments)
|
||
|
|
return resp
|
||
|
|
finally:
|
||
|
|
await client.close()
|
||
|
|
|
||
|
|
|
||
|
|
def fine_tune_model(
|
||
|
|
base_model: str,
|
||
|
|
training_files: List[str],
|
||
|
|
hyperparams: Dict[str, Any],
|
||
|
|
name: str,
|
||
|
|
version: str,
|
||
|
|
) -> Dict[str, Any]:
|
||
|
|
"""Synchronously request a fine-tune run on the MCP server.
|
||
|
|
|
||
|
|
Expects the MCP tool `fine_tune` to accept: {base_model, training_files, hyperparams, name, version}
|
||
|
|
and to return a JSON-like dict containing at least `status` and on success `model_path` and `version`.
|
||
|
|
"""
|
||
|
|
return asyncio.run(_call_mcp("fine_tune", {
|
||
|
|
"base_model": base_model,
|
||
|
|
"training_files": training_files,
|
||
|
|
"hyperparams": hyperparams,
|
||
|
|
"name": name,
|
||
|
|
"version": version,
|
||
|
|
}))
|
||
|
|
|
||
|
|
|
||
|
|
def load_model_for_inference(model_path: str) -> Dict[str, Any]:
|
||
|
|
"""Tell the MCP server to load a model into memory/serving for inference.
|
||
|
|
|
||
|
|
Expects the MCP tool `load_model` with {model_path} returning status info.
|
||
|
|
"""
|
||
|
|
return asyncio.run(_call_mcp("load_model", {"model_path": model_path}))
|
||
|
|
|
||
|
|
|
||
|
|
def infer_with_model(model_path: str, prompt: str, options: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||
|
|
"""Request inference from the MCP server using a previously fine-tuned model.
|
||
|
|
|
||
|
|
Calls the MCP tool `infer` with {model_path, prompt, options}.
|
||
|
|
"""
|
||
|
|
return asyncio.run(_call_mcp("infer", {"model_path": model_path, "prompt": prompt, "options": options or {}}))
|
||
|
|
|
||
|
|
|
||
|
|
def register_model_in_db(name: str, version: str, model_path: str) -> AgentModel:
|
||
|
|
"""Convenience DB helper: create and return an AgentModel record.
|
||
|
|
|
||
|
|
NOTE: migrations are required after the model field change prior to using this in production.
|
||
|
|
"""
|
||
|
|
return AgentModel.objects.create(name=name, version=version, path=model_path)
|