Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python | |
| """LLM backend for the DIY Creator multi-agent team (#17): local Ollama (default) or Claude. | |
| Two entry points: | |
| chat(system, user) -> free text (the writer / reviser produce markdown guides) | |
| json_call(system, user) -> parsed JSON (the researcher's brief, the critic's verdict) | |
| Each call takes an optional `model` so different ROLES can run on different models — which is the | |
| whole point of the critic: it can be a *different* model from the writer, so the review is | |
| independent (the "diamond"), not the author grading its own work. Errors log the exception TYPE only. | |
| """ | |
| import json | |
| import os | |
| import urllib.request | |
| BACKEND = os.environ.get("DIY_LLM_BACKEND", "ollama").lower() | |
| OLLAMA_URL = os.environ.get("DIY_OLLAMA_URL", "http://127.0.0.1:11434").rstrip("/") | |
| OLLAMA_MODEL = os.environ.get("DIY_OLLAMA_MODEL", "qwen2.5:7b") | |
| # The critic defaults to the same model, but set DIY_CRITIC_MODEL to a *different* one for a truly | |
| # independent review (e.g. a different family, or a bigger model where the quality bar matters). | |
| CRITIC_MODEL = os.environ.get("DIY_CRITIC_MODEL", "").strip() or OLLAMA_MODEL | |
| ANTHROPIC_KEY = (os.environ.get("ANTHROPIC_API_KEY") or "").strip() or None | |
| ANTHROPIC_MODEL = os.environ.get("ANTHROPIC_MODEL", "claude-haiku-4-5") | |
| ANTHROPIC_CRITIC = os.environ.get("ANTHROPIC_CRITIC_MODEL", "").strip() or ANTHROPIC_MODEL | |
| _USE_CLAUDE = BACKEND == "claude" and ANTHROPIC_KEY | |
| def _ollama(system, user, model, temperature, max_tokens, as_json): | |
| payload = {"model": model, "stream": False, | |
| "messages": [{"role": "system", "content": system}, {"role": "user", "content": user}], | |
| "options": {"temperature": temperature, "num_predict": max_tokens}} | |
| if as_json: | |
| payload["format"] = "json" # Ollama constrains output to valid JSON | |
| body = json.dumps(payload).encode() | |
| req = urllib.request.Request(OLLAMA_URL + "/api/chat", data=body, | |
| headers={"Content-Type": "application/json"}) | |
| return json.loads(urllib.request.urlopen(req, timeout=180).read())["message"]["content"].strip() | |
| def _claude(system, user, model, temperature, max_tokens, as_json): | |
| if as_json: | |
| system = system + "\nReply with ONLY valid JSON — no prose, no code fences." | |
| body = json.dumps({"model": model, "max_tokens": max_tokens, "temperature": temperature, | |
| "system": system, "messages": [{"role": "user", "content": user}]}).encode() | |
| req = urllib.request.Request("https://api.anthropic.com/v1/messages", data=body, headers={ | |
| "x-api-key": ANTHROPIC_KEY, "anthropic-version": "2023-06-01", "content-type": "application/json"}) | |
| data = json.loads(urllib.request.urlopen(req, timeout=180).read()) | |
| return "".join(b.get("text", "") for b in data.get("content", [])).strip() | |
| def _resolve(model, role): | |
| """Pick the concrete model id for a role given the active backend.""" | |
| if _USE_CLAUDE: | |
| return ANTHROPIC_CRITIC if role == "critic" else ANTHROPIC_MODEL | |
| if role == "critic": | |
| return CRITIC_MODEL | |
| return model or OLLAMA_MODEL | |
| def chat(system: str, user: str, role: str = "", model: str = None, | |
| temperature: float = 0.5, max_tokens: int = 1400) -> str: | |
| m = _resolve(model, role) | |
| fn = _claude if _USE_CLAUDE else _ollama | |
| return fn(system, user, m, temperature, max_tokens, as_json=False) | |
| def json_call(system: str, user: str, role: str = "", model: str = None, | |
| temperature: float = 0.2, max_tokens: int = 1200): | |
| m = _resolve(model, role) | |
| fn = _claude if _USE_CLAUDE else _ollama | |
| raw = fn(system, user, m, temperature, max_tokens, as_json=True) | |
| return _parse_json(raw) | |
| def _parse_json(raw: str): | |
| s = (raw or "").strip() | |
| if s.startswith("```"): | |
| s = s.split("```", 2)[1].lstrip("json").strip() if s.count("```") >= 2 else s.strip("`") | |
| try: | |
| return json.loads(s) | |
| except Exception: | |
| pass | |
| # salvage the first {...} or [...] block | |
| for open_c, close_c in (("{", "}"), ("[", "]")): | |
| i, j = s.find(open_c), s.rfind(close_c) | |
| if 0 <= i < j: | |
| try: | |
| return json.loads(s[i:j + 1]) | |
| except Exception: | |
| continue | |
| return None | |
| def backend_label() -> str: | |
| if _USE_CLAUDE: | |
| extra = f" · critic: {ANTHROPIC_CRITIC}" if ANTHROPIC_CRITIC != ANTHROPIC_MODEL else "" | |
| return f"Claude ({ANTHROPIC_MODEL}){extra}" | |
| extra = f" · critic: {CRITIC_MODEL}" if CRITIC_MODEL != OLLAMA_MODEL else "" | |
| return f"Ollama ({OLLAMA_MODEL} @ {OLLAMA_URL}){extra}" | |