vasiuuu/DGX_AI
0
1from __future__ import annotations2 3import json4import logging5from collections.abc import Callable6from pathlib import Path7from typing import Any8from uuid import uuid49 10from codeforge.environment import CodeForgeEnvironment11from codeforge.models import CodeForgeAction, CodeForgeActionType12from codeforge.ralph.synthesizer import Synthesizer13from codeforge.tasks import TASKS14 15_log = logging.getLogger(__name__)16_VERSION = "0.2.0"17 18_Handler = Callable[19 ["CodeForgeMCPServer", "dict[str, Any]"],20 "dict[str, Any]",21]22 23_SID_DESC = "Session ID from codeforge_reset."24_SID_PROP: dict[str, str] = {25 "type": "string",26 "description": _SID_DESC,27}28 29# -------------------------------------------------------------------30# Tool schema definitions (SYSTEM_DESIGN §9.1)31# -------------------------------------------------------------------32 33_TOOL_DEFS: tuple[dict[str, Any], ...] = (34 {35 "name": "codeforge_reset",36 "description": (37 "Start a new CodeForge episode. You will receive a task "38 "brief and initial files. Your goal is to produce working "39 "Python code that passes sandbox verification. Budget is "40 "limited — plan your actions carefully."41 ),42 "inputSchema": {43 "type": "object",44 "properties": {45 "task_level": {46 "type": "string",47 "enum": ["easy", "medium", "hard"],48 "description": (49 "Difficulty level. Easy: single file, budget "50 "4. Medium: multi-file with tests, budget 6. "51 "Hard: three-file module, budget 10."52 ),53 },54 },55 "required": ["task_level"],56 },57 },58 {59 "name": "codeforge_query_kb",60 "description": (61 "Search the coding skills knowledge base. Returns real "62 "documentation from 1006 skill nodes. Use this to find "63 "patterns, best practices, and guidance BEFORE writing "64 "code. Costs 1 budget unit. DO NOT guess library APIs — "65 "search for them here or verify via documentation first."66 ),67 "inputSchema": {68 "type": "object",69 "properties": {70 "session_id": _SID_PROP,71 "claim": {72 "type": "string",73 "description": (74 "What you want to find guidance on. Be "75 "specific. Example: 'pytest fixture patterns "76 "for testing greet functions'"77 ),78 },79 "top_k": {80 "type": "integer",81 "default": 5,82 "minimum": 1,83 "maximum": 20,84 "description": "Number of results to return",85 },86 "required_tags": {87 "type": "array",88 "items": {"type": "string"},89 "default": [],90 "description": (91 "Only return nodes that have ALL of these "92 "tags"93 ),94 },95 },96 "required": ["session_id", "claim"],97 },98 },99 {100 "name": "codeforge_query_cluster",101 "description": (102 "Browse a skill cluster by label. Clusters are communities "103 "of related skill nodes grouped by Jaccard similarity. Use "104 "this to explore a topic area deeply. Costs 1 budget unit."105 ),106 "inputSchema": {107 "type": "object",108 "properties": {109 "session_id": _SID_PROP,110 "cluster_label": {111 "type": "string",112 "description": (113 "The cluster label to look up. Example: "114 "'python_testing_pytest_fixtures'"115 ),116 },117 "top_k": {118 "type": "integer",119 "default": 10,120 "minimum": 1,121 "maximum": 50,122 },123 },124 "required": ["session_id", "cluster_label"],125 },126 },127 {128 "name": "codeforge_interrogate",129 "description": (130 "Get Socratic questions about the task that cite real skill "131 "corpus nodes. Use this BEFORE writing code to identify "132 "edge cases, success criteria, and assumptions you might be "133 "wrong about. Costs 1 budget unit."134 ),135 "inputSchema": {136 "type": "object",137 "properties": {138 "session_id": _SID_PROP,139 "brief_override": {140 "type": "string",141 "description": (142 "Optional override for the task brief. "143 "If omitted, uses the current task brief."144 ),145 },146 },147 "required": ["session_id"],148 },149 },150 {151 "name": "codeforge_run_ralph",152 "description": (153 "Run autonomous improvement iterations on your current "154 "code. Each iteration: synthesize improvement → "155 "sandbox-score → keep if better. Costs max_iters budget "156 "units. Wasted iterations (no improvement) cost 0.05 "157 "penalty each. Use when you want the environment to "158 "iteratively improve your code."159 ),160 "inputSchema": {161 "type": "object",162 "properties": {163 "session_id": _SID_PROP,164 "max_iters": {165 "type": "integer",166 "default": 3,167 "minimum": 1,168 "maximum": 10,169 "description": (170 "Maximum iterations. Each costs 1 budget. "171 "Choose carefully."172 ),173 },174 },175 "required": ["session_id", "max_iters"],176 },177 },178 {179 "name": "codeforge_submit",180 "description": (181 "Submit Python files for grading. Your code will be: "182 "(1) written to a sandbox and checked by ruff, mypy "183 "--strict, pytest, and import resolution — these are REAL "184 "tools, not mocks; (2) AST-grounded to verify every "185 "import and attribute access resolves to a real Python "186 "module/attribute; (3) scored via quality = 0.6*sandbox + "187 "0.4*groundedness; (4) if you provide confidence, "188 "Brier-penalized: reward = quality * (1 - "189 "min((confidence-quality)^2, 0.5)). DO NOT fabricate "190 "library names or API signatures — the grounder WILL "191 "catch them and your score WILL drop."192 ),193 "inputSchema": {194 "type": "object",195 "properties": {196 "session_id": _SID_PROP,197 "files": {198 "type": "object",199 "additionalProperties": {"type": "string"},200 "description": (201 "Map of filename to file content. Example: "202 '{"main.py": "def greet(name: str) -> str:'203 "\\n return f'Hello, {name}!'\\n\"}"204 ),205 },206 "confidence": {207 "type": "number",208 "minimum": 0.0,209 "maximum": 1.0,210 "description": (211 "Your confidence that this submission is "212 "correct (0.0 = no idea, 1.0 = certain). "213 "Overconfidence on bad code is PENALIZED. "214 "Honest uncertainty is treated more "215 "favorably. If you are unsure, say so."216 ),217 },218 },219 "required": ["session_id", "files"],220 },221 },222 {223 "name": "codeforge_get_audit",224 "description": (225 "Read the audit trail for the current episode (or a "226 "specific run). Returns every action taken, every citation "227 "made, every reward earned, and the evidence behind each. "228 "Costs 0 budget. Use this to review your progress and "229 "understand what worked."230 ),231 "inputSchema": {232 "type": "object",233 "properties": {234 "session_id": _SID_PROP,235 "target_run_id": {236 "type": "string",237 "description": (238 "Optional run ID to audit. "239 "Defaults to current episode."240 ),241 },242 },243 "required": ["session_id"],244 },245 },246 {247 "name": "codeforge_state",248 "description": (249 "Get current episode state without taking an action. Shows "250 "task brief, current files, budget remaining, last reward, "251 "and whether the episode is done. Costs 0 budget."252 ),253 "inputSchema": {254 "type": "object",255 "properties": {"session_id": _SID_PROP},256 "required": ["session_id"],257 },258 },259 {260 "name": "codeforge_list_clusters",261 "description": (262 "List all available cluster labels and their node counts. "263 "Use this to discover what topic areas exist before "264 "calling codeforge_query_cluster. Costs 0 budget."265 ),266 "inputSchema": {267 "type": "object",268 "properties": {"session_id": _SID_PROP},269 },270 },271 {272 "name": "codeforge_list_tags",273 "description": (274 "List all available tags in the skill corpus. Use this to "275 "discover valid values for the required_tags parameter. "276 "Costs 0 budget."277 ),278 "inputSchema": {279 "type": "object",280 "properties": {"session_id": _SID_PROP},281 },282 },283)284 285# -------------------------------------------------------------------286# Resource definitions287# -------------------------------------------------------------------288 289_RESOURCE_DEFS: tuple[dict[str, str], ...] = (290 {291 "uri": "codeforge://corpus/stats",292 "name": "Corpus Statistics",293 "description": (294 "Corpus statistics (node count, vocab size, cluster count)"295 ),296 "mimeType": "application/json",297 },298 {299 "uri": "codeforge://corpus/node/{node_id}",300 "name": "Skill Node",301 "description": (302 "Full content of a specific skill node (free, no budget)"303 ),304 "mimeType": "application/json",305 },306 {307 "uri": "codeforge://tasks",308 "name": "Task Definitions",309 "description": (310 "Task definitions with briefs, budgets, targets, tools"311 ),312 "mimeType": "application/json",313 },314 {315 "uri": "codeforge://audit/{episode_id}",316 "name": "Audit Ledger",317 "description": (318 "Serialized audit ledger for a completed episode"319 ),320 "mimeType": "application/json",321 },322)323 324# -------------------------------------------------------------------325# Prompt text326# -------------------------------------------------------------------327 328_SYSTEM_PROMPT_TEXT = (329 "You are solving a CodeForge episode. Your code is graded by "330 "REAL tools (ruff, mypy --strict, pytest, import resolution) in "331 "a sandbox. Every import and attribute access is AST-grounded "332 "against the real Python runtime. Overconfidence is penalized "333 "via Brier scoring. Honest uncertainty about genuinely uncertain "334 "results is rewarded.\n\n"335 "Rules:\n"336 "- DO NOT fabricate library names or API signatures — "337 "the grounder catches them.\n"338 "- DO NOT submit stubs (pass, ..., NotImplementedError) — "339 "they score zero.\n"340 "- Use codeforge_query_kb to find patterns BEFORE writing code.\n"341 "- Use codeforge_interrogate to identify edge cases.\n"342 "- Budget is limited. Plan actions carefully.\n"343 "- If unsure of your confidence, set it low — "344 "the grader rewards honesty.\n"345)346 347_SESSION_ERR = "Invalid session_id: {sid!r}. Call codeforge_reset first."348 349 350# -------------------------------------------------------------------351# Helpers352# -------------------------------------------------------------------353 354 355def _obs_to_dict(obs: Any) -> dict[str, Any]:356 """Convert a CodeForgeObservation to a serializable dict."""357 result: dict[str, Any] = json.loads(obs.model_dump_json())358 return result359 360 361def _make_response(362 obs: Any,363 *,364 session_id: str | None = None,365 extra: dict[str, Any] | None = None,366) -> dict[str, Any]:367 """Build a versioned response dict from an observation."""368 result: dict[str, Any] = {"_codeforge_version": _VERSION}369 if session_id is not None:370 result["session_id"] = session_id371 result["observation"] = _obs_to_dict(obs)372 if extra:373 result.update(extra)374 budget = result["observation"].get("budget_remaining", 0)375 if isinstance(budget, int) and 0 < budget <= 2:376 result["budget_warning"] = (377 f"WARNING: {budget} budget remaining — plan carefully."378 )379 return result380 381 382def _session_error(sid: str) -> dict[str, Any]:383 """Return an isError response for a missing session."""384 return {385 "isError": True,386 "error": _SESSION_ERR.format(sid=sid),387 "_codeforge_version": _VERSION,388 }389 390 391def _require_session(392 server: CodeForgeMCPServer,393 arguments: dict[str, Any],394) -> tuple[str, CodeForgeEnvironment | None]:395 """Extract session_id and look up the environment."""396 sid: str = arguments.get("session_id", "")397 return sid, server._get_session(sid)398 399 400# -------------------------------------------------------------------401# CodeForgeMCPServer — embedded mode402# -------------------------------------------------------------------403 404 405class CodeForgeMCPServer:406 """MCP server wrapping CodeForgeEnvironment (SYSTEM_DESIGN §9).407 408 Embedded mode: imports CodeForgeEnvironment directly.409 Each tool call is routed to a session-keyed environment.410 411 **Universal LLM support:** The Ralph loop's synthesizer is configurable.412 Whichever LLM connects to this MCP server can provide its own config::413 414 # Ollama (local, free)415 server = CodeForgeMCPServer(llm_provider="ollama", llm_model="llama3")416 417 # OpenAI418 server = CodeForgeMCPServer(llm_provider="openai", llm_model="gpt-4o")419 420 # Anthropic421 server = CodeForgeMCPServer(llm_provider="anthropic", llm_model="claude-sonnet-4-20250514")422 423 # Any OpenAI-compatible (vLLM, LM Studio, Together, Groq)424 server = CodeForgeMCPServer(425 llm_provider="openai",426 llm_base_url="http://localhost:8000/v1",427 llm_model="my-model",428 )429 430 When no LLM config is provided, Ralph uses the deterministic431 StubSynthesizer (no API calls needed).432 """433 434 def __init__(435 self,436 *,437 corpus_path: Path | None = None,438 max_sessions: int = 10,439 llm_provider: str | None = None,440 llm_api_key: str | None = None,441 llm_base_url: str | None = None,442 llm_model: str | None = None,443 ) -> None:444 self._corpus_path = corpus_path445 self._sessions: dict[str, CodeForgeEnvironment] = {}446 self._max_sessions = max_sessions447 self._llm_provider = llm_provider448 self._llm_api_key = llm_api_key449 self._llm_base_url = llm_base_url450 self._llm_model = llm_model451 452 # -- Session management ------------------------------------------453 454 def _get_session(455 self,456 session_id: str,457 ) -> CodeForgeEnvironment | None:458 return self._sessions.get(session_id)459 460 def _create_session(461 self,462 ) -> tuple[str, CodeForgeEnvironment]:463 sid = uuid4().hex[:16]464 # Build synthesizer from LLM config (if provided)465 synth: Synthesizer | None = None466 if self._llm_provider:467 from codeforge.ralph.synthesizer import LLMSynthesizer468 469 synth = LLMSynthesizer(470 provider=self._llm_provider,471 api_key=self._llm_api_key,472 base_url=self._llm_base_url,473 model=self._llm_model,474 )475 476 env = CodeForgeEnvironment(477 corpus_path=self._corpus_path,478 synthesizer=synth,479 )480 if len(self._sessions) >= self._max_sessions:481 oldest = next(iter(self._sessions))482 del self._sessions[oldest]483 self._sessions[sid] = env484 return sid, env485 486 # -- Public: definitions -----------------------------------------487 488 def tool_definitions(self) -> list[dict[str, Any]]:489 """Return tool schemas matching SYSTEM_DESIGN §9.1."""490 return [dict(d) for d in _TOOL_DEFS]491 492 def resource_definitions(self) -> list[dict[str, str]]:493 """Return MCP resource definitions."""494 return [dict(r) for r in _RESOURCE_DEFS]495 496 def prompt_definitions(self) -> list[dict[str, Any]]:497 """Return MCP prompt definitions."""498 return [499 {500 "name": "codeforge_system",501 "description": (502 "System prompt injected at session start. "503 "Contains task rules, budget constraints, "504 "grading explanation."505 ),506 "arguments": [],507 },508 {509 "name": "codeforge_task_brief",510 "description": (511 "Dynamic prompt populated with the current "512 "task's brief, initial files, budget, target "513 "score, and tool config."514 ),515 "arguments": [516 {517 "name": "session_id",518 "description": _SID_DESC,519 "required": True,520 },521 ],522 },523 ]524 525 # -- Public: handle_tool -----------------------------------------526 527 def handle_tool(528 self,529 tool_name: str,530 arguments: dict[str, Any],531 ) -> dict[str, Any]:532 """Route tool call to handler, return result dict."""533 handler = _HANDLERS.get(tool_name)534 if handler is None:535 return {536 "isError": True,537 "error": f"Unknown tool: {tool_name!r}",538 "_codeforge_version": _VERSION,539 }540 return handler(self, arguments)541 542 # -- Public: resources -------------------------------------------543 544 def read_resource(545 self,546 uri: str,547 *,548 session_id: str | None = None,549 ) -> dict[str, Any]:550 """Read an MCP resource by URI."""551 if uri == "codeforge://corpus/stats":552 if session_id is None:553 return {554 "_codeforge_version": _VERSION,555 "error": "session_id required for corpus stats",556 }557 env = self._get_session(session_id)558 if env is None:559 return {"_codeforge_version": _VERSION, "error": "Invalid session_id"}560 idx = env._ensure_index()561 stats = idx.stats()562 cluster_count = len(idx.all_cluster_labels())563 return {564 "node_count": stats["node_count"],565 "vocab_size": stats["vocab_size"],566 "avg_doc_len": stats["avg_doc_len"],567 "cluster_count": cluster_count,568 }569 if uri.startswith("codeforge://corpus/node/"):570 node_id = uri.removeprefix("codeforge://corpus/node/")571 if session_id is None:572 return {573 "_codeforge_version": _VERSION,574 "error": "session_id required for node lookup",575 }576 env = self._get_session(session_id)577 if env is None:578 return {"_codeforge_version": _VERSION, "error": "Invalid session_id"}579 idx = env._ensure_index()580 for node in idx._nodes:581 if node.get("id") == node_id:582 return {"_codeforge_version": _VERSION, "node": node}583 return {"_codeforge_version": _VERSION, "error": f"Node {node_id!r} not found"}584 if uri == "codeforge://tasks":585 return {586 "_codeforge_version": _VERSION,587 "tasks": [588 {589 "id": t.task_id,590 "difficulty": t.task_level,591 "brief": t.brief,592 "target_score": t.target_score,593 "max_budget": t.max_budget,594 "tools": list(t.tools),595 }596 for t in TASKS597 ],598 }599 if uri.startswith("codeforge://audit/"):600 episode_id = uri.removeprefix("codeforge://audit/")601 if session_id is None:602 return {603 "_codeforge_version": _VERSION,604 "error": "session_id required for audit lookup",605 }606 env = self._get_session(session_id)607 if env is None:608 return {"_codeforge_version": _VERSION, "error": "Invalid session_id"}609 if env._ledger is not None:610 return {611 "_codeforge_version": _VERSION,612 "episode_id": episode_id,613 "audit": env._ledger.serialize(),614 }615 return {616 "_codeforge_version": _VERSION,617 "error": "No audit data for this session",618 }619 return {"_codeforge_version": _VERSION, "error": f"Unknown resource URI: {uri!r}"}620 621 # -- Public: prompts ---------------------------------------------622 623 def get_prompt(624 self,625 name: str,626 *,627 session_id: str | None = None,628 ) -> list[dict[str, str]]:629 """Return prompt messages for the given prompt name."""630 if name == "codeforge_system":631 return [632 {"role": "system", "content": _SYSTEM_PROMPT_TEXT},633 ]634 if name == "codeforge_task_brief":635 if session_id is None:636 return [637 {638 "role": "system",639 "content": (640 "Error: session_id required for "641 "task_brief prompt."642 ),643 },644 ]645 env = self._get_session(session_id)646 if env is None:647 return [648 {649 "role": "system",650 "content": "Error: invalid session_id.",651 },652 ]653 obs = env.state654 task = env._task655 target = task.target_score if task is not None else 0.0656 content = (657 f"## Task: {obs.task_id}\n"658 f"**Level:** {obs.task_level}\n"659 f"**Brief:** {obs.task_brief}\n"660 f"**Budget:** {obs.budget_remaining}\n"661 f"**Target score:** {target}\n\n"662 "### Initial files\n"663 )664 for fname, body in obs.initial_files.items():665 content += (666 f"\n**{fname}:**\n```python\n{body}\n```\n"667 )668 return [{"role": "system", "content": content}]669 return [670 {671 "role": "system",672 "content": f"Unknown prompt: {name!r}",673 },674 ]675 676 677# -------------------------------------------------------------------678# Tool handlers (private, keyed by tool name)679# -------------------------------------------------------------------680 681 682def _handle_reset(683 server: CodeForgeMCPServer,684 arguments: dict[str, Any],685) -> dict[str, Any]:686 task_level = arguments.get("task_level", "easy")687 sid, env = server._create_session()688 obs = env.reset(task_level=task_level)689 return _make_response(obs, session_id=sid)690 691 692def _handle_query_kb(693 server: CodeForgeMCPServer,694 arguments: dict[str, Any],695) -> dict[str, Any]:696 sid, env = _require_session(server, arguments)697 if env is None:698 return _session_error(sid)699 action = CodeForgeAction(700 action_type=CodeForgeActionType.QUERY_KB,701 claim=arguments.get("claim"),702 top_k=arguments.get("top_k", 5),703 required_tags=tuple(arguments.get("required_tags", ())),704 )705 obs = env.step(action)706 return _make_response(obs, session_id=sid)707 708 709def _handle_query_cluster(710 server: CodeForgeMCPServer,711 arguments: dict[str, Any],712) -> dict[str, Any]:713 sid, env = _require_session(server, arguments)714 if env is None:715 return _session_error(sid)716 action = CodeForgeAction(717 action_type=CodeForgeActionType.QUERY_CLUSTER,718 cluster_label=arguments.get("cluster_label"),719 top_k=arguments.get("top_k", 10),720 )721 obs = env.step(action)722 return _make_response(obs, session_id=sid)723 724 725def _handle_interrogate(726 server: CodeForgeMCPServer,727 arguments: dict[str, Any],728) -> dict[str, Any]:729 sid, env = _require_session(server, arguments)730 if env is None:731 return _session_error(sid)732 action = CodeForgeAction(733 action_type=CodeForgeActionType.INTERROGATE,734 )735 obs = env.step(action)736 return _make_response(obs, session_id=sid)737 738 739def _handle_run_ralph(740 server: CodeForgeMCPServer,741 arguments: dict[str, Any],742) -> dict[str, Any]:743 sid, env = _require_session(server, arguments)744 if env is None:745 return _session_error(sid)746 action = CodeForgeAction(747 action_type=CodeForgeActionType.RUN_RALPH,748 max_iters=arguments.get("max_iters", 3),749 )750 obs = env.step(action)751 return _make_response(obs, session_id=sid)752 753 754def _handle_submit(755 server: CodeForgeMCPServer,756 arguments: dict[str, Any],757) -> dict[str, Any]:758 sid, env = _require_session(server, arguments)759 if env is None:760 return _session_error(sid)761 action = CodeForgeAction(762 action_type=CodeForgeActionType.SUBMIT,763 files=arguments.get("files"),764 confidence=arguments.get("confidence"),765 )766 obs = env.step(action)767 return _make_response(obs, session_id=sid)768 769 770def _handle_get_audit(771 server: CodeForgeMCPServer,772 arguments: dict[str, Any],773) -> dict[str, Any]:774 sid, env = _require_session(server, arguments)775 if env is None:776 return _session_error(sid)777 action = CodeForgeAction(778 action_type=CodeForgeActionType.GET_AUDIT,779 target_run_id=arguments.get("target_run_id"),780 )781 obs = env.step(action)782 return _make_response(obs, session_id=sid)783 784 785def _handle_state(786 server: CodeForgeMCPServer,787 arguments: dict[str, Any],788) -> dict[str, Any]:789 sid, env = _require_session(server, arguments)790 if env is None:791 return _session_error(sid)792 obs = env.state793 return _make_response(obs, session_id=sid)794 795 796def _handle_list_clusters(797 server: CodeForgeMCPServer,798 arguments: dict[str, Any],799) -> dict[str, Any]:800 _sid, env = _require_session(server, arguments)801 if env is None:802 return {"_codeforge_version": _VERSION, "clusters": []}803 try:804 idx = env._ensure_index()805 except FileNotFoundError:806 return {"_codeforge_version": _VERSION, "clusters": []}807 labels = idx.all_cluster_labels()808 cluster_info: list[dict[str, Any]] = []809 for label in labels:810 cluster = idx.cluster_by_label(label)811 if cluster is not None:812 cluster_info.append({813 "label": cluster.label,814 "node_count": cluster.node_count,815 })816 return {"_codeforge_version": _VERSION, "clusters": cluster_info}817 818 819def _handle_list_tags(820 server: CodeForgeMCPServer,821 arguments: dict[str, Any],822) -> dict[str, Any]:823 _sid, env = _require_session(server, arguments)824 if env is None:825 return {"_codeforge_version": _VERSION, "tags": []}826 try:827 idx = env._ensure_index()828 except FileNotFoundError:829 return {"_codeforge_version": _VERSION, "tags": []}830 return {831 "_codeforge_version": _VERSION,832 "tags": sorted(idx.all_tags()),833 }834 835 836# Handler dispatch table837_HANDLERS: dict[str, _Handler] = {838 "codeforge_reset": _handle_reset,839 "codeforge_query_kb": _handle_query_kb,840 "codeforge_query_cluster": _handle_query_cluster,841 "codeforge_interrogate": _handle_interrogate,842 "codeforge_run_ralph": _handle_run_ralph,843 "codeforge_submit": _handle_submit,844 "codeforge_get_audit": _handle_get_audit,845 "codeforge_state": _handle_state,846 "codeforge_list_clusters": _handle_list_clusters,847 "codeforge_list_tags": _handle_list_tags,848}849