54 lines
1.9 KiB
Python
54 lines
1.9 KiB
Python
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 |