feat(models): 组织、用户、角色、客户、会话与消息模型
按组织边界定义数据关系，补齐仓储层与通用工具，迁移脚本可重复执行。
This commit is contained in:
WangShen007 committed
1 parent
977783c126
commit
9eb8613c05
126 files changed
+17436
No files matched your search
@@ -0,0 +1,214 @@
|
||||
"""
|
||||
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 datetime import datetime, timedelta, timezone
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
from sqlalchemy import and_, or_
|
||||
from app.models.knowledge import Knowledge
|
||||
from typing import List, Optional
|
||||
from sqlalchemy import func
|
||||
from sqlalchemy.sql import text
|
||||
import logging
|
||||
from uuid import UUID
|
||||
from app.models.knowledge_to_agent import KnowledgeToAgent
|
||||
from app.models.knowledge_queue import KnowledgeQueue, QueueStatus
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class KnowledgeRepository:
|
||||
def __init__(self, db: Session):
|
||||
self.db = db
|
||||
|
||||
def get_by_id(self, knowledge_id: int) -> Optional[Knowledge]:
|
||||
"""Get knowledge source by ID"""
|
||||
return self.db.query(Knowledge)\
|
||||
.filter(Knowledge.id == knowledge_id)\
|
||||
.first()
|
||||
|
||||
def get_by_agent(self, agent_id: UUID, skip: int = 0, limit: int = 100):
|
||||
"""Get paginated knowledge items for an agent"""
|
||||
|
||||
query = (self.db.query(Knowledge)
|
||||
.join(KnowledgeToAgent)
|
||||
.filter(
|
||||
KnowledgeToAgent.knowledge_id == Knowledge.id,
|
||||
KnowledgeToAgent.agent_id == agent_id
|
||||
)
|
||||
.order_by(Knowledge.id)
|
||||
.offset(skip)
|
||||
.limit(limit))
|
||||
|
||||
result = query.all()
|
||||
|
||||
return result
|
||||
|
||||
def count_by_agent(self, agent_id: UUID) -> int:
|
||||
"""Get total count of knowledge items for an agent"""
|
||||
|
||||
query = (self.db.query(func.count(Knowledge.id))
|
||||
.join(KnowledgeToAgent)
|
||||
.filter(
|
||||
KnowledgeToAgent.knowledge_id == Knowledge.id,
|
||||
KnowledgeToAgent.agent_id == agent_id
|
||||
))
|
||||
|
||||
result = query.scalar() or 0
|
||||
return result
|
||||
|
||||
def get_by_org(self, org_id: UUID) -> List[Knowledge]:
|
||||
"""Get all knowledge sources for an organization"""
|
||||
return self.db.query(Knowledge)\
|
||||
.filter(Knowledge.organization_id == org_id)\
|
||||
.all()
|
||||
|
||||
def create(self, knowledge: Knowledge) -> Knowledge:
|
||||
"""Create a new knowledge source or get existing"""
|
||||
existing = self.db.query(Knowledge)\
|
||||
.filter(
|
||||
Knowledge.organization_id == knowledge.organization_id,
|
||||
Knowledge.source == knowledge.source
|
||||
).first()
|
||||
|
||||
if existing:
|
||||
return existing
|
||||
|
||||
self.db.add(knowledge)
|
||||
self.db.commit()
|
||||
self.db.refresh(knowledge)
|
||||
return knowledge
|
||||
|
||||
def delete(self, knowledge_id: int) -> bool:
|
||||
"""Delete a knowledge source"""
|
||||
knowledge = self.db.query(Knowledge)\
|
||||
.filter(Knowledge.id == knowledge_id)\
|
||||
.first()
|
||||
if knowledge:
|
||||
self.db.delete(knowledge)
|
||||
self.db.commit()
|
||||
return True
|
||||
return False
|
||||
|
||||
def count_by_organization(self, org_id: UUID) -> int:
|
||||
"""Get total count of indexed knowledge items for an organization"""
|
||||
|
||||
query = (self.db.query(func.count(Knowledge.id))
|
||||
.filter(Knowledge.organization_id == org_id))
|
||||
|
||||
result = query.scalar() or 0
|
||||
return result
|
||||
|
||||
# A crawl occupies its plan slot the moment it is queued, not when it
|
||||
# finishes. Counting only indexed rows let an organization submit its
|
||||
# whole backlog in one sitting - every request saw a count of zero - and
|
||||
# end up over the plan limit once the queue drained.
|
||||
IN_FLIGHT_QUEUE_STATUSES = (QueueStatus.PENDING.value, QueueStatus.PROCESSING.value)
|
||||
# The worker only ever picks up PENDING, so a row left PROCESSING by a
|
||||
# restart is never retried and never completes. Stop charging for one
|
||||
# after this long, or a redeploy mid-crawl would cost a slot for ever.
|
||||
STALE_PROCESSING_AFTER = timedelta(hours=6)
|
||||
|
||||
def count_sources_in_use(self, org_id: UUID, now: datetime = None) -> int:
|
||||
"""Knowledge sources occupying a plan slot: the indexed ones plus the
|
||||
crawls still queued or running.
|
||||
|
||||
Every queue row becomes its own Knowledge row - the worker does not
|
||||
merge repeats - so rows are counted, not distinct sources: otherwise
|
||||
submitting one URL repeatedly walked past the limit unchecked. A row
|
||||
that has just written its Knowledge row but is not yet marked complete
|
||||
is counted twice for that instant, which errs toward refusing, the
|
||||
safe direction for a limit."""
|
||||
now = now or datetime.now(timezone.utc)
|
||||
fresh_enough = now - self.STALE_PROCESSING_AFTER
|
||||
in_flight = (
|
||||
self.db.query(func.count(KnowledgeQueue.id))
|
||||
.filter(
|
||||
KnowledgeQueue.organization_id == org_id,
|
||||
or_(
|
||||
KnowledgeQueue.status == QueueStatus.PENDING.value,
|
||||
and_(
|
||||
KnowledgeQueue.status == QueueStatus.PROCESSING.value,
|
||||
func.coalesce(
|
||||
KnowledgeQueue.updated_at, KnowledgeQueue.created_at
|
||||
) >= fresh_enough,
|
||||
),
|
||||
),
|
||||
)
|
||||
.scalar() or 0
|
||||
)
|
||||
return self.count_by_organization(org_id) + in_flight
|
||||
|
||||
def get_by_organization(
|
||||
self,
|
||||
org_id: UUID,
|
||||
skip: int = 0,
|
||||
limit: int = 10
|
||||
) -> List[Knowledge]:
|
||||
"""Get paginated knowledge items for an organization"""
|
||||
|
||||
query = (self.db.query(Knowledge)
|
||||
.filter(Knowledge.organization_id == org_id)
|
||||
.order_by(Knowledge.id)
|
||||
.offset(skip)
|
||||
.limit(limit))
|
||||
|
||||
result = query.all()
|
||||
return result
|
||||
|
||||
def get_by_sources(self, organization_id: UUID, sources: List[str]):
|
||||
"""Get knowledge items by source URLs for an organization"""
|
||||
return self.db.query(Knowledge)\
|
||||
.filter(
|
||||
Knowledge.organization_id == organization_id,
|
||||
Knowledge.source.in_(sources)
|
||||
).all()
|
||||
|
||||
def delete_with_data(self, knowledge_id: int) -> bool:
|
||||
"""Delete knowledge source and its associated data"""
|
||||
try:
|
||||
knowledge = self.get_by_id(knowledge_id)
|
||||
if not knowledge:
|
||||
return False
|
||||
|
||||
# Delete from knowledge_to_agent table first
|
||||
self.db.execute(
|
||||
text("DELETE FROM knowledge_to_agents WHERE knowledge_id = :kid"),
|
||||
{"kid": knowledge_id}
|
||||
)
|
||||
|
||||
# If table_name and schema exist, delete the actual data
|
||||
if knowledge.table_name and knowledge.schema:
|
||||
try:
|
||||
# Delete data from the specific table
|
||||
self.db.execute(
|
||||
text(f'DELETE FROM {knowledge.schema}."{ knowledge.table_name}" WHERE name = :source'),
|
||||
{"source": knowledge.source}
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting data from {
|
||||
knowledge.table_name}: {str(e)}")
|
||||
# Continue with knowledge deletion even if data deletion fails
|
||||
|
||||
# Delete the knowledge entry
|
||||
self.db.delete(knowledge)
|
||||
self.db.commit()
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
self.db.rollback()
|
||||
logger.error(f"Error in delete_with_data: {str(e)}")
|
||||
raise
|
||||
Reference in new issue
Block a user