forked from Rockachopa/Timmy-time-dashboard
feat: code quality audit + autoresearch integration + infra hardening (#150)
This commit is contained in:
Alexander Whitestone
committed by
GitHub
GitHub
parent
fd0ede0d51
commit
ae3bb1cc21
@@ -21,52 +21,54 @@ DEFAULT_RQLITE_URL = "http://localhost:4001"
|
||||
|
||||
class BrainClient:
|
||||
"""Client for distributed brain (rqlite).
|
||||
|
||||
|
||||
Connects to local rqlite instance, which handles replication.
|
||||
All writes go to leader, reads can come from local node.
|
||||
"""
|
||||
|
||||
|
||||
def __init__(self, rqlite_url: Optional[str] = None, node_id: Optional[str] = None):
|
||||
from config import settings
|
||||
|
||||
self.rqlite_url = rqlite_url or settings.rqlite_url or DEFAULT_RQLITE_URL
|
||||
self.node_id = node_id or f"{socket.gethostname()}-{os.getpid()}"
|
||||
self.source = self._detect_source()
|
||||
self._client = httpx.AsyncClient(timeout=30)
|
||||
|
||||
|
||||
def _detect_source(self) -> str:
|
||||
"""Detect what component is using the brain."""
|
||||
# Could be 'timmy', 'zeroclaw', 'worker', etc.
|
||||
# For now, infer from context or env
|
||||
from config import settings
|
||||
|
||||
return settings.brain_source
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────
|
||||
# Memory Operations
|
||||
# ──────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def remember(
|
||||
self,
|
||||
content: str,
|
||||
tags: Optional[List[str]] = None,
|
||||
source: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Store a memory with embedding.
|
||||
|
||||
|
||||
Args:
|
||||
content: Text content to remember
|
||||
tags: Optional list of tags (e.g., ['shell', 'result'])
|
||||
source: Source identifier (defaults to self.source)
|
||||
metadata: Additional JSON-serializable metadata
|
||||
|
||||
|
||||
Returns:
|
||||
Dict with 'id' and 'status'
|
||||
"""
|
||||
from brain.embeddings import get_embedder
|
||||
|
||||
|
||||
embedder = get_embedder()
|
||||
embedding_bytes = embedder.encode_single(content)
|
||||
|
||||
|
||||
query = """
|
||||
INSERT INTO memories (content, embedding, source, tags, metadata, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
@@ -77,100 +79,90 @@ class BrainClient:
|
||||
source or self.source,
|
||||
json.dumps(tags or []),
|
||||
json.dumps(metadata or {}),
|
||||
datetime.utcnow().isoformat()
|
||||
datetime.utcnow().isoformat(),
|
||||
]
|
||||
|
||||
|
||||
try:
|
||||
resp = await self._client.post(
|
||||
f"{self.rqlite_url}/db/execute",
|
||||
json=[query, params]
|
||||
)
|
||||
resp = await self._client.post(f"{self.rqlite_url}/db/execute", json=[query, params])
|
||||
resp.raise_for_status()
|
||||
result = resp.json()
|
||||
|
||||
|
||||
# Extract inserted ID
|
||||
last_id = None
|
||||
if "results" in result and result["results"]:
|
||||
last_id = result["results"][0].get("last_insert_id")
|
||||
|
||||
|
||||
logger.debug(f"Stored memory {last_id}: {content[:50]}...")
|
||||
return {"id": last_id, "status": "stored"}
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to store memory: {e}")
|
||||
raise
|
||||
|
||||
|
||||
async def recall(
|
||||
self,
|
||||
query: str,
|
||||
limit: int = 5,
|
||||
sources: Optional[List[str]] = None
|
||||
self, query: str, limit: int = 5, sources: Optional[List[str]] = None
|
||||
) -> List[str]:
|
||||
"""Semantic search for memories.
|
||||
|
||||
|
||||
Args:
|
||||
query: Search query text
|
||||
limit: Max results to return
|
||||
sources: Filter by source(s) (e.g., ['timmy', 'user'])
|
||||
|
||||
|
||||
Returns:
|
||||
List of memory content strings
|
||||
"""
|
||||
from brain.embeddings import get_embedder
|
||||
|
||||
|
||||
embedder = get_embedder()
|
||||
query_emb = embedder.encode_single(query)
|
||||
|
||||
|
||||
# rqlite with sqlite-vec extension for vector search
|
||||
sql = "SELECT content, source, metadata, distance FROM memories WHERE embedding MATCH ?"
|
||||
params = [query_emb]
|
||||
|
||||
|
||||
if sources:
|
||||
placeholders = ",".join(["?"] * len(sources))
|
||||
sql += f" AND source IN ({placeholders})"
|
||||
params.extend(sources)
|
||||
|
||||
|
||||
sql += " ORDER BY distance LIMIT ?"
|
||||
params.append(limit)
|
||||
|
||||
|
||||
try:
|
||||
resp = await self._client.post(
|
||||
f"{self.rqlite_url}/db/query",
|
||||
json=[sql, params]
|
||||
)
|
||||
resp = await self._client.post(f"{self.rqlite_url}/db/query", json=[sql, params])
|
||||
resp.raise_for_status()
|
||||
result = resp.json()
|
||||
|
||||
|
||||
results = []
|
||||
if "results" in result and result["results"]:
|
||||
for row in result["results"][0].get("rows", []):
|
||||
results.append({
|
||||
"content": row[0],
|
||||
"source": row[1],
|
||||
"metadata": json.loads(row[2]) if row[2] else {},
|
||||
"distance": row[3]
|
||||
})
|
||||
|
||||
results.append(
|
||||
{
|
||||
"content": row[0],
|
||||
"source": row[1],
|
||||
"metadata": json.loads(row[2]) if row[2] else {},
|
||||
"distance": row[3],
|
||||
}
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to search memories: {e}")
|
||||
# Graceful fallback - return empty list
|
||||
return []
|
||||
|
||||
|
||||
async def get_recent(
|
||||
self,
|
||||
hours: int = 24,
|
||||
limit: int = 20,
|
||||
sources: Optional[List[str]] = None
|
||||
self, hours: int = 24, limit: int = 20, sources: Optional[List[str]] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Get recent memories by time.
|
||||
|
||||
|
||||
Args:
|
||||
hours: Look back this many hours
|
||||
limit: Max results
|
||||
sources: Optional source filter
|
||||
|
||||
|
||||
Returns:
|
||||
List of memory dicts
|
||||
"""
|
||||
@@ -180,84 +172,83 @@ class BrainClient:
|
||||
WHERE created_at > datetime('now', ?)
|
||||
"""
|
||||
params = [f"-{hours} hours"]
|
||||
|
||||
|
||||
if sources:
|
||||
placeholders = ",".join(["?"] * len(sources))
|
||||
sql += f" AND source IN ({placeholders})"
|
||||
params.extend(sources)
|
||||
|
||||
|
||||
sql += " ORDER BY created_at DESC LIMIT ?"
|
||||
params.append(limit)
|
||||
|
||||
|
||||
try:
|
||||
resp = await self._client.post(
|
||||
f"{self.rqlite_url}/db/query",
|
||||
json=[sql, params]
|
||||
)
|
||||
resp = await self._client.post(f"{self.rqlite_url}/db/query", json=[sql, params])
|
||||
resp.raise_for_status()
|
||||
result = resp.json()
|
||||
|
||||
|
||||
memories = []
|
||||
if "results" in result and result["results"]:
|
||||
for row in result["results"][0].get("rows", []):
|
||||
memories.append({
|
||||
"id": row[0],
|
||||
"content": row[1],
|
||||
"source": row[2],
|
||||
"tags": json.loads(row[3]) if row[3] else [],
|
||||
"metadata": json.loads(row[4]) if row[4] else {},
|
||||
"created_at": row[5]
|
||||
})
|
||||
|
||||
memories.append(
|
||||
{
|
||||
"id": row[0],
|
||||
"content": row[1],
|
||||
"source": row[2],
|
||||
"tags": json.loads(row[3]) if row[3] else [],
|
||||
"metadata": json.loads(row[4]) if row[4] else {},
|
||||
"created_at": row[5],
|
||||
}
|
||||
)
|
||||
|
||||
return memories
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get recent memories: {e}")
|
||||
return []
|
||||
|
||||
|
||||
async def get_context(self, query: str) -> str:
|
||||
"""Get formatted context for system prompt.
|
||||
|
||||
|
||||
Combines recent memories + relevant memories.
|
||||
|
||||
|
||||
Args:
|
||||
query: Current user query to find relevant context
|
||||
|
||||
|
||||
Returns:
|
||||
Formatted context string for prompt injection
|
||||
"""
|
||||
recent = await self.get_recent(hours=24, limit=10)
|
||||
relevant = await self.recall(query, limit=5)
|
||||
|
||||
|
||||
lines = ["Recent activity:"]
|
||||
for m in recent[:5]:
|
||||
lines.append(f"- {m['content'][:100]}")
|
||||
|
||||
|
||||
lines.append("\nRelevant memories:")
|
||||
for r in relevant[:5]:
|
||||
lines.append(f"- {r['content'][:100]}")
|
||||
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────
|
||||
# Task Queue Operations
|
||||
# ──────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def submit_task(
|
||||
self,
|
||||
content: str,
|
||||
task_type: str = "general",
|
||||
priority: int = 0,
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Submit a task to the distributed queue.
|
||||
|
||||
|
||||
Args:
|
||||
content: Task description/prompt
|
||||
task_type: Type of task (shell, creative, code, research, general)
|
||||
priority: Higher = processed first
|
||||
metadata: Additional task data
|
||||
|
||||
|
||||
Returns:
|
||||
Dict with task 'id'
|
||||
"""
|
||||
@@ -270,50 +261,45 @@ class BrainClient:
|
||||
task_type,
|
||||
priority,
|
||||
json.dumps(metadata or {}),
|
||||
datetime.utcnow().isoformat()
|
||||
datetime.utcnow().isoformat(),
|
||||
]
|
||||
|
||||
|
||||
try:
|
||||
resp = await self._client.post(
|
||||
f"{self.rqlite_url}/db/execute",
|
||||
json=[query, params]
|
||||
)
|
||||
resp = await self._client.post(f"{self.rqlite_url}/db/execute", json=[query, params])
|
||||
resp.raise_for_status()
|
||||
result = resp.json()
|
||||
|
||||
|
||||
last_id = None
|
||||
if "results" in result and result["results"]:
|
||||
last_id = result["results"][0].get("last_insert_id")
|
||||
|
||||
|
||||
logger.info(f"Submitted task {last_id}: {content[:50]}...")
|
||||
return {"id": last_id, "status": "queued"}
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to submit task: {e}")
|
||||
raise
|
||||
|
||||
|
||||
async def claim_task(
|
||||
self,
|
||||
capabilities: List[str],
|
||||
node_id: Optional[str] = None
|
||||
self, capabilities: List[str], node_id: Optional[str] = None
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Atomically claim next available task.
|
||||
|
||||
|
||||
Uses UPDATE ... RETURNING pattern for atomic claim.
|
||||
|
||||
|
||||
Args:
|
||||
capabilities: List of capabilities this node has
|
||||
node_id: Identifier for claiming node
|
||||
|
||||
|
||||
Returns:
|
||||
Task dict or None if no tasks available
|
||||
"""
|
||||
claimer = node_id or self.node_id
|
||||
|
||||
|
||||
# Try to claim a matching task atomically
|
||||
# This works because rqlite uses Raft consensus - only one node wins
|
||||
placeholders = ",".join(["?"] * len(capabilities))
|
||||
|
||||
|
||||
query = f"""
|
||||
UPDATE tasks
|
||||
SET status = 'claimed',
|
||||
@@ -330,15 +316,12 @@ class BrainClient:
|
||||
RETURNING id, content, task_type, priority, metadata
|
||||
"""
|
||||
params = [claimer, datetime.utcnow().isoformat()] + capabilities
|
||||
|
||||
|
||||
try:
|
||||
resp = await self._client.post(
|
||||
f"{self.rqlite_url}/db/execute",
|
||||
json=[query, params]
|
||||
)
|
||||
resp = await self._client.post(f"{self.rqlite_url}/db/execute", json=[query, params])
|
||||
resp.raise_for_status()
|
||||
result = resp.json()
|
||||
|
||||
|
||||
if "results" in result and result["results"]:
|
||||
rows = result["results"][0].get("rows", [])
|
||||
if rows:
|
||||
@@ -348,24 +331,20 @@ class BrainClient:
|
||||
"content": row[1],
|
||||
"type": row[2],
|
||||
"priority": row[3],
|
||||
"metadata": json.loads(row[4]) if row[4] else {}
|
||||
"metadata": json.loads(row[4]) if row[4] else {},
|
||||
}
|
||||
|
||||
|
||||
return None
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to claim task: {e}")
|
||||
return None
|
||||
|
||||
|
||||
async def complete_task(
|
||||
self,
|
||||
task_id: int,
|
||||
success: bool,
|
||||
result: Optional[str] = None,
|
||||
error: Optional[str] = None
|
||||
self, task_id: int, success: bool, result: Optional[str] = None, error: Optional[str] = None
|
||||
) -> None:
|
||||
"""Mark task as completed or failed.
|
||||
|
||||
|
||||
Args:
|
||||
task_id: Task ID
|
||||
success: True if task succeeded
|
||||
@@ -373,7 +352,7 @@ class BrainClient:
|
||||
error: Error message if failed
|
||||
"""
|
||||
status = "done" if success else "failed"
|
||||
|
||||
|
||||
query = """
|
||||
UPDATE tasks
|
||||
SET status = ?,
|
||||
@@ -383,23 +362,20 @@ class BrainClient:
|
||||
WHERE id = ?
|
||||
"""
|
||||
params = [status, result, error, datetime.utcnow().isoformat(), task_id]
|
||||
|
||||
|
||||
try:
|
||||
await self._client.post(
|
||||
f"{self.rqlite_url}/db/execute",
|
||||
json=[query, params]
|
||||
)
|
||||
await self._client.post(f"{self.rqlite_url}/db/execute", json=[query, params])
|
||||
logger.debug(f"Task {task_id} marked {status}")
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to complete task {task_id}: {e}")
|
||||
|
||||
|
||||
async def get_pending_tasks(self, limit: int = 100) -> List[Dict[str, Any]]:
|
||||
"""Get list of pending tasks (for dashboard/monitoring).
|
||||
|
||||
|
||||
Args:
|
||||
limit: Max tasks to return
|
||||
|
||||
|
||||
Returns:
|
||||
List of pending task dicts
|
||||
"""
|
||||
@@ -410,33 +386,32 @@ class BrainClient:
|
||||
ORDER BY priority DESC, created_at ASC
|
||||
LIMIT ?
|
||||
"""
|
||||
|
||||
|
||||
try:
|
||||
resp = await self._client.post(
|
||||
f"{self.rqlite_url}/db/query",
|
||||
json=[sql, [limit]]
|
||||
)
|
||||
resp = await self._client.post(f"{self.rqlite_url}/db/query", json=[sql, [limit]])
|
||||
resp.raise_for_status()
|
||||
result = resp.json()
|
||||
|
||||
|
||||
tasks = []
|
||||
if "results" in result and result["results"]:
|
||||
for row in result["results"][0].get("rows", []):
|
||||
tasks.append({
|
||||
"id": row[0],
|
||||
"content": row[1],
|
||||
"type": row[2],
|
||||
"priority": row[3],
|
||||
"metadata": json.loads(row[4]) if row[4] else {},
|
||||
"created_at": row[5]
|
||||
})
|
||||
|
||||
tasks.append(
|
||||
{
|
||||
"id": row[0],
|
||||
"content": row[1],
|
||||
"type": row[2],
|
||||
"priority": row[3],
|
||||
"metadata": json.loads(row[4]) if row[4] else {},
|
||||
"created_at": row[5],
|
||||
}
|
||||
)
|
||||
|
||||
return tasks
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get pending tasks: {e}")
|
||||
return []
|
||||
|
||||
|
||||
async def close(self):
|
||||
"""Close HTTP client."""
|
||||
await self._client.aclose()
|
||||
|
||||
Reference in New Issue
Block a user