company-ai-platform/app/modules/personalization/services/context.py

from dataclasses import dataclass, field
from typing import Any

from sqlalchemy import select
from sqlalchemy.orm import Session

from app.core.constants import ActorValue
from app.modules.ai_memory.constants import AIMemoryScope
from app.modules.ai_memory.service import AIMemoryService
from app.modules.business.models import MarketWatchlist
from app.modules.personalization.constants import (
    PersonalizationContextKey,
    PreferenceCategory,
)
from app.modules.personalization.services.conversations import ConversationService
from app.modules.personalization.services.preferences import PreferenceService


@dataclass(frozen=True)
class PersonalizationContext:
    """Ordered, provider-neutral context sections for one AI request."""

    system_constraints: str
    company_rules: list[dict[str, Any]]
    personal_rules: list[dict[str, Any]]
    current_request: str
    preferences: list[dict[str, Any]]
    interests: list[dict[str, Any]]
    personal_memory: list[dict[str, Any]]
    conversation_history: list[dict[str, Any]]
    provider_session_id: str | None = field(default=None)

    def as_ordered_dict(self) -> dict[str, Any]:
        return {
            PersonalizationContextKey.SYSTEM_CONSTRAINTS: self.system_constraints,
            PersonalizationContextKey.COMPANY_RULES: self.company_rules,
            PersonalizationContextKey.PERSONAL_RULES: self.personal_rules,
            PersonalizationContextKey.CURRENT_REQUEST: self.current_request,
            PersonalizationContextKey.PREFERENCES: self.preferences,
            PersonalizationContextKey.INTERESTS: self.interests,
            PersonalizationContextKey.PERSONAL_MEMORY: self.personal_memory,
            PersonalizationContextKey.CONVERSATION_HISTORY: self.conversation_history,
        }


class PersonalizationContextService:
    """Load only the explicitly requested company and owner-scoped context layers."""

    def __init__(self, db: Session):
        self.db = db
        self.memory = AIMemoryService(db)
        self.preferences = PreferenceService(db)
        self.conversations = ConversationService(db)

    def build(
        self,
        *,
        owner_id: int | None,
        request: str,
        system_constraints: str,
        chat_type: str | None = None,
        chat_key: str | None = None,
        scope: str = AIMemoryScope.GLOBAL,
        subject: str | None = None,
        actor: str = ActorValue.SYSTEM,
        include_company_rules: bool = True,
        include_personal_context: bool = True,
        include_history: bool = True,
    ) -> PersonalizationContext:
        company_rules = (
            self.memory.active_rules(
                scope=scope,
                subject=subject,
                owner_id=None,
            )
            if include_company_rules
            else []
        )
        personal_rules: list[dict[str, Any]] = []
        preference_items: list[dict[str, Any]] = []
        interests: list[dict[str, Any]] = []
        personal_memory: list[dict[str, Any]] = []
        history: list[dict[str, Any]] = []
        provider_session_id: str | None = None

        if owner_id is not None and include_personal_context:
            personal_rules = self.memory.active_rules(
                scope=scope,
                subject=subject,
                owner_id=owner_id,
            )
            all_preferences = self.preferences.list_preferences(owner_id)
            interest_categories = {
                PreferenceCategory.TOPIC,
                PreferenceCategory.INTEREST,
            }
            for preference in all_preferences:
                if preference["category"] in interest_categories:
                    interests.append(preference)
                else:
                    preference_items.append(preference)
            interests.extend(self._watchlist_interests(owner_id))
            personal_memory = self.memory.recall(
                query=request,
                scope=scope,
                subject=subject,
                actor=actor,
                owner_id=owner_id,
            )
            if include_history and chat_type and chat_key:
                history = self.conversations.history(owner_id, chat_type, chat_key)
                provider_session_id = self.conversations.provider_session_id(
                    owner_id,
                    chat_type,
                    chat_key,
                )

        return PersonalizationContext(
            system_constraints=system_constraints,
            company_rules=company_rules,
            personal_rules=personal_rules,
            current_request=request,
            preferences=preference_items,
            interests=interests,
            personal_memory=personal_memory,
            conversation_history=history,
            provider_session_id=provider_session_id,
        )

    def build_private_scheduled(
        self,
        *,
        owner_id: int,
        request: str,
        system_constraints: str,
        scope: str = AIMemoryScope.GLOBAL,
        subject: str | None = None,
        actor: str = ActorValue.SYSTEM,
    ) -> PersonalizationContext:
        """Build private scheduled context without company data or conversation history."""

        return self.build(
            owner_id=owner_id,
            request=request,
            system_constraints=system_constraints,
            scope=scope,
            subject=subject,
            actor=actor,
            include_company_rules=False,
            include_personal_context=True,
            include_history=False,
        )

    def build_group_scheduled(
        self,
        *,
        request: str,
        system_constraints: str,
        scope: str = AIMemoryScope.GLOBAL,
        subject: str | None = None,
    ) -> PersonalizationContext:
        """Build group scheduled context without any creator profile."""

        return self.build(
            owner_id=None,
            request=request,
            system_constraints=system_constraints,
            scope=scope,
            subject=subject,
            include_company_rules=True,
            include_personal_context=False,
            include_history=False,
        )

    def _watchlist_interests(self, owner_id: int) -> list[dict[str, Any]]:
        symbols = self.db.execute(
            select(MarketWatchlist.symbol)
            .where(
                MarketWatchlist.owner_id == owner_id,
                MarketWatchlist.enabled.is_(True),
            )
            .order_by(MarketWatchlist.symbol.asc())
        ).scalars()
        return [
            {
                "category": "watchlist",
                "value": symbol,
                "source": "market_watchlist",
            }
            for symbol in symbols
        ]