130 lines
5.2 KiB
Python
130 lines
5.2 KiB
Python
"""
|
|
tools/brain_tool.py — 2nd Brain tool for Jarvis.
|
|
|
|
Gives Jarvis (and any agent with brain access) two operations:
|
|
capture — save a thought to the brain
|
|
search — retrieve thoughts by semantic similarity
|
|
browse — list recent thoughts (optional type filter)
|
|
stats — database statistics
|
|
|
|
Capture is the only write operation and requires no confirmation (it's
|
|
non-destructive and the user expects it to be instant).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
|
|
from ..context_vars import current_user as _current_user_var
|
|
from .base import BaseTool, ToolResult
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class BrainTool(BaseTool):
|
|
name = "brain"
|
|
description = (
|
|
"Access the 2nd Brain knowledge base. "
|
|
"Operations: capture (save a thought), search (semantic search by meaning), "
|
|
"browse (recent thoughts), stats (database overview). "
|
|
"Use 'capture' to save anything worth remembering. "
|
|
"Use 'search' to find relevant past thoughts by meaning, not just keywords."
|
|
)
|
|
input_schema = {
|
|
"type": "object",
|
|
"properties": {
|
|
"operation": {
|
|
"type": "string",
|
|
"enum": ["capture", "search", "browse", "stats"],
|
|
"description": "Operation to perform.",
|
|
},
|
|
"content": {
|
|
"type": "string",
|
|
"description": "Text to capture (required for 'capture').",
|
|
},
|
|
"query": {
|
|
"type": "string",
|
|
"description": "Search query (required for 'search').",
|
|
},
|
|
"threshold": {
|
|
"type": "number",
|
|
"description": "Similarity threshold 0-1 for 'search' (default 0.7). Lower = broader results.",
|
|
"default": 0.7,
|
|
},
|
|
"limit": {
|
|
"type": "integer",
|
|
"description": "Max results for 'search' or 'browse' (default 10).",
|
|
"default": 10,
|
|
},
|
|
"type_filter": {
|
|
"type": "string",
|
|
"description": "Filter 'browse' by thought type: insight, person_note, task, reference, idea, other.",
|
|
},
|
|
},
|
|
"required": ["operation"],
|
|
}
|
|
|
|
requires_confirmation = False
|
|
allowed_in_scheduled_tasks = True
|
|
|
|
async def execute(self, **kwargs) -> ToolResult:
|
|
operation = kwargs.get("operation")
|
|
# Resolve current user for brain namespace scoping (3-G)
|
|
_user = _current_user_var.get()
|
|
user_id = _user.id if _user else None
|
|
|
|
try:
|
|
from ..brain.database import get_pool
|
|
if get_pool() is None:
|
|
return ToolResult(
|
|
success=False,
|
|
error="Brain DB is not available. Check BRAIN_DB_URL in .env.",
|
|
)
|
|
|
|
if operation == "capture":
|
|
return await self._capture(kwargs.get("content", ""), user_id=user_id)
|
|
elif operation == "search":
|
|
return await self._search(
|
|
kwargs.get("query", ""),
|
|
float(kwargs.get("threshold", 0.7)),
|
|
int(kwargs.get("limit", 10)),
|
|
user_id=user_id,
|
|
)
|
|
elif operation == "browse":
|
|
return await self._browse(
|
|
int(kwargs.get("limit", 10)),
|
|
kwargs.get("type_filter"),
|
|
user_id=user_id,
|
|
)
|
|
elif operation == "stats":
|
|
return await self._stats(user_id=user_id)
|
|
else:
|
|
return ToolResult(success=False, error=f"Unknown operation: {operation}")
|
|
|
|
except Exception as e:
|
|
logger.error("BrainTool error (%s): %s", operation, e)
|
|
return ToolResult(success=False, error=str(e))
|
|
|
|
async def _capture(self, content: str, user_id: str | None = None) -> ToolResult:
|
|
if not content.strip():
|
|
return ToolResult(success=False, error="content is required for capture")
|
|
from ..brain.ingest import ingest_thought
|
|
result = await ingest_thought(content, user_id=user_id)
|
|
return ToolResult(success=True, data=result)
|
|
|
|
async def _search(self, query: str, threshold: float, limit: int, user_id: str | None = None) -> ToolResult:
|
|
if not query.strip():
|
|
return ToolResult(success=False, error="query is required for search")
|
|
from ..brain.search import semantic_search
|
|
results = await semantic_search(query, threshold=threshold, limit=limit, user_id=user_id)
|
|
return ToolResult(success=True, data={"results": results, "count": len(results)})
|
|
|
|
async def _browse(self, limit: int, type_filter: str | None, user_id: str | None = None) -> ToolResult:
|
|
from ..brain.database import browse_thoughts
|
|
results = await browse_thoughts(limit=limit, type_filter=type_filter, user_id=user_id)
|
|
return ToolResult(success=True, data={"results": results, "count": len(results)})
|
|
|
|
async def _stats(self, user_id: str | None = None) -> ToolResult:
|
|
from ..brain.database import get_stats
|
|
stats = await get_stats(user_id=user_id)
|
|
return ToolResult(success=True, data=stats)
|