company-ai-platform/app/modules/ai_agent/service.py

from typing import Any

from fastapi import HTTPException, status
from sqlalchemy.orm import Session

from app.core.constants import ActorValue
from app.core.config import get_settings
from app.modules.ai_agent.adapters import HermesAdapter, OpenClawAdapter, get_adapter
from app.modules.ai_agent.constants import (
    AI_UNAVAILABLE_ANSWER,
    AI_AUDIT_MAX_DEPTH,
    AI_AUDIT_MAX_SEQUENCE_ITEMS,
    AI_AUDIT_MAX_TEXT_LENGTH,
    AI_AUDIT_REDACTED_VALUE,
    AI_AUDIT_SENSITIVE_KEYS,
    AI_AUDIT_TRUNCATED_VALUE,
    COMPANY_MANAGEMENT_SYSTEM_INSTRUCTIONS,
    PREFERENCE_EXTRACTION_INSTRUCTIONS,
    AIContextKey,
    AIExecutionMode,
    AIProviderName,
    AIRequestKey,
    AIResponseKey,
)
from app.modules.ai_memory.constants import AIMemoryPayloadKey, AIMemoryScope
from app.modules.ai_memory.service import AIMemoryService
from app.modules.ai_agent.skills import AISkillId, get_ai_skill
from app.modules.audit.constants import (
    AuditAction,
    AuditRiskLevel,
    AuditSource,
    AuditTargetType,
)
from app.modules.audit.schemas import AuditLogCreate
from app.modules.audit.service import AuditService
from app.modules.personalization.services import (
    ConversationService,
    PersonalizationContext,
    PersonalizationContextService,
    PreferenceService,
)
from app.modules.personalization.services.preferences import contains_preference_signal


