iov_data_analysis_agent/utils/extract_code.py

from typing import Optional
import yaml


def extract_code_from_response(response: str) -> Optional[str]:
        """从LLM响应中提取代码"""
        try:
            # 尝试解析YAML
            if '```yaml' in response:
                start = response.find('```yaml') + 7
                end = response.find('```', start)
                yaml_content = response[start:end].strip()
            elif '```' in response:
                start = response.find('```') + 3
                end = response.find('```', start)
                yaml_content = response[start:end].strip()
            else:
                yaml_content = response.strip()
            
            yaml_data = yaml.safe_load(yaml_content)
            if 'code' in yaml_data:
                return yaml_data['code']
        except:
            pass
        
        # 如果YAML解析失败，尝试提取```python代码块
        if '```python' in response:
            start = response.find('```python') + 9
            end = response.find('```', start)
            if end != -1:
                return response[start:end].strip()
        
        # 尝试提取 code: | 形式的代码块（针对YAML格式错误但结构清晰的情况）
        import re
        # 匹配 code: | 后面的内容，直到遇到下一个键（next_key:）或结尾
        # 假设代码块至少缩进2个空格
        pattern = r'code:\s*\|\s*\n((?: {2,}.*\n?)+)'
        match = re.search(pattern, response)
        if match:
            code_block = match.group(1)
            # 尝试去除公共缩进
            try:
                import textwrap
                return textwrap.dedent(code_block).strip()
            except:
                return code_block.strip()

        elif '```' in response:
            start = response.find('```') + 3
            end = response.find('```', start)
            if end != -1:
                return response[start:end].strip()
        
        return None