tools/database_tool.py

"""
Database tool: wraps Neo4j and MongoDB for agent read/write operations.
"""
from typing import Any, Dict, List, Optional
from tools.base_tool import BaseTool, ToolResult
from services.neo4j_service import Neo4jService
from services.mongo_service import MongoService
from utils.logger import get_logger

logger = get_logger("database_tool")


class Neo4jTool(BaseTool):
    name = "neo4j_query"
    description = "Run Cypher queries against the knowledge graph."

    def __init__(self, neo4j: Neo4jService):
        self.neo4j = neo4j

    async def run(self, query: str, write: bool = False, **params) -> ToolResult:
        try:
            if write:
                results = self.neo4j.run_write(query, **params)
            else:
                results = self.neo4j.run(query, **params)
            return self._ok(results)
        except Exception as e:
            logger.error(f"Neo4j query failed: {e}")
            return self._err(str(e))


class MongoTool(BaseTool):
    name = "mongo_job"
    description = "Read/write job state and results in MongoDB."

    def __init__(self, mongo: MongoService):
        self.mongo = mongo

    async def run(
        self,
        operation: str,
        job_id: str,
        data: Optional[Dict] = None,
    ) -> ToolResult:
        try:
            if operation == "get_job":
                result = await self.mongo.get_job(job_id)
            elif operation == "get_result":
                result = await self.mongo.get_result(job_id)
            elif operation == "save_result":
                await self.mongo.save_result(job_id, data or {})
                result = {"saved": True}
            elif operation == "upsert_job":
                await self.mongo.upsert_job(job_id, data or {})
                result = {"upserted": True}
            else:
                return self._err(f"Unknown operation: {operation}")
            return self._ok(result)
        except Exception as e:
            logger.error(f"Mongo operation '{operation}' failed: {e}")
            return self._err(str(e))