class AIService:
    """Coordinate AI provider calls and audit logging."""

    def __init__(self, db: Session):
        self.db = db
        self.audit = AuditService(db)

    def ask(
        self,
        prompt: str,
        context: dict[str, Any] | None = None,
        actor: str = ActorValue.API,
        source: str = AuditSource.API,
    ) -> dict[str, Any]:
        request_context = dict(context or {})
        _validate_external_context(request_context)
        adapter = get_adapter()
        if _is_unavailable_adapter(adapter):
            response = _unavailable_response(adapter.provider_name)
            self._audit_ai_response(
                actor=actor,
                source=source,
                request_payload={
                    AIRequestKey.PROMPT: prompt,
                    AIRequestKey.CONTEXT: request_context,
                },
                response=response,
            )
            return response

        adapter_context = dict(request_context)
        memory_service = AIMemoryService(self.db)
        memory_scope = _memory_scope(request_context)
        memory_subject = _memory_subject(request_context)
        user_rules = memory_service.active_rules(
            scope=memory_scope,
            subject=memory_subject,
            owner_id=None,
        )
        if user_rules:
            adapter_context[AIContextKey.USER_RULES] = user_rules
        local_memory = memory_service.recall(
            query=prompt,
            scope=memory_scope,
            subject=memory_subject,
            actor=actor,
            owner_id=None,
        )
        if local_memory:
            adapter_context[AIContextKey.LOCAL_MEMORY] = local_memory
        result = adapter.ask(prompt, adapter_context)
        answer = result[AIResponseKey.ANSWER]
        raw = dict(result.get(AIResponseKey.RAW, {}))
        if local_memory:
            raw[AIResponseKey.LOCAL_MEMORY] = local_memory
        memory_record = memory_service.auto_write(
            prompt=prompt,
            context=request_context,
            answer=answer,
            actor=actor,
            owner_id=None,
        )
        if memory_record is not None:
            raw[AIResponseKey.MEMORY_WRITE] = {
                AIMemoryPayloadKey.CODE: memory_record.code,
                AIMemoryPayloadKey.STATUS: memory_record.status,
            }
        response = {
            AIResponseKey.OK: True,
            AIResponseKey.PROVIDER: adapter.provider_name,
            AIResponseKey.ANSWER: answer,
            AIResponseKey.RAW: raw,
        }
        self._audit_ai_response(
            actor=actor,
            source=source,
            request_payload={
                AIRequestKey.PROMPT: prompt,
                AIRequestKey.CONTEXT: request_context,
            },
            response=response,
        )
        return response

    def ask_personalized(
        self,
        owner_id: int,
        chat_type: str,
        chat_key: str,
        prompt: str,
        actor: str = ActorValue.FEISHU,
        source: str = AuditSource.FEISHU,
        scope: str = AIMemoryScope.USER,
        subject: str | None = None,
    ) -> dict[str, Any]:
        """Answer with one verified owner's isolated personalization context."""

        _validate_owner_id(owner_id)
        _validate_prompt(prompt)
        session_id = ConversationService.provider_session_id(
            owner_id,
            chat_type,
            chat_key,
        )
        adapter = get_adapter()
        if _is_unavailable_adapter(adapter):
            response = _unavailable_response(adapter.provider_name)
            self._audit_ai_response(
                actor=actor,
                source=source,
                request_payload={
                    "owner_id": owner_id,
                    "chat_type": chat_type,
                    AIContextKey.EXECUTION_MODE: AIExecutionMode.PERSONALIZED,
                },
                response=response,
            )
            return response

        memory_subject = subject or f"owner:{owner_id}"
        personalization = PersonalizationContextService(self.db).build(
            owner_id=owner_id,
            request=prompt,
            system_constraints=COMPANY_MANAGEMENT_SYSTEM_INSTRUCTIONS,
            chat_type=chat_type,
            chat_key=chat_key,
            scope=scope,
            subject=memory_subject,
            actor=actor,
            include_company_rules=True,
            include_personal_context=True,
            include_history=True,
        )
        adapter_context = _adapter_context(
            personalization,
            execution_mode=AIExecutionMode.PERSONALIZED,
            allow_provider_memory=False,
            provider_session_id=session_id,
        )
        result = adapter.ask(prompt, adapter_context)
        answer = _required_answer(result)
        raw = dict(result.get(AIResponseKey.RAW, {}))

        ConversationService(self.db).record_turn(
            owner_id,
            chat_type,
            chat_key,
            user_content=prompt,
            assistant_content=answer,
            provider_name=str(adapter.provider_name),
        )
        memory_record = AIMemoryService(self.db).auto_write(
            prompt=prompt,
            context={
                AIMemoryPayloadKey.SCOPE: scope,
                AIMemoryPayloadKey.SUBJECT: memory_subject,
            },
            answer=answer,
            actor=actor,
            owner_id=owner_id,
        )
        if memory_record is not None:
            raw[AIResponseKey.MEMORY_WRITE] = {
                AIMemoryPayloadKey.CODE: memory_record.code,
                AIMemoryPayloadKey.STATUS: memory_record.status,
            }
        saved_preferences = self._extract_preferences(
            adapter=adapter,
            owner_id=owner_id,
            user_text=prompt,
            provider_session_id=session_id,
        )
        if saved_preferences:
            raw["preferences_saved"] = [
                {
                    "code": item["code"],
                    "category": item["category"],
                }
                for item in saved_preferences
            ]

        response = {
            AIResponseKey.OK: True,
            AIResponseKey.PROVIDER: adapter.provider_name,
            AIResponseKey.ANSWER: answer,
            AIResponseKey.RAW: raw,
        }
        self._audit_ai_response(
            actor=actor,
            source=source,
            request_payload={
                "owner_id": owner_id,
                "chat_type": chat_type,
                AIContextKey.EXECUTION_MODE: AIExecutionMode.PERSONALIZED,
            },
            response={
                AIResponseKey.OK: True,
                AIResponseKey.PROVIDER: adapter.provider_name,
            },
        )
        return response

    def generate_scheduled(
        self,
        prompt: str,
        owner_id: int | None,
        group: bool,
        actor: str = "subscription-system",
    ) -> dict[str, Any]:
        """Generate side-effect-free scheduled content within its target boundary."""

        _validate_prompt(prompt)
        if not group:
            if owner_id is None:
                raise HTTPException(
                    status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                    detail="Private scheduled generation requires an owner",
                )
            _validate_owner_id(owner_id)

        adapter = get_adapter()
        execution_mode = (
            AIExecutionMode.SCHEDULED_GROUP
            if group
            else AIExecutionMode.SCHEDULED_PRIVATE
        )
        if _is_unavailable_adapter(adapter):
            response = _unavailable_response(adapter.provider_name)
            self._audit_ai_response(
                actor=actor,
                source=AuditSource.FEISHU,
                request_payload={AIContextKey.EXECUTION_MODE: execution_mode},
                response=response,
            )
            return response

        context_service = PersonalizationContextService(self.db)
        if group:
            personalization = context_service.build_group_scheduled(
                request=prompt,
                system_constraints=COMPANY_MANAGEMENT_SYSTEM_INSTRUCTIONS,
            )
        else:
            personalization = context_service.build_private_scheduled(
                owner_id=owner_id,
                request=prompt,
                system_constraints=COMPANY_MANAGEMENT_SYSTEM_INSTRUCTIONS,
                actor=actor,
                scope=AIMemoryScope.USER,
                subject=f"owner:{owner_id}",
            )
        adapter_context = _adapter_context(
            personalization,
            execution_mode=execution_mode,
            allow_provider_memory=False,
            provider_session_id=None,
        )
        result = adapter.ask(prompt, adapter_context)
        response = {
            AIResponseKey.OK: True,
            AIResponseKey.PROVIDER: adapter.provider_name,
            AIResponseKey.ANSWER: _required_answer(result),
            AIResponseKey.RAW: dict(result.get(AIResponseKey.RAW, {})),
        }
        self._audit_ai_response(
            actor=actor,
            source=AuditSource.FEISHU,
            request_payload={AIContextKey.EXECUTION_MODE: execution_mode},
            response={
                AIResponseKey.OK: True,
                AIResponseKey.PROVIDER: adapter.provider_name,
            },
        )
        return response

    def _extract_preferences(
        self,
        *,
        adapter: Any,
        owner_id: int,
        user_text: str,
        provider_session_id: str,
    ) -> list[dict[str, Any]]:
        if not contains_preference_signal(user_text):
            return []
        extraction_context = {
            "preference_source_text": user_text,
            AIContextKey.PROVIDER_SESSION_ID: f"{provider_session_id}-preferences",
            AIContextKey.ALLOW_PROVIDER_MEMORY: False,
            AIContextKey.EXECUTION_MODE: AIExecutionMode.PREFERENCE_EXTRACTION,
        }
        try:
            extraction = adapter.ask(
                PREFERENCE_EXTRACTION_INSTRUCTIONS,
                extraction_context,
            )
            structured_payload = extraction.get(AIResponseKey.ANSWER, "")
            return PreferenceService(self.db).save_auto_extraction(
                owner_id,
                provider_name=str(adapter.provider_name),
                user_text=user_text,
                structured_payload=structured_payload,
            )
        except Exception:
            # Preference inference must never block or change the primary answer.
            return []

    def _audit_ai_response(
        self,
        *,
        actor: str,
        source: str,
        request_payload: dict[str, Any],
        response: dict[str, Any],
    ) -> None:
        self.audit.log(
            AuditLogCreate(
                actor=actor,
                source=source,
                action=AuditAction.AI_ASK,
                target_type=AuditTargetType.AI,
                risk_level=AuditRiskLevel.MEDIUM,
                request_payload=_audit_safe_payload(request_payload),
                response_payload=_audit_safe_payload(response),
            )
        )

    def run_skill(
        self,
        skill_id: AISkillId | str,
        context: dict[str, Any] | None = None,
        variables: dict[str, Any] | None = None,
        actor: str = ActorValue.API,
    ) -> dict[str, Any]:
        skill = get_ai_skill(skill_id)
        return self.ask(
            skill.render(variables),
            context=context or {},
            actor=actor,
            source=skill.source,
        )

    def provider_health(self, actor: str = ActorValue.API) -> dict[str, Any]:
        settings = get_settings()
        openclaw = self._health_result(OpenClawAdapter(settings).health)
        hermes = self._health_result(HermesAdapter(settings).health)
        response = {
            "model_provider": settings.model_provider,
            AIProviderName.OPENCLAW: openclaw,
            AIProviderName.HERMES: hermes,
        }
        self.audit.log(
            AuditLogCreate(
                actor=actor,
                source=AuditSource.API,
                action=AuditAction.AI_PROVIDER_HEALTH,
                target_type=AuditTargetType.AI,
                risk_level=AuditRiskLevel.LOW,
                response_payload=_audit_safe_payload(response),
            )
        )
        return response

    @staticmethod
    def _health_result(check: Any) -> dict[str, Any]:
        try:
            return check()
        except Exception as exc:
            # Health checks should report failures, not mask the other provider.
            return {
                AIResponseKey.OK: False,
                AIResponseKey.ERROR: str(exc),
                AIResponseKey.TYPE: type(exc).__name__,
            }

    def draft_policy(
        self,
        title: str,
        policy_type: str,
        requirements: list[str],
        actor: str,
    ) -> dict[str, Any]:
        return self.run_skill(
            AISkillId.DRAFT_POLICY,
            variables={
                "title": title,
                "policy_type": policy_type,
                "requirements": requirements,
            },
            actor=actor,
        )

    def draft_investment_research(
        self,
        symbol_or_topic: str,
        risk_preference: str,
        actor: str,
    ) -> dict[str, Any]:
        return self.run_skill(
            AISkillId.INVESTMENT_RESEARCH,
            variables={
                "symbol_or_topic": symbol_or_topic,
                "risk_preference": risk_preference,
            },
            actor=actor,
        )


