133 lines
4.6 KiB
Python
133 lines
4.6 KiB
Python
"""LLM interaction, tool dispatch, and memory management."""
|
|
|
|
import os
|
|
import re
|
|
import json
|
|
import urllib.request
|
|
import urllib.error
|
|
|
|
|
|
def log(msg):
|
|
print(f"[tools] {msg}", flush=True)
|
|
|
|
|
|
def set_logger(fn):
|
|
global log
|
|
log = fn
|
|
|
|
|
|
# ─── Memory ──────────────────────────────────────────────────────────
|
|
|
|
def load_memory(workspace):
|
|
"""Load all memory files from workspace."""
|
|
memory = ""
|
|
try:
|
|
with open(f"{workspace}/MEMORY.md") as f:
|
|
memory = f.read().strip()
|
|
mem_dir = f"{workspace}/memory"
|
|
if os.path.isdir(mem_dir):
|
|
for fname in sorted(os.listdir(mem_dir)):
|
|
if fname.endswith(".md"):
|
|
try:
|
|
with open(f"{mem_dir}/{fname}") as f:
|
|
topic = fname.replace(".md", "")
|
|
memory += f"\n\n## {topic}\n{f.read().strip()}"
|
|
except Exception:
|
|
pass
|
|
except FileNotFoundError:
|
|
pass
|
|
return memory
|
|
|
|
|
|
# ─── Tool Call Parsing ───────────────────────────────────────────────
|
|
|
|
def try_parse_tool_call(text):
|
|
"""Parse text-based tool calls (model dumps JSON as text)."""
|
|
text = re.sub(r"</?tool_call>", "", text).strip()
|
|
for start in range(len(text)):
|
|
if text[start] == "{":
|
|
for end in range(len(text), start, -1):
|
|
if text[end - 1] == "}":
|
|
try:
|
|
obj = json.loads(text[start:end])
|
|
name = obj.get("name")
|
|
args = obj.get("arguments", {})
|
|
if name and isinstance(args, dict):
|
|
return (name, args)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
return None
|
|
|
|
|
|
# ─── LLM Interaction ────────────────────────────────────────────────
|
|
|
|
def ollama_request(ollama_url, payload):
|
|
data = json.dumps(payload).encode("utf-8")
|
|
req = urllib.request.Request(
|
|
f"{ollama_url}/api/chat",
|
|
data=data,
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
with urllib.request.urlopen(req, timeout=120) as resp:
|
|
return json.loads(resp.read(2_000_000))
|
|
|
|
|
|
def query_ollama(messages, runtime, tools, skill_scripts, dispatch_fn, ollama_url, max_rounds):
|
|
"""Call Ollama chat API with skill-based tool support."""
|
|
payload = {
|
|
"model": runtime["model"],
|
|
"messages": messages,
|
|
"stream": False,
|
|
"options": {"num_predict": 512},
|
|
}
|
|
|
|
if tools:
|
|
payload["tools"] = tools
|
|
|
|
for round_num in range(max_rounds):
|
|
remaining = max_rounds - round_num
|
|
try:
|
|
data = ollama_request(ollama_url, payload)
|
|
except (urllib.error.URLError, TimeoutError) as e:
|
|
return f"[error: {e}]"
|
|
|
|
msg = data.get("message", {})
|
|
|
|
# Structured tool calls
|
|
tool_calls = msg.get("tool_calls")
|
|
if tool_calls:
|
|
messages.append(msg)
|
|
for tc in tool_calls:
|
|
fn = tc.get("function", {})
|
|
result = dispatch_fn(
|
|
fn.get("name", ""),
|
|
fn.get("arguments", {}),
|
|
round_num + 1,
|
|
)
|
|
if remaining <= 2:
|
|
result += f"\n[warning: {remaining - 1} tool rounds remaining — wrap up]"
|
|
messages.append({"role": "tool", "content": result})
|
|
payload["messages"] = messages
|
|
continue
|
|
|
|
# Text-based tool calls
|
|
content = msg.get("content", "").strip()
|
|
parsed_tool = try_parse_tool_call(content)
|
|
if parsed_tool:
|
|
fn_name, fn_args = parsed_tool
|
|
if fn_name in skill_scripts:
|
|
messages.append({"role": "assistant", "content": content})
|
|
result = dispatch_fn(fn_name, fn_args, round_num + 1)
|
|
if remaining <= 2:
|
|
result += f"\n[warning: {remaining - 1} tool rounds remaining — wrap up]"
|
|
messages.append({
|
|
"role": "user",
|
|
"content": f"Tool result:\n{result}\n\nNow respond to the user based on this result.",
|
|
})
|
|
payload["messages"] = messages
|
|
continue
|
|
|
|
return content
|
|
|
|
return "[max tool rounds reached]"
|