agent-foundry-runtime 0.4.0__py3-none-any.whl

agent_foundry/mcp_server.py

@@ -0,0 +1,191 @@
+"""Minimal MCP stdio server exposing plan/execute/list_skills/get_skill.
+
+Implements JSON-RPC 2.0 over stdin/stdout with no extra dependencies.
+"""
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+from typing import Any
+
+from .config import Config, LOG_DB_PATH
+from .executor import execute as execute_skill
+from .indexer import get_index_cached
+from .models import ExecuteRequest
+from .planner import rank_skills
+
+PROTOCOL_VERSION = "2025-06-18"
+SERVER_VERSION = "0.4.0"
+
+TOOLS = [
+    {
+        "name": "plan",
+        "description": "Rank Agent Foundry skills for a prompt.",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "prompt": {"type": "string"},
+                "max_results": {"type": "integer", "default": 5},
+            },
+            "required": ["prompt"],
+        },
+    },
+    {
+        "name": "execute",
+        "description": "Run one skill against a prompt via the configured LLM.",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "skill_id": {"type": "string"},
+                "prompt": {"type": "string"},
+                "model": {"type": "string"},
+            },
+            "required": ["skill_id", "prompt"],
+        },
+    },
+    {
+        "name": "list_skills",
+        "description": "List indexed skills with category and description.",
+        "inputSchema": {"type": "object", "properties": {}},
+    },
+    {
+        "name": "get_skill",
+        "description": "Return the full SKILL.md body for one skill.",
+        "inputSchema": {
+            "type": "object",
+            "properties": {"skill_id": {"type": "string"}},
+            "required": ["skill_id"],
+        },
+    },
+]
+
+
+class McpServer:
+    def __init__(self, cfg: Config | None = None, db_path: Path | None = None):
+        self.cfg = cfg or Config.load()
+        self.db_path = db_path or LOG_DB_PATH
+
+    def _index(self):
+        idx = get_index_cached(Path(self.cfg.core.index_path).expanduser(),
+                               ttl_seconds=self.cfg.core.index_cache_ttl_seconds)
+        if idx is None:
+            raise RuntimeError("No skill index. Run `agent-foundry index` first.")
+        return idx
+
+    def call_tool(self, name: str, arguments: dict[str, Any]) -> dict[str, Any]:
+        if name == "plan":
+            idx = self._index()
+            prompt = str(arguments.get("prompt", ""))
+            max_results = int(arguments.get("max_results", 5))
+            ranked = rank_skills(prompt, idx, self.cfg.planner,
+                                 weights=_weights(self.db_path),
+                                 rerank_model=(self.cfg.core.default_llm
+                                               if self.cfg.planner.rerank_enabled else None))
+            result = {
+                "prompt": prompt,
+                "results": [r.model_dump() for r in ranked[:max_results]],
+                "total_available": len(idx.skills),
+            }
+            return _text(result)
+
+        if name == "execute":
+            idx = self._index()
+            req = ExecuteRequest(
+                skill_id=str(arguments["skill_id"]),
+                prompt=str(arguments["prompt"]),
+                model=arguments.get("model"),
+            )
+            resp = execute_skill(
+                req, model=req.model or self.cfg.core.default_llm,
+                skill_paths=idx.skill_paths,
+                db_path=self.db_path,
+                skill_hashes={s.id: s.sha256 for s in idx.skills},
+                timeout_seconds=self.cfg.core.executor_timeout_seconds,
+                max_tokens=self.cfg.core.executor_max_tokens,
+                retries=self.cfg.core.executor_retries,
+            )
+            return _text(resp.model_dump())
+
+        if name == "list_skills":
+            idx = self._index()
+            skills = [
+                {"id": s.id, "name": s.name, "category": s.category,
+                 "description": s.description}
+                for s in idx.skills if not s.is_fallback
+            ]
+            return _text({"skills": skills, "count": len(skills)})
+
+        if name == "get_skill":
+            idx = self._index()
+            skill_id = str(arguments["skill_id"])
+            path = idx.skill_paths.get(skill_id)
+            if not path or path.startswith("__virtual__:"):
+                raise RuntimeError(f"unknown skill: {skill_id}")
+            return _text({"skill_id": skill_id,
+                          "content": Path(path).read_text(encoding="utf-8")})
+
+        raise RuntimeError(f"unknown tool: {name}")
+
+    def handle(self, msg: dict[str, Any]) -> dict[str, Any] | None:
+        method = msg.get("method")
+        msg_id = msg.get("id")
+
+        if method == "initialize":
+            return _result(msg_id, {
+                "protocolVersion": PROTOCOL_VERSION,
+                "capabilities": {"tools": {}},
+                "serverInfo": {"name": "agent-foundry", "version": SERVER_VERSION},
+            })
+        if method in ("notifications/initialized", "initialized"):
+            return None
+        if method == "ping":
+            return _result(msg_id, {})
+        if method == "tools/list":
+            return _result(msg_id, {"tools": TOOLS})
+        if method == "tools/call":
+            params = msg.get("params") or {}
+            try:
+                payload = self.call_tool(params.get("name", ""), params.get("arguments") or {})
+                return _result(msg_id, payload)
+            except Exception as exc:
+                return _result(msg_id, {"content": [{"type": "text", "text": str(exc)}],
+                                        "isError": True})
+        return {
+            "jsonrpc": "2.0",
+            "id": msg_id,
+            "error": {"code": -32601, "message": f"method not found: {method}"},
+        }
+
+
+def _weights(db_path: Path) -> dict[str, float]:
+    try:
+        from .logging_db import get_skill_weights
+        return get_skill_weights(db_path)
+    except Exception:
+        return {}
+
+
+def _text(value: Any) -> dict[str, Any]:
+    return {"content": [{"type": "text", "text": json.dumps(value, ensure_ascii=False)}]}
+
+
+def _result(msg_id: Any, result: dict[str, Any]) -> dict[str, Any]:
+    return {"jsonrpc": "2.0", "id": msg_id, "result": result}
+
+
+def run_stdio(cfg: Config | None = None) -> None:
+    """Serve MCP over stdin/stdout until EOF."""
+    server = McpServer(cfg)
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            msg = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+        response = server.handle(msg)
+        if response is not None:
+            sys.stdout.write(json.dumps(response, ensure_ascii=False) + "\n")
+            sys.stdout.flush()

