209 lines
8.4 KiB
Python
209 lines
8.4 KiB
Python
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
from datetime import datetime
|
|
|
|
from sqlalchemy import Boolean, DateTime, Integer, String, Text, create_engine, select
|
|
from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column, sessionmaker
|
|
|
|
from app.config import get_settings
|
|
|
|
|
|
DEFAULT_TOOLS = {
|
|
"brave_search": True,
|
|
"second_brain": True,
|
|
"browser_use": True,
|
|
"searxng_search": False,
|
|
"web_fetch": True,
|
|
"apple_notes": True,
|
|
"files": True,
|
|
"terminal": True,
|
|
}
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
|
|
class SettingORM(Base):
|
|
__tablename__ = "settings"
|
|
|
|
key: Mapped[str] = mapped_column(String(100), primary_key=True)
|
|
value: Mapped[str] = mapped_column(Text, nullable=False)
|
|
updated_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)
|
|
|
|
|
|
class ToolStateORM(Base):
|
|
__tablename__ = "tool_states"
|
|
|
|
name: Mapped[str] = mapped_column(String(100), primary_key=True)
|
|
enabled: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)
|
|
updated_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)
|
|
|
|
|
|
class AuthorizedUserORM(Base):
|
|
__tablename__ = "authorized_users"
|
|
|
|
telegram_user_id: Mapped[int] = mapped_column(Integer, primary_key=True)
|
|
username: Mapped[str | None] = mapped_column(String(255))
|
|
display_name: Mapped[str | None] = mapped_column(String(255))
|
|
is_active: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)
|
|
created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)
|
|
updated_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)
|
|
|
|
|
|
class MemoryItemORM(Base):
|
|
__tablename__ = "memory_items"
|
|
|
|
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
|
content: Mapped[str] = mapped_column(Text, nullable=False)
|
|
kind: Mapped[str] = mapped_column(String(50), nullable=False, default="message")
|
|
created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)
|
|
|
|
|
|
class AuditLogORM(Base):
|
|
__tablename__ = "audit_logs"
|
|
|
|
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
|
category: Mapped[str] = mapped_column(String(50), nullable=False)
|
|
message: Mapped[str] = mapped_column(Text, nullable=False)
|
|
created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)
|
|
|
|
|
|
class SecretORM(Base):
|
|
__tablename__ = "secrets"
|
|
|
|
key: Mapped[str] = mapped_column(String(100), primary_key=True)
|
|
value: Mapped[str] = mapped_column(Text, nullable=False)
|
|
updated_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)
|
|
|
|
|
|
class TelegramUserProfileORM(Base):
|
|
__tablename__ = "telegram_user_profiles"
|
|
|
|
telegram_user_id: Mapped[int] = mapped_column(Integer, primary_key=True)
|
|
display_name: Mapped[str | None] = mapped_column(String(255))
|
|
bio: Mapped[str | None] = mapped_column(Text)
|
|
occupation: Mapped[str | None] = mapped_column(String(255))
|
|
primary_use_cases: Mapped[str] = mapped_column(Text, nullable=False, default="[]")
|
|
answer_priorities: Mapped[str] = mapped_column(Text, nullable=False, default="[]")
|
|
tone_preference: Mapped[str | None] = mapped_column(String(100))
|
|
response_length: Mapped[str | None] = mapped_column(String(50))
|
|
language_preference: Mapped[str | None] = mapped_column(String(100))
|
|
workflow_preference: Mapped[str | None] = mapped_column(String(100))
|
|
interests: Mapped[str] = mapped_column(Text, nullable=False, default="[]")
|
|
approval_preferences: Mapped[str] = mapped_column(Text, nullable=False, default="[]")
|
|
avoid_preferences: Mapped[str | None] = mapped_column(Text)
|
|
onboarding_completed: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
|
|
last_onboarding_step: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
|
created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)
|
|
updated_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)
|
|
|
|
|
|
class AutomationORM(Base):
|
|
__tablename__ = "automations"
|
|
|
|
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
|
telegram_user_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True)
|
|
name: Mapped[str] = mapped_column(String(255), nullable=False)
|
|
prompt: Mapped[str] = mapped_column(Text, nullable=False)
|
|
schedule_type: Mapped[str] = mapped_column(String(50), nullable=False)
|
|
interval_hours: Mapped[int | None] = mapped_column(Integer)
|
|
time_of_day: Mapped[str | None] = mapped_column(String(20))
|
|
days_of_week: Mapped[str] = mapped_column(Text, nullable=False, default="[]")
|
|
status: Mapped[str] = mapped_column(String(20), nullable=False, default="active")
|
|
last_run_at: Mapped[datetime | None] = mapped_column(DateTime)
|
|
next_run_at: Mapped[datetime | None] = mapped_column(DateTime)
|
|
last_result: Mapped[str | None] = mapped_column(Text)
|
|
created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)
|
|
updated_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)
|
|
|
|
|
|
class AutomationWizardORM(Base):
|
|
__tablename__ = "automation_wizards"
|
|
|
|
telegram_user_id: Mapped[int] = mapped_column(Integer, primary_key=True)
|
|
step: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
|
draft_json: Mapped[str] = mapped_column(Text, nullable=False, default="{}")
|
|
created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)
|
|
updated_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)
|
|
|
|
|
|
class SecondBrainNoteORM(Base):
|
|
__tablename__ = "second_brain_notes"
|
|
|
|
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
|
telegram_user_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True)
|
|
content: Mapped[str] = mapped_column(Text, nullable=False)
|
|
source: Mapped[str] = mapped_column(String(50), nullable=False, default="telegram")
|
|
created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)
|
|
updated_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)
|
|
|
|
|
|
class SecondBrainCaptureORM(Base):
|
|
__tablename__ = "second_brain_captures"
|
|
|
|
telegram_user_id: Mapped[int] = mapped_column(Integer, primary_key=True)
|
|
created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)
|
|
updated_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)
|
|
|
|
|
|
settings = get_settings()
|
|
DEFAULT_SETTINGS = {
|
|
"terminal_mode": "3",
|
|
"search_provider": settings.search_provider,
|
|
"model_provider": settings.model_provider,
|
|
"local_base_url": settings.local_base_url,
|
|
"local_model": settings.local_model,
|
|
"zai_model": settings.zai_model,
|
|
"anythingllm_base_url": settings.anythingllm_base_url,
|
|
"anythingllm_workspace_slug": settings.anythingllm_workspace_slug,
|
|
}
|
|
|
|
engine = create_engine(
|
|
settings.db_url,
|
|
connect_args={"check_same_thread": False} if settings.db_url.startswith("sqlite") else {},
|
|
)
|
|
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False)
|
|
|
|
|
|
def init_db() -> None:
|
|
Base.metadata.create_all(bind=engine)
|
|
with session_scope() as session:
|
|
_seed_defaults(session)
|
|
|
|
|
|
def _seed_defaults(session: Session) -> None:
|
|
for key, value in DEFAULT_SETTINGS.items():
|
|
if session.get(SettingORM, key) is None:
|
|
session.add(SettingORM(key=key, value=value))
|
|
|
|
for name, enabled in DEFAULT_TOOLS.items():
|
|
if session.get(ToolStateORM, name) is None:
|
|
session.add(ToolStateORM(name=name, enabled=enabled))
|
|
|
|
|
|
def get_session() -> Iterator[Session]:
|
|
session = SessionLocal()
|
|
try:
|
|
yield session
|
|
finally:
|
|
session.close()
|
|
|
|
|
|
@contextmanager
|
|
def session_scope() -> Iterator[Session]:
|
|
session = SessionLocal()
|
|
try:
|
|
yield session
|
|
session.commit()
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
session.close()
|
|
|
|
|
|
def list_recent_logs(session: Session, limit: int = 10) -> list[str]:
|
|
stmt = select(AuditLogORM).order_by(AuditLogORM.created_at.desc(), AuditLogORM.id.desc()).limit(limit)
|
|
return [row.message for row in session.scalars(stmt)]
|