enrich.go raw
1 package cartography
2
3 import (
4 "fmt"
5 "os"
6 "path/filepath"
7 "sort"
8 "strings"
9 )
10
11 // composeEnrichQuery builds a prompt describing a declaration.
12 func composeEnrichQuery(entry *Entry, source string, related []string) string {
13 var b strings.Builder
14
15 b.WriteString("Describe this Go declaration. Give a concise answer with these sections:\n\n")
16 b.WriteString("DESCRIPTION: What it does and why (one paragraph)\n")
17 b.WriteString("CONTRACT: Inputs, outputs, error conditions, guarantees (one paragraph)\n")
18 b.WriteString("EDGE CASES: What happens with nil, empty, zero, or invalid inputs (one line)\n")
19 b.WriteString("VALIDATION: How to verify it works correctly (one line)\n\n")
20
21 fmt.Fprintf(&b, "Package: %s\n", entry.Package)
22 if entry.Signature != "" {
23 fmt.Fprintf(&b, "Signature: %s\n", entry.Signature)
24 }
25 if entry.DocComment != "" {
26 fmt.Fprintf(&b, "Doc: %s\n", entry.DocComment)
27 }
28 if len(related) > 0 {
29 fmt.Fprintf(&b, "Related declarations: %s\n", strings.Join(related, ", "))
30 }
31
32 fmt.Fprintf(&b, "\nSource code:\n```go\n%s\n```\n", source)
33
34 return b.String()
35 }
36
37 // parseEnrichResponse extracts structured fields from the oracle's response.
38 func parseEnrichResponse(entry *Entry, answer string) {
39 lines := strings.Split(answer, "\n")
40
41 var section string
42 var desc, contract, edgeCases, validation strings.Builder
43
44 for _, line := range lines {
45 trimmed := strings.TrimSpace(line)
46 upper := strings.ToUpper(trimmed)
47
48 switch {
49 case strings.HasPrefix(upper, "DESCRIPTION:"):
50 section = "desc"
51 rest := strings.TrimPrefix(trimmed, trimmed[:12])
52 rest = strings.TrimSpace(rest)
53 if rest != "" {
54 desc.WriteString(rest)
55 }
56 case strings.HasPrefix(upper, "CONTRACT:"):
57 section = "contract"
58 rest := strings.TrimPrefix(trimmed, trimmed[:9])
59 rest = strings.TrimSpace(rest)
60 if rest != "" {
61 contract.WriteString(rest)
62 }
63 case strings.HasPrefix(upper, "EDGE CASES:") || strings.HasPrefix(upper, "EDGE_CASES:"):
64 section = "edge"
65 rest := trimmed[strings.IndexByte(trimmed, ':')+1:]
66 rest = strings.TrimSpace(rest)
67 if rest != "" {
68 edgeCases.WriteString(rest)
69 }
70 case strings.HasPrefix(upper, "VALIDATION:"):
71 section = "validation"
72 rest := strings.TrimPrefix(trimmed, trimmed[:11])
73 rest = strings.TrimSpace(rest)
74 if rest != "" {
75 validation.WriteString(rest)
76 }
77 default:
78 if trimmed == "" {
79 continue
80 }
81 switch section {
82 case "desc":
83 desc.WriteString(" ")
84 desc.WriteString(trimmed)
85 case "contract":
86 contract.WriteString(" ")
87 contract.WriteString(trimmed)
88 case "edge":
89 edgeCases.WriteString(" ")
90 edgeCases.WriteString(trimmed)
91 case "validation":
92 validation.WriteString(" ")
93 validation.WriteString(trimmed)
94 }
95 }
96 }
97
98 if d := strings.TrimSpace(desc.String()); d != "" {
99 entry.Description = d
100 }
101 if c := strings.TrimSpace(contract.String()); c != "" {
102 entry.Contract = c
103 }
104 if e := strings.TrimSpace(edgeCases.String()); e != "" {
105 entry.EdgeCases = e
106 }
107 if v := strings.TrimSpace(validation.String()); v != "" {
108 entry.Validation = v
109 }
110 }
111
112 // readDeclSourceFrom reads declaration source from an explicit root directory.
113 // For external packages, sourceRoot is the package directory and entry.FilePath
114 // is importPath/filename.go — we use only the filename part.
115 func readDeclSourceFrom(entry *Entry, sourceRoot string) (string, error) {
116 // entry.FilePath is like "go/token/token.go"; we need just the filename.
117 filename := filepath.Base(entry.FilePath)
118 path := filepath.Join(sourceRoot, filename)
119 return readDeclSourcePath(path, entry.Line)
120 }
121
122 // readDeclSource reads the source text of a declaration from disk.
123 // Returns up to 150 lines starting from the declaration's line.
124 func readDeclSource(entry *Entry, projectRoot string) (string, error) {
125 path := filepath.Join(projectRoot, entry.FilePath)
126 return readDeclSourcePath(path, entry.Line)
127 }
128
129 // readDeclSourcePath reads a declaration from a file starting at the given line.
130 func readDeclSourcePath(path string, line int) (string, error) {
131 data, err := os.ReadFile(path)
132 if err != nil {
133 return "", err
134 }
135
136 lines := strings.Split(string(data), "\n")
137 start := line - 1 // 0-indexed
138 if start < 0 {
139 start = 0
140 }
141 end := start + 150
142 if end > len(lines) {
143 end = len(lines)
144 }
145
146 // Find the end of the declaration (matching braces).
147 depth := 0
148 for i := start; i < end; i++ {
149 for _, ch := range lines[i] {
150 if ch == '{' {
151 depth++
152 }
153 if ch == '}' {
154 depth--
155 if depth == 0 {
156 return strings.Join(lines[start:i+1], "\n"), nil
157 }
158 }
159 }
160 }
161
162 return strings.Join(lines[start:end], "\n"), nil
163 }
164
165 // findRelated returns IDs of related entries (same package, referenced in DependsOn).
166 func findRelated(entry *Entry, atlas *Atlas) []string {
167 ids := make([]string, 0, len(atlas.Entries))
168 for id := range atlas.Entries {
169 ids = append(ids, id)
170 }
171 sort.Strings(ids)
172
173 var related []string
174 for _, id := range ids {
175 e := atlas.Entries[id]
176 if e.Package == entry.Package && e.ID != entry.ID && e.Exported {
177 related = append(related, e.ID)
178 if len(related) >= 5 {
179 break
180 }
181 }
182 }
183
184 return related
185 }
186