def _audit_safe_payload(value: Any, depth: int = 0) -> Any:
    if depth >= AI_AUDIT_MAX_DEPTH:
        return AI_AUDIT_TRUNCATED_VALUE
    if isinstance(value, dict):
        safe: dict[str, Any] = {}
        for key, item in value.items():
            key_text = str(key)
            if key_text.lower() in AI_AUDIT_SENSITIVE_KEYS:
                safe[key_text] = AI_AUDIT_REDACTED_VALUE
            else:
                safe[key_text] = _audit_safe_payload(item, depth + 1)
        return safe
    if isinstance(value, (list, tuple)):
        items = list(value[:AI_AUDIT_MAX_SEQUENCE_ITEMS])
        safe_items = [_audit_safe_payload(item, depth + 1) for item in items]
        if len(value) > AI_AUDIT_MAX_SEQUENCE_ITEMS:
            safe_items.append(AI_AUDIT_TRUNCATED_VALUE)
        return safe_items
    if isinstance(value, str) and len(value) > AI_AUDIT_MAX_TEXT_LENGTH:
        return value[:AI_AUDIT_MAX_TEXT_LENGTH] + AI_AUDIT_TRUNCATED_VALUE
    return value


def _memory_scope(context: dict[str, Any]) -> str:
    return str(
        context.get(AIContextKey.MEMORY_SCOPE)
        or context.get(AIMemoryPayloadKey.SCOPE)
        or AIMemoryScope.GLOBAL
    )