agent_foundry/models.py

@@ -0,0 +1,166 @@
+"""Pydantic data models for Agent Foundry."""
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Literal
+from pydantic import BaseModel, Field, field_validator, model_validator
+
+
+class TokenEstimate(BaseModel):
+    """Cost estimate in tokens. Total is auto-computed if not provided."""
+    input: int = Field(ge=0, description="Input tokens (skill body + prompt)")
+    output: int = Field(ge=0, description="Estimated output tokens")
+    total: int = Field(default=0, ge=0, description="input + output (computed)")
+
+    @model_validator(mode="before")
+    @classmethod
+    def _compute_total(cls, values):
+        if isinstance(values, dict):
+            inp = values.get("input", 0) or 0
+            out = values.get("output", 0) or 0
+            values = dict(values)
+            values.setdefault("total", inp + out)
+        return values
+
+
+class SkillManifest(BaseModel):
+    """One skill, as stored in skills/index.json."""
+    id: str = Field(pattern=r"^[a-z0-9-]+$")
+    name: str
+    description: str = Field(max_length=500)
+    trigger_patterns: list[str] = Field(default_factory=list)
+    estimated_token_cost: dict = Field(default_factory=lambda: {"input": 200, "output": 300})
+    estimated_time_seconds: int = 3
+    version: str = "0.1.0"
+    license: str = "MIT"
+    category: Literal["core", "optional"] = "core"
+    location: str
+    tags: list[str] = Field(default_factory=list)
+    is_fallback: bool = False
+    sha256: str = ""
+
+    def total_cost(self) -> int:
+        c = self.estimated_token_cost or {}
+        return (c.get("input", 0) or 0) + (c.get("output", 0) or 0)
+
+
+class SkillIndex(BaseModel):
+    """The full skill index."""
+    version: str = "1"
+    generated_at: datetime
+    skills: list[SkillManifest]
+    skill_paths: dict[str, str] = Field(default_factory=dict)
+
+
+# ===== Planner =====
+
+class PlanRequest(BaseModel):
+    prompt: str
+    max_results: int = 5
+    graph_path: str | None = None
+
+
+class PlanResult(BaseModel):
+    """One ranked skill in a plan."""
+    skill_id: str
+    name: str
+    description: str
+    score: float
+    matched_patterns: list[str] = Field(default_factory=list)
+    estimated_cost: TokenEstimate
+    location: str | None = None
+
+
+class PlanResponse(BaseModel):
+    """Plan endpoint reply: ranked list of skills."""
+    prompt: str
+    results: list[PlanResult]
+    total_available: int
+
+
+# ===== Executor =====
+
+class ExecuteRequest(BaseModel):
+    skill_id: str
+    prompt: str
+    force: bool = False
+    model: str | None = None
+
+
+class ExecuteResponse(BaseModel):
+    """Execute endpoint reply."""
+    skill_id: str
+    output: str = ""
+    tokens_used: int = 0
+    duration_seconds: float = 0.0
+    success: bool = True
+    error: str | None = None
+    requires_confirmation: bool = False
+    estimated_cost: TokenEstimate | None = None
+    was_fallback: bool = False
+    cost_usd: float = 0.0
+
+
+# ===== Loop =====
+
+class LoopRequest(BaseModel):
+    prompt: str
+    force: bool = False
+    model: str | None = None
+    budget: int | None = None
+    judge: bool = False  # if true, score output via af-critic
+
+
+class LoopResponse(BaseModel):
+    """Loop endpoint reply."""
+    plan: list[PlanResult] = Field(default_factory=list)
+    skill_id: str | None = None
+    output: str = ""
+    tokens_used: int = 0
+    duration_seconds: float = 0.0
+    success: bool = True
+    error: str | None = None
+    requires_confirmation: bool = False
+    estimated_cost: TokenEstimate | None = None
+    was_fallback: bool = False
+    judge_score: dict | None = None   # correctness/slop/scope/verdict if --judge
+    judged: bool = False
+    cost_usd: float = 0.0
+    session_id: int | None = None
+
+
+class LoopResumeRequest(BaseModel):
+    session_id: int
+    prompt: str
+    model: str | None = None
+
+
+# ===== Index =====
+
+class StatsResponse(BaseModel):
+    total_executions: int = 0
+    successful: int = 0
+    failed: int = 0
+    fallback_count: int = 0
+    avg_tokens: int = 0
+    avg_duration: float = 0.0
+    thumbs_up: int = 0
+    thumbs_down: int = 0
+    instincts_count: int = 0
+
+
+class FeedbackRequest(BaseModel):
+    execution_id: int
+    feedback: int = Field(..., description="1 = thumbs_up, -1 = thumbs_down")
+
+
+class LearnResponse(BaseModel):
+    new_instincts: int = 0
+    updated_instincts: int = 0
+    patterns_found: list[dict] = Field(default_factory=list)
+
+
+class IndexResponse(BaseModel):
+    total_skills: int
+    took_ms: float
+    regenerated: bool = True

