38 lines
1.3 KiB
Python
38 lines
1.3 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()
|
|||
|
|
elif '```' in response:
|
|||
|
|
start = response.find('```') + 3
|
|||
|
|
end = response.find('```', start)
|
|||
|
|
if end != -1:
|
|||
|
|
return response[start:end].strip()
|
|||
|
|
|
|||
|
|
return None
|