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

from typing import Any

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 (
    AIDefault,
    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,
    AIProviderName,
    AIRequestKey,
    AIResponseKey,
)
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


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]:
        adapter = get_adapter()
        result = adapter.ask(prompt, context or {})
        response = {
            AIResponseKey.PROVIDER: adapter.provider_name,
            AIResponseKey.ANSWER: result[AIResponseKey.ANSWER],
            AIResponseKey.RAW: result.get(AIResponseKey.RAW, {}),
        }
        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({
                    AIRequestKey.PROMPT: prompt,
                    AIRequestKey.CONTEXT: context or {},
                }),
                response_payload=_audit_safe_payload(response),
            )
        )
        return 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

    def invoke_openclaw_tool(
        self,
        tool: str,
        action: str = AIDefault.ACTION_JSON,
        args: dict[str, Any] | None = None,
        session_key: str = AIDefault.SESSION_KEY_MAIN,
        actor: str = ActorValue.API,
    ) -> dict[str, Any]:
        result = OpenClawAdapter(get_settings()).invoke_tool(tool, action, args or {}, session_key)
        response = {AIResponseKey.PROVIDER: AIProviderName.OPENCLAW, AIResponseKey.RESULT: result}
        self.audit.log(
            AuditLogCreate(
                actor=actor,
                source=AuditSource.OPENCLAW,
                action=AuditAction.OPENCLAW_TOOLS_INVOKE,
                target_type=AuditTargetType.OPENCLAW_TOOL,
                target_id=tool,
                risk_level=AuditRiskLevel.HIGH,
                request_payload=_audit_safe_payload({
                    "tool": tool,
                    "action": action,
                    "args": args or {},
                    "session_key": session_key,
                }),
                response_payload=_audit_safe_payload(result),
            )
        )
        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