fix: strengthen memory tool prompt — Gemma 4B needs explicit examples
The 4B model was responding conversationally instead of calling memory_save/ memory_recall. Added imperative language (MUST call), concrete examples of trigger phrases, and explicit instructions to never skip the tool call. Verified: model now reliably generates tool_calls for save/recall/list.
This commit is contained in:
Shane
@@ -0,0 +1,240 @@
|
||||
"""Background task worker — runs an autonomous LLM loop with tools.
|
||||
|
||||
Each dispatched task gets its own TaskWorker instance running in a separate
|
||||
asyncio task. The worker calls Gemma iteratively, executing tool calls and
|
||||
logging each step to the registry so the UI can display live progress.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from datetime import datetime, timezone
|
||||
|
||||
logger = logging.getLogger("voice-agent.tasks")
|
||||
|
||||
|
||||
class TaskWorker:
|
||||
"""Runs a background research/lookup task using the LLM + tools."""
|
||||
|
||||
def __init__(self, task_id: str, description: str):
|
||||
self.task_id = task_id
|
||||
self.description = description
|
||||
self._cancelled = False
|
||||
|
||||
async def run(self):
|
||||
from openai import AsyncOpenAI
|
||||
from agent.task_registry import registry as reg
|
||||
|
||||
base_url = os.environ.get("GEMMA_BASE_URL", "http://192.168.86.2:8023/v1")
|
||||
model = os.environ.get("GEMMA_MODEL", "gemma-4-e4b")
|
||||
api_key = os.environ.get("GEMMA_API_KEY", "not-needed")
|
||||
|
||||
await reg.add_step(self.task_id, "thinking", f"Starting: {self.description}")
|
||||
|
||||
client = AsyncOpenAI(base_url=base_url, api_key=api_key)
|
||||
tools = self._load_tools()
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": self._system_prompt()},
|
||||
{"role": "user", "content": self.description},
|
||||
]
|
||||
|
||||
max_iterations = 10
|
||||
for iteration in range(max_iterations):
|
||||
if self._cancelled:
|
||||
await reg.fail(self.task_id, "Cancelled")
|
||||
return
|
||||
|
||||
await reg.add_step(
|
||||
self.task_id, "thinking", f"Step {iteration + 1}: analyzing..."
|
||||
)
|
||||
|
||||
try:
|
||||
response = await client.chat.completions.create(
|
||||
model=model,
|
||||
messages=messages,
|
||||
tools=tools or None,
|
||||
max_tokens=500,
|
||||
temperature=0.7,
|
||||
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
|
||||
)
|
||||
except Exception as e:
|
||||
await reg.fail(self.task_id, f"LLM error: {e}")
|
||||
return
|
||||
|
||||
msg = response.choices[0].message
|
||||
|
||||
if msg.tool_calls:
|
||||
messages.append(msg)
|
||||
for tc in msg.tool_calls:
|
||||
tool_name = tc.function.name
|
||||
try:
|
||||
tool_args = json.loads(tc.function.arguments)
|
||||
except json.JSONDecodeError:
|
||||
tool_args = {}
|
||||
await reg.add_step(
|
||||
self.task_id,
|
||||
"tool_call",
|
||||
f"{tool_name}({json.dumps(tool_args)})",
|
||||
)
|
||||
result = await self._execute_tool(tool_name, tool_args)
|
||||
await reg.add_step(self.task_id, "tool_result", str(result)[:500])
|
||||
messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tc.id,
|
||||
"content": str(result),
|
||||
}
|
||||
)
|
||||
else:
|
||||
text = (msg.content or "").strip()
|
||||
if not text:
|
||||
await reg.fail(self.task_id, "LLM returned empty response")
|
||||
return
|
||||
await reg.add_step(self.task_id, "text", text)
|
||||
await reg.complete(self.task_id, text)
|
||||
return
|
||||
|
||||
await reg.fail(self.task_id, f"Reached max iterations ({max_iterations})")
|
||||
|
||||
def _system_prompt(self) -> str:
|
||||
return (
|
||||
"You are Hope's research assistant working on a background task. "
|
||||
"Use your tools to gather information thoroughly. When you have "
|
||||
"enough, provide a final answer in 2-4 sentences suitable for "
|
||||
"speaking aloud. Plain text only, no markdown or formatting."
|
||||
)
|
||||
|
||||
def _load_tools(self) -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get current weather and today's forecast for a location in Fahrenheit.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"location": {"type": "string"}},
|
||||
"required": ["location"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_time",
|
||||
"description": "Get the current date and time, optionally for a specific location.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"location": {"type": "string"}},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "memory_recall",
|
||||
"description": "Search saved memories for relevant information.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"query": {"type": "string"}},
|
||||
"required": ["query"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "web_search",
|
||||
"description": "Search the web for current information using Firecrawl.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"query": {"type": "string"}},
|
||||
"required": ["query"],
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
async def _execute_tool(self, name: str, args: dict) -> str:
|
||||
import httpx
|
||||
|
||||
try:
|
||||
if name == "get_weather":
|
||||
location = args.get("location", "")
|
||||
async with httpx.AsyncClient(timeout=30) as client:
|
||||
resp = await client.get(
|
||||
f"https://wttr.in/{location}?format=j1",
|
||||
headers={"User-Agent": "curl/8.0"},
|
||||
)
|
||||
data = resp.json()
|
||||
current = data.get("current_condition", [{}])[0] or {}
|
||||
temp_c = current.get("temp_C")
|
||||
if temp_c is None:
|
||||
return f"Could not get weather for {location}."
|
||||
temp_f = round(int(temp_c) * 9 / 5 + 32)
|
||||
desc = (current.get("weatherDesc", [{}])[0].get("value") or "unknown").lower()
|
||||
feels_c = current.get("FeelsLikeC")
|
||||
feels_f = round(int(feels_c) * 9 / 5 + 32) if feels_c else None
|
||||
today = (data.get("weather") or [{}])[0] or {}
|
||||
high_f = round(int(today["maxtempC"]) * 9 / 5 + 32) if today.get("maxtempC") else None
|
||||
low_f = round(int(today["mintempC"]) * 9 / 5 + 32) if today.get("mintempC") else None
|
||||
parts = [f"It's {temp_f} degrees and {desc} in {location}"]
|
||||
if feels_f and feels_f != temp_f:
|
||||
parts[0] += f", feels like {feels_f}"
|
||||
if high_f and low_f:
|
||||
parts.append(f"High {high_f}, low {low_f}.")
|
||||
return " ".join(parts)
|
||||
|
||||
elif name == "get_time":
|
||||
now = datetime.now(timezone.utc)
|
||||
loc = args.get("location", "")
|
||||
suffix = f" in {loc}" if loc else ", UTC"
|
||||
return f"It's {now.strftime('%I:%M %p')} on {now.strftime('%A, %B %d')}{suffix}."
|
||||
|
||||
elif name == "memory_recall":
|
||||
query = args.get("query", "").lower()
|
||||
memory_dir = os.environ.get("MEMORY_DIR", "/memory")
|
||||
results = []
|
||||
try:
|
||||
for fname in os.listdir(memory_dir):
|
||||
if not fname.endswith(".md"):
|
||||
continue
|
||||
with open(os.path.join(memory_dir, fname)) as f:
|
||||
content = f.read()
|
||||
if any(w in content.lower() for w in query.split()):
|
||||
results.append(f"[{fname}] {content[:300]}")
|
||||
except OSError:
|
||||
pass
|
||||
return "\n".join(results) or "No memories found."
|
||||
|
||||
elif name == "web_search":
|
||||
query = args.get("query", "")
|
||||
firecrawl_base = os.environ.get("FIRECRAWL_BASE", "")
|
||||
if not firecrawl_base:
|
||||
return "Web search is not configured."
|
||||
async with httpx.AsyncClient(timeout=30) as client:
|
||||
resp = await client.post(
|
||||
f"{firecrawl_base}/v1/search",
|
||||
json={"query": query, "limit": 3},
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
data = resp.json()
|
||||
items = data.get("data", [])
|
||||
lines = [
|
||||
f"{r.get('title', '')}: {r.get('description', '')}"
|
||||
for r in items[:3]
|
||||
]
|
||||
return "\n".join(lines) or "No results found."
|
||||
return f"Search failed (HTTP {resp.status_code})."
|
||||
|
||||
else:
|
||||
return f"Unknown tool: {name}"
|
||||
|
||||
except Exception as e:
|
||||
return f"Tool error: {e}"
|
||||
|
||||
def cancel(self):
|
||||
self._cancelled = True
|
||||
Reference in New Issue
Block a user