Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python | |
| """LLM backends for the Web Research Agent (#13): local Ollama (default, free/private) or Claude. | |
| On-demand only — these are plain function calls, nothing runs in the background. Errors log the | |
| exception TYPE (never repr) so a key can't leak into logs. | |
| """ | |
| import json | |
| import os | |
| import urllib.request | |
| BACKEND = os.environ.get("RESEARCH_LLM_BACKEND", "ollama").lower() | |
| OLLAMA_URL = os.environ.get("RESEARCH_OLLAMA_URL", "http://127.0.0.1:11434") # set to your Spark/fleet host | |
| OLLAMA_MODEL = os.environ.get("RESEARCH_OLLAMA_MODEL", "qwen2.5:7b") | |
| ANTHROPIC_KEY = (os.environ.get("ANTHROPIC_API_KEY") or "").strip() or None | |
| ANTHROPIC_MODEL = os.environ.get("ANTHROPIC_MODEL", "claude-haiku-4-5") | |
| def _ollama(system, user, temperature, max_tokens, model=None): | |
| body = json.dumps({ | |
| "model": model or OLLAMA_MODEL, | |
| "messages": [{"role": "system", "content": system}, {"role": "user", "content": user}], | |
| "stream": False, "options": {"temperature": temperature, "num_predict": max_tokens}, | |
| }).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, temperature, max_tokens, model=None): | |
| body = json.dumps({ | |
| "model": model or ANTHROPIC_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 chat(system: str, user: str, temperature: float = 0.3, max_tokens: int = 900, model: str = None) -> str: | |
| """One completion. Routes to Claude if RESEARCH_LLM_BACKEND=claude (and a key is set), else Ollama.""" | |
| fn = _claude if (BACKEND == "claude" and ANTHROPIC_KEY) else _ollama | |
| return fn(system, user, temperature, max_tokens, model) | |
| def chat_json(system: str, user: str, temperature: float = 0.0, max_tokens: int = 500, model: str = None): | |
| """Completion parsed as JSON. Tolerant: strips code fences and pulls the first {...} or [...] block.""" | |
| raw = chat(system + "\nReply with ONLY valid JSON, no prose, no code fences.", | |
| user, temperature, max_tokens, model) | |
| return _parse_json(raw) | |
| def _parse_json(raw: str): | |
| s = raw.strip() | |
| if s.startswith("```"): | |
| s = s.split("```", 2)[1] if s.count("```") >= 2 else s.strip("`") | |
| if s.lower().startswith("json"): | |
| s = s[4:] | |
| s = s.strip() | |
| try: | |
| return json.loads(s) | |
| except Exception: | |
| # pull first array/object | |
| for op, cl in (("[", "]"), ("{", "}")): | |
| i, j = s.find(op), s.rfind(cl) | |
| if 0 <= i < j: | |
| try: | |
| return json.loads(s[i:j + 1]) | |
| except Exception: | |
| pass | |
| return None | |
| def backend_label() -> str: | |
| return "Claude" if (BACKEND == "claude" and ANTHROPIC_KEY) else f"Ollama ({OLLAMA_MODEL} @ {OLLAMA_URL})" | |
| if __name__ == "__main__": | |
| print("backend:", backend_label()) | |
| print(chat("You are terse.", "Say 'research llm ok' and nothing else.", max_tokens=20)) | |
| print("json:", chat_json("You split questions.", "Give 2 search queries for: impact of SLMs on enterprise cost. As a JSON array of strings.")) | |