def _memory_subject(context: dict[str, Any]) -> str | None:
    value = context.get(AIContextKey.MEMORY_SUBJECT) or context.get(AIMemoryPayloadKey.SUBJECT)
    return str(value) if value else None


def _validate_external_context(context: dict[str, Any]) -> None:
    controlled_keys = {
        AIContextKey.OPENCLAW_TOOL,
        AIContextKey.OPENCLAW_ACTION,
        AIContextKey.OPENCLAW_ARGS,
        AIContextKey.OPENCLAW_SESSION_KEY,
        AIContextKey.AGENT_PIPELINE,
        AIContextKey.HERMES_MEMORY,
        AIContextKey.LOCAL_MEMORY,
        AIContextKey.OPENCLAW,
        AIContextKey.MODE,
        AIContextKey.USER_PROMPT,
        AIContextKey.REQUEST_CONTEXT,
        AIContextKey.ASSISTANT_ANSWER,
        AIContextKey.USER_RULES,
        AIContextKey.COMPANY_RULES,
        AIContextKey.PERSONAL_RULES,
        AIContextKey.PREFERENCES,
        AIContextKey.INTERESTS,
        AIContextKey.CONVERSATION_HISTORY,
        AIContextKey.PROVIDER_SESSION_ID,
        AIContextKey.ALLOW_PROVIDER_MEMORY,
        AIContextKey.EXECUTION_MODE,
        "system_constraints",
        "current_request",
        "personal_memory",
        "owner_id",
        "tenant_key",
        "open_id",
    }
    supplied = {str(key) for key in context}
    if supplied.intersection(str(key) for key in controlled_keys):
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Controlled AI context fields are not allowed",
        )


def _adapter_context(
    personalization: PersonalizationContext,
    *,
    execution_mode: AIExecutionMode,
    allow_provider_memory: bool,
    provider_session_id: str | None,
) -> dict[str, Any]:
    context: dict[str, Any] = {
        AIContextKey.COMPANY_RULES: personalization.company_rules,
        AIContextKey.PERSONAL_RULES: personalization.personal_rules,
        AIContextKey.PREFERENCES: personalization.preferences,
        AIContextKey.INTERESTS: personalization.interests,
        AIContextKey.LOCAL_MEMORY: personalization.personal_memory,
        AIContextKey.CONVERSATION_HISTORY: personalization.conversation_history,
    }
    if provider_session_id:
        context[AIContextKey.PROVIDER_SESSION_ID] = provider_session_id
    context[AIContextKey.ALLOW_PROVIDER_MEMORY] = allow_provider_memory
    context[AIContextKey.EXECUTION_MODE] = execution_mode
    return context


def _is_unavailable_adapter(adapter: Any) -> bool:
    return str(getattr(adapter, "provider_name", "")).strip().lower() == AIProviderName.NOOP


def _unavailable_response(provider_name: Any) -> dict[str, Any]:
    return {
        AIResponseKey.OK: False,
        AIResponseKey.PROVIDER: str(provider_name or AIProviderName.NOOP),
        AIResponseKey.ANSWER: AI_UNAVAILABLE_ANSWER,
        AIResponseKey.RAW: {"reason": "provider_unavailable"},
        AIResponseKey.ERROR: "AI provider unavailable",
    }


def _required_answer(result: dict[str, Any]) -> str:
    answer = str(result.get(AIResponseKey.ANSWER) or "").strip()
    if not answer:
        raise HTTPException(
            status_code=status.HTTP_502_BAD_GATEWAY,
            detail="AI provider returned an empty answer",
        )
    return answer


def _validate_owner_id(owner_id: int) -> None:
    if owner_id <= 0:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="Valid AI owner is required",
        )


def _validate_prompt(prompt: str) -> None:
    if not str(prompt).strip():
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="AI prompt is required",
        )