Spaces:
Running
Running
File size: 2,061 Bytes
0e38162 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | """
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))
|