devplacepy/devplacepy/services/devii/hub.py

# retoor <retoor@molodetz.nl>

from __future__ import annotations

import logging
from typing import Any, Callable

from devplacepy.database import db

from .agentic import LessonStore
from .behavior import BehaviorStore
from .config import Settings
from .cost.tracker import Pricing
from .llm import LLMClient
from .session import DeviiSession
from .store import ConversationStore, TurnAudit, UsageLedger
from .tasks.store import TaskStore, memory_db
from .virtual_tools import VirtualToolStore

logger = logging.getLogger("devii.hub")

SettingsBuilder = Callable[[str, str, str, bool], "tuple[Settings, Pricing]"]


class DeviiHub:
    def __init__(self, settings_builder: SettingsBuilder) -> None:
        self._build = settings_builder
        self._stores: dict[str, Any] = {
            "conversations": ConversationStore(),
            "ledger": UsageLedger(),
            "turns": TurnAudit(),
        }
        self._sessions: dict[tuple[str, str, str], DeviiSession] = {}

    @property
    def ledger(self) -> UsageLedger:
        return self._stores["ledger"]

    def get_or_create(
        self,
        owner_kind: str,
        owner_id: str,
        username: str,
        api_key: str,
        base_url: str,
        is_admin: bool = False,
        is_primary_admin: bool = False,
        channel: str = "main",
    ) -> DeviiSession:
        key = (owner_kind, owner_id, channel)
        session = self._sessions.get(key)
        if session is not None:
            return session
        settings, pricing = self._build(api_key, base_url, owner_kind, is_admin)
        llm = LLMClient(settings)
        # Persistent, owner-isolated stores for signed-in users; ephemeral in-memory for guests.
        # The `docs` channel (Docii) is a self-contained documentation assistant: it gets its own
        # ephemeral stores so Devii's tasks, lessons, behavior and virtual tools never leak into it
        # (and a Docii reflection never pollutes the user's Devii memory). Only the conversation
        # thread persists, keyed per channel in the shared ConversationStore.
        owned_db = (
            db
            if (owner_kind == "user" and channel in ("main", "telegram"))
            else memory_db()
        )
        task_store = TaskStore(owned_db, owner_kind, owner_id)
        lessons = LessonStore(owned_db, owner_kind, owner_id)
        virtual_tool_store = VirtualToolStore(owned_db, owner_kind, owner_id)
        behavior_store = BehaviorStore(owned_db, owner_kind, owner_id)
        session = DeviiSession(
            owner_kind,
            owner_id,
            username,
            settings,
            llm,
            lessons,
            pricing,
            task_store,
            virtual_tool_store,
            behavior_store,
            self._stores,
            is_admin=is_admin,
            is_primary_admin=is_primary_admin,
            channel=channel,
        )
        if owner_kind == "user":
            saved = self._stores["conversations"].load(owner_kind, owner_id, channel)
            if saved:
                session.restore_history(saved)
        self._sessions[key] = session
        logger.info(
            "Created session %s/%s [%s] (total %d)",
            owner_kind,
            owner_id,
            channel,
            len(self._sessions),
        )
        return session

    def find(
        self, owner_kind: str, owner_id: str, channel: str = "main"
    ) -> DeviiSession | None:
        return self._sessions.get((owner_kind, owner_id, channel))

    def active_sessions(self) -> int:
        return len(self._sessions)

    def connection_count(self) -> int:
        return sum(s.connection_count for s in self._sessions.values())

    async def gc_idle(self) -> int:
        removed = 0
        for key, session in list(self._sessions.items()):
            if session.connection_count == 0 and not session.has_pending_tasks():
                await session.aclose()
                del self._sessions[key]
                removed += 1
        if removed:
            logger.info("GC removed %d idle session(s)", removed)
        return removed

    async def aclose(self) -> None:
        for session in list(self._sessions.values()):
            await session.aclose()
        self._sessions.clear()