203 lines
7.7 KiB
Python
203 lines
7.7 KiB
Python
"""
|
|
Copyright 2024-2026 ChatterMate
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
"""
|
|
|
|
from sqlalchemy.orm import Session, joinedload
|
|
from typing import List, Optional
|
|
from app.models.agent import Agent
|
|
from app.models.mcp_tool import MCPTool, MCPToolToAgent, MCPTransportType
|
|
from app.models.ticket_settings import OrganizationTicketSettings
|
|
from uuid import UUID
|
|
from app.core.logger import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class MCPToolRepository:
|
|
def __init__(self, db: Session):
|
|
self.db = db
|
|
|
|
def create_mcp_tool(self, **kwargs) -> MCPTool:
|
|
"""Create a new MCP tool"""
|
|
# Handle org_id to organization_id conversion for backward compatibility
|
|
if 'org_id' in kwargs:
|
|
kwargs['organization_id'] = kwargs.pop('org_id')
|
|
|
|
# Create new MCP tool
|
|
mcp_tool = MCPTool(**kwargs)
|
|
self.db.add(mcp_tool)
|
|
self.db.commit()
|
|
self.db.refresh(mcp_tool)
|
|
return mcp_tool
|
|
|
|
def get_by_name(self, name: str, organization_id: UUID) -> Optional[MCPTool]:
|
|
"""Get MCP tool by name within an organization"""
|
|
return self.db.query(MCPTool).filter(
|
|
MCPTool.name == name,
|
|
MCPTool.organization_id == organization_id
|
|
).first()
|
|
|
|
def get_mcp_tool(self, mcp_tool_id: int) -> Optional[MCPTool]:
|
|
"""Get MCP tool by ID"""
|
|
return self.db.query(MCPTool).filter(
|
|
MCPTool.id == mcp_tool_id
|
|
).first()
|
|
|
|
def get_by_ids(self, org_id: UUID, tool_ids: List[int], enabled_only: bool = True) -> List[MCPTool]:
|
|
"""Get specific MCP tools by id, scoped to the organization — ids
|
|
from another org are silently dropped."""
|
|
if not tool_ids:
|
|
return []
|
|
query = self.db.query(MCPTool).filter(
|
|
MCPTool.organization_id == org_id,
|
|
MCPTool.id.in_(tool_ids),
|
|
)
|
|
if enabled_only:
|
|
query = query.filter(MCPTool.enabled == True)
|
|
return query.all()
|
|
|
|
def get_org_mcp_tools(self, org_id: UUID, enabled_only: bool = True) -> List[MCPTool]:
|
|
"""Get all MCP tools for an organization"""
|
|
query = self.db.query(MCPTool).filter(
|
|
MCPTool.organization_id == org_id
|
|
)
|
|
if enabled_only:
|
|
query = query.filter(MCPTool.enabled == True)
|
|
return query.all()
|
|
|
|
def update_mcp_tool(self, mcp_tool_id: int, **kwargs) -> Optional[MCPTool]:
|
|
"""Update an existing MCP tool"""
|
|
mcp_tool = self.get_mcp_tool(mcp_tool_id)
|
|
if not mcp_tool:
|
|
return None
|
|
|
|
for key, value in kwargs.items():
|
|
setattr(mcp_tool, key, value)
|
|
|
|
self.db.commit()
|
|
self.db.refresh(mcp_tool)
|
|
return mcp_tool
|
|
|
|
def _get_ticket_settings(self, organization_id: UUID) -> Optional[OrganizationTicketSettings]:
|
|
return self.db.query(OrganizationTicketSettings).filter(
|
|
OrganizationTicketSettings.organization_id == organization_id
|
|
).first()
|
|
|
|
def is_used_in_investigations(self, mcp_tool_id: int, organization_id: UUID) -> bool:
|
|
"""Whether the org's investigation connector selection points at this tool."""
|
|
settings = self._get_ticket_settings(organization_id)
|
|
return bool(settings and mcp_tool_id in (settings.investigation_mcp_tool_ids or []))
|
|
|
|
def get_mcp_tool_agent_names(self, mcp_tool_id: int) -> List[str]:
|
|
"""Names of the agents this tool is linked to."""
|
|
rows = self.db.query(Agent.name)\
|
|
.join(MCPToolToAgent, MCPToolToAgent.agent_id == Agent.id)\
|
|
.filter(MCPToolToAgent.mcp_tool_id == mcp_tool_id)\
|
|
.order_by(Agent.name)\
|
|
.all()
|
|
return [name for (name,) in rows]
|
|
|
|
def delete_mcp_tool(self, mcp_tool_id: int) -> bool:
|
|
"""Delete an MCP tool, along with every reference to it."""
|
|
mcp_tool = self.get_mcp_tool(mcp_tool_id)
|
|
if not mcp_tool:
|
|
return False
|
|
|
|
# investigation_mcp_tool_ids is a plain JSON id list with no foreign
|
|
# key, so a deleted connector would otherwise stay "configured"
|
|
# forever and every later run would report fewer connectors loaded
|
|
# than configured — the exact signal #271 added to catch real
|
|
# failures. Agent links go with the tool via cascade.
|
|
settings = self._get_ticket_settings(mcp_tool.organization_id)
|
|
if settings and mcp_tool_id in (settings.investigation_mcp_tool_ids or []):
|
|
settings.investigation_mcp_tool_ids = [
|
|
tool_id for tool_id in settings.investigation_mcp_tool_ids
|
|
if tool_id != mcp_tool_id
|
|
]
|
|
|
|
self.db.delete(mcp_tool)
|
|
self.db.commit()
|
|
return True
|
|
|
|
def get_enabled_mcp_tools(self, org_id: UUID) -> List[MCPTool]:
|
|
"""Get enabled MCP tools for an organization"""
|
|
return self.db.query(MCPTool)\
|
|
.filter(MCPTool.organization_id == org_id)\
|
|
.filter(MCPTool.enabled == True)\
|
|
.all()
|
|
|
|
def get_all_mcp_tools(self, org_id: UUID) -> List[MCPTool]:
|
|
"""Get all MCP tools for organization"""
|
|
return self.db.query(MCPTool)\
|
|
.filter(MCPTool.organization_id == org_id)\
|
|
.all()
|
|
|
|
# Agent-MCPTool association methods
|
|
def add_mcp_tool_to_agent(self, mcp_tool_id: int, agent_id: UUID) -> Optional[MCPToolToAgent]:
|
|
"""Add MCP tool to agent"""
|
|
# Check if association already exists
|
|
existing = self.db.query(MCPToolToAgent).filter(
|
|
MCPToolToAgent.mcp_tool_id == mcp_tool_id,
|
|
MCPToolToAgent.agent_id == agent_id
|
|
).first()
|
|
|
|
if existing:
|
|
return existing
|
|
|
|
# Create new association
|
|
association = MCPToolToAgent(
|
|
mcp_tool_id=mcp_tool_id,
|
|
agent_id=agent_id
|
|
)
|
|
self.db.add(association)
|
|
self.db.commit()
|
|
self.db.refresh(association)
|
|
return association
|
|
|
|
def remove_mcp_tool_from_agent(self, mcp_tool_id: int, agent_id: UUID) -> bool:
|
|
"""Remove MCP tool from agent"""
|
|
association = self.db.query(MCPToolToAgent).filter(
|
|
MCPToolToAgent.mcp_tool_id == mcp_tool_id,
|
|
MCPToolToAgent.agent_id == agent_id
|
|
).first()
|
|
|
|
if not association:
|
|
return False
|
|
|
|
self.db.delete(association)
|
|
self.db.commit()
|
|
return True
|
|
|
|
def get_agent_mcp_tools(self, agent_id: UUID) -> List[MCPTool]:
|
|
"""Get all MCP tools associated with an agent"""
|
|
return self.db.query(MCPTool)\
|
|
.join(MCPToolToAgent)\
|
|
.filter(MCPToolToAgent.agent_id == agent_id)\
|
|
.filter(MCPTool.enabled == True)\
|
|
.all()
|
|
|
|
def get_mcp_tool_agents(self, mcp_tool_id: int) -> List[UUID]:
|
|
"""Get all agent IDs associated with an MCP tool"""
|
|
associations = self.db.query(MCPToolToAgent).filter(
|
|
MCPToolToAgent.mcp_tool_id == mcp_tool_id
|
|
).all()
|
|
return [assoc.agent_id for assoc in associations]
|
|
|
|
def get_mcp_tool_with_agents(self, mcp_tool_id: int) -> Optional[MCPTool]:
|
|
"""Get MCP tool with associated agents loaded"""
|
|
return self.db.query(MCPTool)\
|
|
.options(joinedload(MCPTool.agent_links))\
|
|
.filter(MCPTool.id == mcp_tool_id)\
|
|
.first() |