214 lines
7.8 KiB
Python
214 lines
7.8 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 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
|