erp-ass/backend/app/core/ai_engine.py

"""AI Engine for ERP AI Assistant.

This module provides the ClaudeEngine class that wraps Claude API calls
and provides JSON parsing utilities.
"""

import json
import re
from typing import Any

import anthropic
from loguru import logger

from app.config import get_settings


class ClaudeEngine:
    """Engine for interacting with Claude API.

    This class wraps the Anthropic Claude API client and provides
    utilities for parsing JSON responses from Claude.
    """

    def __init__(self) -> None:
        """Initialize Claude engine with settings."""
        settings = get_settings()

        # Initialize Anthropic client with optional custom base_url
        client_kwargs = {"api_key": settings.ANTHROPIC_API_KEY}
        if settings.ANTHROPIC_BASE_URL:
            client_kwargs["base_url"] = settings.ANTHROPIC_BASE_URL
            logger.info(f"Using custom Anthropic base URL: {settings.ANTHROPIC_BASE_URL}")

        self.client = anthropic.AsyncAnthropic(**client_kwargs)
        self.model = settings.CLAUDE_MODEL
        self.max_tokens = settings.CLAUDE_MAX_TOKENS
        self.temperature = settings.CLAUDE_TEMPERATURE

    def parse_json_response(self, content: str) -> dict[str, Any]:
        """Parse JSON from Claude responses.

        Attempts multiple parsing strategies:
        1. Direct JSON parse
        2. Extract from markdown code blocks
        3. Extract any {...} block

        Args:
            content: The response content from Claude

        Returns:
            Parsed JSON as a dictionary

        Raises:
            ValueError: If JSON cannot be parsed using any strategy
        """
        if not content or not content.strip():
            raise ValueError("Empty content provided")

        # Strategy 1: Try direct JSON parse
        try:
            return json.loads(content)
        except json.JSONDecodeError:
            pass

        # Strategy 2: Try extracting from markdown code blocks
        json_code_block_pattern = r'```json\s*(\{.*?\})\s*```'
        json_match = re.search(json_code_block_pattern, content, re.DOTALL)
        if json_match:
            try:
                return json.loads(json_match.group(1))
            except json.JSONDecodeError:
                pass

        # Also try any code block (not just json tagged)
        code_block_pattern = r'```\s*(\{.*?\})\s*```'
        code_block_match = re.search(code_block_pattern, content, re.DOTALL)
        if code_block_match:
            try:
                return json.loads(code_block_match.group(1))
            except json.JSONDecodeError:
                pass

        # Strategy 3: Try extracting any {...} block
        # Find balanced braces
        brace_pattern = r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}'
        json_blocks = re.findall(brace_pattern, content, re.DOTALL)
        for json_block in json_blocks:
            try:
                return json.loads(json_block)
            except json.JSONDecodeError:
                continue

        # All strategies failed
        logger.error(f"无法解析 Claude 返回的 JSON: {content[:200]}")
        raise ValueError("无法解析 Claude 返回的 JSON，请检查响应格式")

    async def call_claude(
        self,
        messages: list[dict[str, str]],
        temperature: float | None = None
    ) -> str:
        """Call Claude API.

        Args:
            messages: List of message dictionaries with 'role' and 'content'
            temperature: Optional temperature override (0-2)

        Returns:
            The text content from Claude's response

        Raises:
            Exception: If the API call fails
        """
        response = await self.client.messages.create(
            model=self.model,
            max_tokens=self.max_tokens,
            temperature=temperature if temperature is not None else self.temperature,
            messages=messages
        )
        return response.content[0].text