152 lines
3.9 KiB
Go
152 lines
3.9 KiB
Go
package extractor
|
||
|
||
import (
|
||
"fmt"
|
||
"os"
|
||
"path/filepath"
|
||
"strings"
|
||
)
|
||
|
||
type File struct {
|
||
Path string
|
||
Content string
|
||
}
|
||
|
||
// Extract 从模型响应中提取代码文件
|
||
// 使用纯字符串解析，不依赖 goldmark AST
|
||
func Extract(response string) ([]File, error) {
|
||
var files []File
|
||
seen := make(map[string]bool)
|
||
|
||
// 按 // File: 分割
|
||
parts := strings.Split(response, "// File:")
|
||
|
||
for _, part := range parts {
|
||
part = strings.TrimSpace(part)
|
||
if part == "" {
|
||
continue
|
||
}
|
||
|
||
// 第一行是文件名
|
||
lines := strings.SplitN(part, "\n", 2)
|
||
if len(lines) < 2 {
|
||
continue
|
||
}
|
||
|
||
fileName := strings.TrimSpace(lines[0])
|
||
code := strings.TrimSpace(lines[1])
|
||
|
||
if fileName == "" || code == "" {
|
||
continue
|
||
}
|
||
|
||
// 清理代码块标记（移除 ```go 或 ```）
|
||
code = strings.TrimPrefix(code, "```go")
|
||
code = strings.TrimPrefix(code, "```")
|
||
code = strings.TrimSuffix(code, "```")
|
||
code = strings.TrimSpace(code)
|
||
|
||
// 智能识别文件类型
|
||
fileName = smartFileName(fileName, code)
|
||
|
||
// 规范化路径
|
||
fullPath := normalizePath(fileName)
|
||
|
||
// 去重
|
||
if seen[fullPath] {
|
||
continue
|
||
}
|
||
seen[fullPath] = true
|
||
|
||
files = append(files, File{
|
||
Path: fullPath,
|
||
Content: code,
|
||
})
|
||
}
|
||
|
||
if len(files) == 0 {
|
||
return nil, fmt.Errorf("未找到任何代码块，请确保响应格式为：// File: 文件名\n```go\n代码\n```")
|
||
}
|
||
|
||
return files, nil
|
||
}
|
||
|
||
// smartFileName 智能识别文件名
|
||
func smartFileName(fileName, code string) string {
|
||
fileName = strings.TrimSpace(fileName)
|
||
fileName = strings.TrimLeft(fileName, "0123456789.- ")
|
||
|
||
// 如果文件名是 code.go 或类似通用名，根据内容推断
|
||
base := filepath.Base(fileName)
|
||
if base == "code.go" || base == "code" || base == "main.go" {
|
||
// 根据内容推断
|
||
if strings.Contains(code, "package sdk_test") {
|
||
return "example_test.go"
|
||
}
|
||
if strings.Contains(code, "module ") {
|
||
return "go.mod"
|
||
}
|
||
if strings.Contains(code, "package sdk") && strings.Contains(code, "func (c *Client)") {
|
||
// 检查是哪个 API 文件
|
||
if strings.Contains(code, "CreateKeyOrder") || strings.Contains(code, "OrderKey") {
|
||
return "api_key.go"
|
||
}
|
||
if strings.Contains(code, "BatchOrder") || strings.Contains(code, "BatchQuery") {
|
||
return "api_batch.go"
|
||
}
|
||
if strings.Contains(code, "QueryKey") || strings.Contains(code, "DiscardKey") {
|
||
return "api_key.go"
|
||
}
|
||
}
|
||
if strings.Contains(code, "package sdk") && strings.Contains(code, "type Client struct") {
|
||
return "client.go"
|
||
}
|
||
if strings.Contains(code, "package sdk") && (strings.Contains(code, "type APIError") || strings.Contains(code, "CodeSuccess")) {
|
||
return "errors.go"
|
||
}
|
||
if strings.Contains(code, "package sdk") && (strings.Contains(code, "ParseRSAPrivateKey") || strings.Contains(code, "AESEncryptECB")) {
|
||
return "crypto.go"
|
||
}
|
||
if strings.Contains(code, "package sdk") && (strings.Contains(code, "type RequestEnvelope") || strings.Contains(code, "type KeyOrderInfo")) {
|
||
return "types.go"
|
||
}
|
||
}
|
||
|
||
return fileName
|
||
}
|
||
|
||
func normalizePath(path string) string {
|
||
path = strings.TrimSpace(path)
|
||
path = strings.TrimLeft(path, "0123456789.- ")
|
||
|
||
// 确保文件扩展名正确
|
||
if !strings.HasSuffix(path, ".go") && !strings.HasSuffix(path, ".mod") {
|
||
// 如果是 go.mod 或 go.work，保持原样
|
||
if path != "go.mod" && path != "go.work" {
|
||
path = path + ".go"
|
||
}
|
||
}
|
||
|
||
// 确保路径以 sdk/ 开头
|
||
//if !strings.HasPrefix(path, "sdk/") && !strings.HasPrefix(path, "output/") {
|
||
// path = filepath.Join("sdk", path)
|
||
//}
|
||
return path
|
||
}
|
||
|
||
// WriteFiles 将文件写入本地
|
||
func WriteFiles(files []File, outputDir string) error {
|
||
for _, f := range files {
|
||
fullPath := filepath.Join(outputDir, f.Path)
|
||
|
||
if err := os.MkdirAll(filepath.Dir(fullPath), 0755); err != nil {
|
||
return fmt.Errorf("创建目录失败: %w", err)
|
||
}
|
||
|
||
if err := os.WriteFile(fullPath, []byte(f.Content), 0644); err != nil {
|
||
return fmt.Errorf("写入文件失败: %w", fullPath, err)
|
||
}
|
||
}
|
||
return nil
|
||
}
|