183 lines
6.4 KiB
Python
183 lines
6.4 KiB
Python
"""
|
|
Copyright 2024-2026 ReplyPier
|
|
|
|
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
|
|
from typing import List, Optional
|
|
from app.models.agent import Agent
|
|
import json
|
|
from sqlalchemy.orm import joinedload
|
|
from uuid import UUID
|
|
from app.core.logger import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
class AgentRepository:
|
|
def __init__(self, db: Session):
|
|
self.db = db
|
|
|
|
def create_agent(self, **kwargs) -> Agent:
|
|
"""Create a new agent"""
|
|
# Handle instructions
|
|
# We don't need to do any conversion here, as the Agent model will handle it
|
|
# through its property setter
|
|
|
|
# Handle tools if present
|
|
if 'tools' in kwargs and kwargs['tools']:
|
|
if not isinstance(kwargs['tools'], str):
|
|
kwargs['tools'] = json.dumps(kwargs['tools'])
|
|
|
|
# Handle org_id to organization_id conversion for backward compatibility
|
|
if 'org_id' in kwargs:
|
|
kwargs['organization_id'] = kwargs.pop('org_id')
|
|
|
|
# Create new agent
|
|
agent = Agent(**kwargs)
|
|
self.db.add(agent)
|
|
self.db.commit()
|
|
self.db.refresh(agent)
|
|
return agent
|
|
|
|
def get_by_name(self, name: str, organization_id: UUID) -> Optional[Agent]:
|
|
"""Get agent by name within an organization"""
|
|
return self.db.query(Agent).filter(
|
|
Agent.name == name,
|
|
Agent.organization_id == organization_id
|
|
).first()
|
|
|
|
def get_agent(self, agent_id: UUID) -> Optional[Agent]:
|
|
"""Get template by ID"""
|
|
return self.db.query(Agent).filter(
|
|
Agent.id == agent_id
|
|
).first()
|
|
|
|
def get_agent_in_org(self, agent_id, org_id) -> Optional[Agent]:
|
|
"""Get an agent only if it belongs to org_id.
|
|
|
|
The single choke point for "is this agent mine?" — endpoints that take
|
|
an agent_id from the caller use this so an agent from another tenant
|
|
can never be read or reconfigured.
|
|
"""
|
|
return self.db.query(Agent).filter(
|
|
Agent.id == agent_id,
|
|
Agent.organization_id == org_id
|
|
).first()
|
|
|
|
def get_org_agents(self, org_id: UUID, active_only: bool = True) -> List[Agent]:
|
|
"""Get all templates for an organization"""
|
|
query = self.db.query(Agent).filter(
|
|
Agent.organization_id == org_id
|
|
)
|
|
if active_only:
|
|
query = query.filter(Agent.is_active == True)
|
|
return query.all()
|
|
|
|
def get_default_agents(self, org_id: UUID) -> Optional[Agent]:
|
|
"""Get default template for an organization"""
|
|
return self.db.query(Agent).filter(
|
|
Agent.organization_id == org_id,
|
|
Agent.is_default == True,
|
|
Agent.is_active == True
|
|
).first()
|
|
|
|
def update_agent(self, agent_id: UUID, **kwargs) -> Optional[Agent]:
|
|
"""Update an existing template"""
|
|
agent = self.get_agent(agent_id)
|
|
if not agent:
|
|
return None
|
|
|
|
# Handle special fields
|
|
if 'instructions' in kwargs:
|
|
instructions = kwargs['instructions']
|
|
if isinstance(instructions, str):
|
|
try:
|
|
# Try to parse if it's a JSON string
|
|
parsed = json.loads(instructions)
|
|
if isinstance(parsed, list):
|
|
kwargs['instructions'] = parsed
|
|
else:
|
|
# If it's a single instruction, wrap in list
|
|
kwargs['instructions'] = [instructions]
|
|
except json.JSONDecodeError:
|
|
# If not JSON, treat as single instruction
|
|
kwargs['instructions'] = [instructions]
|
|
if 'tools' in kwargs:
|
|
kwargs['tools'] = json.dumps(kwargs['tools'])
|
|
|
|
# Handle default flag
|
|
if kwargs.get('is_default'):
|
|
existing_defaults = self.db.query(Agent).filter(
|
|
Agent.organization_id == agent.organization_id,
|
|
Agent.is_default == True,
|
|
Agent.id != agent_id
|
|
).all()
|
|
for existing in existing_defaults:
|
|
existing.is_default = False
|
|
|
|
for key, value in kwargs.items():
|
|
setattr(agent, key, value)
|
|
|
|
self.db.commit()
|
|
self.db.refresh(agent)
|
|
return agent
|
|
|
|
def delete_agent(self, agent_id: str) -> bool:
|
|
"""Soft delete a template by setting is_active to False"""
|
|
template = self.get_agent(agent_id)
|
|
if not template:
|
|
return False
|
|
|
|
template.is_active = False
|
|
self.db.commit()
|
|
return True
|
|
|
|
def get_active_agents(self, org_id: UUID) -> List[Agent]:
|
|
"""Get the active agent template for an organization"""
|
|
return self.db.query(Agent)\
|
|
.filter(Agent.organization_id == org_id)\
|
|
.filter(Agent.is_active == True)\
|
|
.all()
|
|
|
|
def get_by_agent_id(self, agent_id: str):
|
|
"""Get agent by ID with relationships loaded"""
|
|
if isinstance(agent_id, str):
|
|
agent_id = UUID(agent_id)
|
|
return self.db.query(Agent)\
|
|
.options(
|
|
joinedload(Agent.groups),
|
|
joinedload(Agent.organization)
|
|
)\
|
|
.filter(Agent.id == agent_id)\
|
|
.first()
|
|
|
|
def get_all(self) -> list[Agent]:
|
|
return self.db.query(Agent).all()
|
|
|
|
def get_by_id(self, agent_id: str) -> Agent | None:
|
|
return self.db.query(Agent).filter(Agent.id == agent_id).first()
|
|
|
|
def get_all_agents(self, org_id: UUID) -> List[Agent]:
|
|
"""Get all agents for organization"""
|
|
return self.db.query(Agent)\
|
|
.filter(Agent.organization_id == org_id)\
|
|
.all()
|
|
|
|
def count_by_organization(self, org_id: UUID) -> int:
|
|
"""Count agents for an organization"""
|
|
from sqlalchemy import func
|
|
return self.db.query(func.count(Agent.id))\
|
|
.filter(Agent.organization_id == org_id)\
|
|
.scalar() or 0
|