agent_foundry/planner.py

@@ -0,0 +1,288 @@
+"""Planner: rank skills by prompt relevance + cost penalty + graph relationships."""
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from pathlib import Path
+from typing import Iterable
+
+from .config import PlannerConfig
+from .indexer import GENERIC_REASONING_ID, get_index_cached
+from .models import (
+    PlanRequest,
+    PlanResponse,
+    PlanResult,
+    SkillIndex,
+    SkillManifest,
+    TokenEstimate,
+)
+
+
+_STOPWORDS = frozenset({
+    "the", "and", "for", "that", "this", "with", "from", "how", "what",
+    "when", "why", "where", "which", "your", "have", "will", "can", "are",
+    "not", "but", "all", "was", "out", "use", "into", "any", "its", "our",
+})
+
+
+def _keyword_overlap(prompt: str, skill: SkillManifest, limit: int = 3) -> int:
+    """Cheap lexical fallback when trigger regexes miss. Count shared content words."""
+    p = {w for w in re.findall(r"[a-z]{3,}", prompt.lower()) if w not in _STOPWORDS}
+    d = {w for w in re.findall(r"[a-z]{3,}", (skill.description or "").lower())
+         if w not in _STOPWORDS}
+    return min(limit, len(p & d))
+
+
+def _cost_penalty(total_cost: int, divisor: float) -> float:
+    return 1.0 / (1.0 + total_cost / divisor)
+
+
+def _match_patterns(prompt: str, patterns: Iterable[str]) -> list[str]:
+    matched: list[str] = []
+    p_lower = prompt.lower()
+    for pat in patterns or []:
+        if not pat:
+            continue
+        try:
+            if re.search(pat, prompt, re.IGNORECASE | re.DOTALL):
+                matched.append(pat)
+                continue
+        except re.error:
+            pass
+        raw = pat.replace(r"\b", "").replace("\\\\", "")
+        if raw and raw.lower() in p_lower:
+            matched.append(pat)
+    return matched
+
+
+def _name_in_prompt(prompt: str, skill: SkillManifest) -> bool:
+    p_lower = prompt.lower()
+    name_lower = (skill.name or "").lower()
+    if name_lower and name_lower in p_lower:
+        return True
+    desc_lower = (skill.description or "").lower()
+    if desc_lower and desc_lower[:60] in p_lower:
+        return True
+    return False
+
+
+_GRAPH_CACHE: dict[str, tuple[float, dict[str, list[str]]]] = {}
+_SEM_CACHE: dict = {"key": None, "model": None, "ids": [], "vectors": []}
+
+
+def _graph_adjacency(graph_path: Path) -> dict[str, list[str]]:
+    """Parse the graph once per mtime and return an undirected adjacency map."""
+    try:
+        mtime = graph_path.stat().st_mtime
+    except OSError:
+        return {}
+    cached = _GRAPH_CACHE.get(str(graph_path))
+    if cached and cached[0] == mtime:
+        return cached[1]
+
+    try:
+        data = json.loads(graph_path.read_text())
+        edges = data.get("links", data.get("edges", []))
+        adj: dict[str, list[str]] = {}
+        for e in edges:
+            src = e.get("source", "")
+            tgt = e.get("target", "")
+            if not src or not tgt:
+                continue
+            adj.setdefault(src, []).append(tgt)
+            adj.setdefault(tgt, []).append(src)
+    except Exception:
+        return {}
+    _GRAPH_CACHE[str(graph_path)] = (mtime, adj)
+    return adj
+
+
+def _graph_related(skill_id: str, graph_path: Path | None,
+                   depth: int = 1) -> list[str]:
+    """Return skill IDs connected in the knowledge graph (max 3 per depth)."""
+    if not graph_path or not graph_path.exists():
+        return []
+    adj = _graph_adjacency(graph_path)
+    related: list[str] = []
+    frontier = [skill_id]
+    seen = {skill_id}
+    for _ in range(max(1, depth)):
+        nxt: list[str] = []
+        for node in frontier:
+            for neigh in adj.get(node, []):
+                if neigh not in seen:
+                    seen.add(neigh)
+                    nxt.append(neigh)
+        related.extend(nxt[:3])
+        frontier = nxt
+    return related
+
+
+def _semantic_scores(prompt: str, skills: list[SkillManifest], model_name: str) -> dict[str, float]:
+    """Cosine similarity of prompt vs skill descriptions. Empty if unavailable."""
+    try:
+        from sentence_transformers import SentenceTransformer  # type: ignore
+    except Exception:
+        return {}
+
+    signature = hashlib.sha1(
+        (model_name + "|" + "|".join(f"{s.id}:{s.description}" for s in skills)).encode()
+    ).hexdigest()
+    if _SEM_CACHE["key"] != signature:
+        try:
+            model = _SEM_CACHE["model"]
+            if model is None or _SEM_CACHE.get("model_name") != model_name:
+                model = SentenceTransformer(model_name)
+            vectors = model.encode([s.description for s in skills], normalize_embeddings=True)
+            _SEM_CACHE.update({
+                "key": signature,
+                "model": model,
+                "model_name": model_name,
+                "ids": [s.id for s in skills],
+                "vectors": [list(map(float, v)) for v in vectors],
+            })
+        except Exception:
+            return {}
+
+    try:
+        q = _SEM_CACHE["model"].encode([prompt], normalize_embeddings=True)[0]
+        q = list(map(float, q))
+    except Exception:
+        return {}
+
+    out: dict[str, float] = {}
+    for sid, vec in zip(_SEM_CACHE["ids"], _SEM_CACHE["vectors"]):
+        out[sid] = sum(a * b for a, b in zip(q, vec))  # vectors are normalized
+    return out
+
+
+def _llm_rerank(prompt: str, results: list[tuple[float, PlanResult]],
+                model: str, top_k: int) -> list[tuple[float, PlanResult]]:
+    """Optional LLM ordering of the top candidates. Returns input on any failure."""
+    try:
+        import litellm
+
+        candidates = [
+            {"id": r.skill_id, "description": (r.description or "")[:160]}
+            for _, r in results[:top_k]
+        ]
+        messages = [
+            {"role": "system",
+             "content": 'Reorder the candidate skills for the prompt. Return only JSON: {"order": ["id", ...]}.'},
+            {"role": "user",
+             "content": json.dumps({"prompt": prompt, "candidates": candidates})},
+        ]
+        resp = litellm.completion(
+            model=model, messages=messages, timeout=30,
+            response_format={"type": "json_object"},
+        )
+        data = json.loads(resp.choices[0].message.content or "{}")
+        order = data.get("order") or []
+        rank = {sid: i for i, sid in enumerate(order)}
+        head = sorted(results[:top_k], key=lambda pair: rank.get(pair[1].skill_id, 999))
+        return head + results[top_k:]
+    except Exception:
+        return results
+
+
+def rank_skills(prompt: str, idx: SkillIndex, cfg: PlannerConfig,
+                graph_path: Path | None = None,
+                weights: dict[str, float] | None = None,
+                rerank_model: str | None = None) -> list[PlanResult]:
+    """Return skills ranked by score (desc). Graph-aware: boosts skills related to the top match."""
+    results: list[tuple[float, PlanResult]] = []
+
+    semantic: dict[str, float] = {}
+    if cfg.semantic_enabled:
+        semantic = _semantic_scores(prompt, [s for s in idx.skills if not s.is_fallback],
+                                    cfg.semantic_model)
+
+    for skill in idx.skills:
+        if skill.is_fallback:
+            continue
+
+        matched = _match_patterns(prompt, skill.trigger_patterns)
+        score = len(matched) * cfg.pattern_match_weight
+
+        if _name_in_prompt(prompt, skill):
+            score += cfg.name_in_prompt_boost
+
+        overlap = _keyword_overlap(prompt, skill)
+        if overlap:
+            score += overlap * 0.04
+
+        if semantic:
+            score += max(0.0, semantic.get(skill.id, 0.0)) * cfg.semantic_weight
+
+        if score <= 0:
+            continue
+
+        total_cost = skill.total_cost()
+        score *= _cost_penalty(total_cost, cfg.cost_penalty_divisor)
+        if weights:
+            score *= max(0.1, float(weights.get(skill.id, 1.0)))
+
+        cost = TokenEstimate(
+            input=int(skill.estimated_token_cost.get("input", 0)),
+            output=int(skill.estimated_token_cost.get("output", 0)),
+        )
+        results.append((
+            score,
+            PlanResult(
+                skill_id=skill.id,
+                name=skill.name,
+                description=skill.description,
+                score=round(score, 4),
+                matched_patterns=matched,
+                estimated_cost=cost,
+                location=skill.location,
+            ),
+        ))
+
+    results.sort(key=lambda x: x[0], reverse=True)
+
+    # Graph-aware boost: top match boosts its 2-hop neighborhood.
+    if graph_path and results:
+        top_id = results[0][1].skill_id
+        related = _graph_related(top_id, graph_path, depth=2)
+        if related:
+            boost = {sid: (1.15 if i < 3 else 1.08) for i, sid in enumerate(related)}
+            for i, (s, r) in enumerate(results):
+                factor = boost.get(r.skill_id)
+                if factor:
+                    results[i] = (s * factor, r)
+
+    results.sort(key=lambda x: x[0], reverse=True)
+    if cfg.rerank_enabled and rerank_model and len(results) > 1:
+        results = _llm_rerank(prompt, results, rerank_model, cfg.rerank_top_k)
+    return [r for _, r in results[:cfg.max_results]]
+
+
+def _safe_graph_path(raw: str | None) -> Path | None:
+    """Only allow graph JSON inside the repo root or current working directory."""
+    if not raw:
+        return None
+    try:
+        path = Path(raw).expanduser().resolve()
+    except OSError:
+        return None
+    if path.suffix.lower() != ".json" or not path.is_file():
+        return None
+    roots = [Path.cwd().resolve(), Path(__file__).resolve().parent.parent]
+    if not any(path.is_relative_to(root) for root in roots):
+        return None
+    return path
+
+
+def plan(request: PlanRequest, idx: SkillIndex, cfg: PlannerConfig,
+         weights: dict[str, float] | None = None,
+         rerank_model: str | None = None) -> PlanResponse:
+    graph_path = _safe_graph_path(request.graph_path)
+    ranked = rank_skills(request.prompt, idx, cfg, graph_path=graph_path,
+                         weights=weights, rerank_model=rerank_model)
+    return PlanResponse(
+        prompt=request.prompt,
+        results=ranked,
+        total_available=len(idx.skills),
+    )