Skip to content

Commit 9379276

Browse files
xixingdexixingde
committed
get batch content
1 parent 095a42e commit 9379276

File tree

5 files changed

+206
-25
lines changed

5 files changed

+206
-25
lines changed

‎etc/conf.yaml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ Cleaner:
3838

3939
Redis:
4040
Addr: 0.0.0.0:6379
41+
DefaultExpiration: 1h
4142

4243
MessageQueue:
4344
Type: redis

‎internal/config/db.go‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,13 @@ type Database struct {
1919

2020
// RedisConfig Redis配置
2121
type RedisConfig struct {
22-
Addr string
23-
Password string `json:",optional"`
24-
DB int `json:",default=0"`
25-
PoolSize int `json:",default=10"`
26-
MinIdleConn int `json:",default=10"`
27-
ConnectTimeout time.Duration `json:",default=10s"`
28-
ReadTimeout time.Duration `json:",default=10s"`
29-
WriteTimeout time.Duration `json:",default=10s"`
22+
Addr string
23+
Password string `json:",optional"`
24+
DB int `json:",default=0"`
25+
PoolSize int `json:",default=10"`
26+
MinIdleConn int `json:",default=10"`
27+
ConnectTimeout time.Duration `json:",default=10s"`
28+
ReadTimeout time.Duration `json:",default=10s"`
29+
WriteTimeout time.Duration `json:",default=10s"`
30+
DefaultExpiration time.Duration `json:",default=24h"` // 默认过期时间
3031
}

‎internal/store/redis/status_manager.go‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,19 +18,27 @@ const (
1818
fileStatusPrefix = "file:status:"
1919
// 请求ID键前缀
2020
requestIdPrefix = "request:id:"
21-
// 默认过期时间 - 24小时
22-
defaultExpiration = 24 * time.Hour
2321
)
2422

2523
// StatusManager 文件状态管理器
2624
type StatusManager struct {
27-
client *redis.Client
25+
client *redis.Client
26+
defaultExpiration time.Duration
2827
}
2928

3029
// NewStatusManager 创建新的状态管理器
3130
func NewStatusManager(client *redis.Client) *StatusManager {
3231
return &StatusManager{
33-
client: client,
32+
client: client,
33+
defaultExpiration: 24 * time.Hour, // 默认24小时，保持向后兼容
34+
}
35+
}
36+
37+
// NewStatusManagerWithExpiration 创建带有自定义过期时间的状态管理器
38+
func NewStatusManagerWithExpiration(client *redis.Client, expiration time.Duration) *StatusManager {
39+
return &StatusManager{
40+
client: client,
41+
defaultExpiration: expiration,
3442
}
3543
}
3644

@@ -44,7 +52,7 @@ func (sm *StatusManager) SetFileStatusByRequestId(ctx context.Context, requestId
4452
}
4553

4654
// 设置到Redis，带过期时间
47-
err = sm.client.Set(ctx, key, data, defaultExpiration).Err()
55+
err = sm.client.Set(ctx, key, data, sm.defaultExpiration).Err()
4856
if err != nil {
4957
return fmt.Errorf("failed to set status in redis: %w", err)
5058
}
@@ -268,7 +276,7 @@ func (sm *StatusManager) ResetPendingAndProcessingTasksToFailed(ctx context.Cont
268276
}
269277

270278
// 更新Redis中的状态
271-
if err := sm.client.Set(ctx, key, updatedData, defaultExpiration).Err(); err != nil {
279+
if err := sm.client.Set(ctx, key, updatedData, sm.defaultExpiration).Err(); err != nil {
272280
logx.Errorf("failed to update status to failed for key %s: %v", key, err)
273281
continue
274282
}

‎internal/store/vector/weaviate_wrapper.go‎

Lines changed: 180 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
package vector
22

33
import (
4+
"bytes"
45
"context"
56
"crypto/md5"
67
"encoding/hex"
8+
"encoding/json"
79
"fmt"
810
"io"
911
"math"
@@ -269,6 +271,72 @@ func (r *weaviateWrapper) unmarshalSimilarSearchResponse(res *models.GraphQLResp
269271
}
270272

271273
items := make([]*types.SemanticFileItem, 0, len(results))
274+
275+
// 如果开启获取源码，则收集所有需要获取的代码片段
276+
var snippets []CodeSnippetRequest
277+
var snippetInfoList []struct {
278+
index int
279+
filePath string
280+
startLine int
281+
endLine int
282+
}
283+
284+
// 第一遍遍历：收集所有需要获取源码的片段信息
285+
for i, result := range results {
286+
obj, ok := result.(map[string]interface{})
287+
if !ok {
288+
continue
289+
}
290+
291+
content := getStringValue(obj, Content)
292+
filePath := getStringValue(obj, MetadataFilePath)
293+
294+
// 如果开启获取源码，则从MetadataRange中提取行号信息
295+
if r.cfg.FetchSourceCode && content != "" && filePath != "" && codebasePath != "" {
296+
// 从MetadataRange中提取startLine和endLine
297+
var startLine, endLine int
298+
if rangeValue, ok := obj[MetadataRange].([]interface{}); ok && len(rangeValue) >= 2 {
299+
if first, ok := rangeValue[0].(float64); ok {
300+
startLine = int(first)
301+
}
302+
if second, ok := rangeValue[2].(float64); ok {
303+
endLine = int(second)
304+
}
305+
}
306+
307+
// 添加到批量获取列表
308+
snippets = append(snippets, CodeSnippetRequest{
309+
FilePath: filePath,
310+
StartLine: startLine,
311+
EndLine: endLine,
312+
})
313+
314+
snippetInfoList = append(snippetInfoList, struct {
315+
index int
316+
filePath string
317+
startLine int
318+
endLine int
319+
}{
320+
index: i,
321+
filePath: filePath,
322+
startLine: startLine,
323+
endLine: endLine,
324+
})
325+
}
326+
}
327+
328+
// 批量获取代码片段内容
329+
var contentMap map[string]string
330+
if len(snippets) > 0 && codebasePath != "" {
331+
var err error
332+
contentMap, err = fetchCodeContentsBatch(context.Background(), r.cfg, clientId, codebasePath, snippets, authorization)
333+
if err != nil {
334+
fmt.Printf("[DEBUG] 批量获取代码片段失败: %v\n", err)
335+
contentMap = make(map[string]string)
336+
}
337+
}
338+
339+
// 第二遍遍历：构建最终的SemanticFileItem列表
272340
for _, result := range results {
273341
obj, ok := result.(map[string]interface{})
274342
if !ok {
@@ -284,9 +352,9 @@ func (r *weaviateWrapper) unmarshalSimilarSearchResponse(res *models.GraphQLResp
284352
content := getStringValue(obj, Content)
285353
filePath := getStringValue(obj, MetadataFilePath)
286354

287-
// 如果开启获取源码，则从MetadataRange中提取行号信息
288-
if r.cfg.FetchSourceCode && content != "" && filePath != "" {
289-
// 从MetadataRange中提取startLine和endLine
355+
// 如果开启获取源码且有批量获取的内容，则使用获取到的内容
356+
if r.cfg.FetchSourceCode && content != "" && filePath != "" && codebasePath != "" {
357+
// 从MetadataRange中提取startLine和endLine（用于构建映射键）
290358
var startLine, endLine int
291359
if rangeValue, ok := obj[MetadataRange].([]interface{}); ok && len(rangeValue) >= 2 {
292360
if first, ok := rangeValue[0].(float64); ok {
@@ -297,12 +365,10 @@ func (r *weaviateWrapper) unmarshalSimilarSearchResponse(res *models.GraphQLResp
297365
}
298366
}
299367

300-
// 通过fetchCodeContent接口获取代码片段
301-
if codebasePath != "" {
302-
fetchedContent, err := fetchCodeContent(context.Background(), r.cfg, clientId, codebasePath, filePath, startLine, endLine, authorization)
303-
if err == nil && fetchedContent != "" {
304-
content = fetchedContent
305-
}
368+
// 构建映射键并查找批量获取的内容
369+
key := fmt.Sprintf("%s:%d-%d", filePath, startLine, endLine)
370+
if fetchedContent, exists := contentMap[key]; exists && fetchedContent != "" {
371+
content = fetchedContent
306372
}
307373
}
308374

@@ -693,6 +759,111 @@ func (r *weaviateWrapper) Query(ctx context.Context, query string, topK int, opt
693759
return rerankedDocs, nil
694760
}
695761

762+
// CodeSnippetRequest 代码片段请求结构
763+
type CodeSnippetRequest struct {
764+
FilePath string `json:"filePath"`
765+
StartLine int `json:"startLine"`
766+
EndLine int `json:"endLine"`
767+
}
768+
769+
// CodeSnippetsBatchRequest 批量代码片段请求结构
770+
type CodeSnippetsBatchRequest struct {
771+
ClientId string `json:"clientId"`
772+
WorkspacePath string `json:"workspacePath"`
773+
CodeSnippets []CodeSnippetRequest `json:"codeSnippets"`
774+
}
775+
776+
// CodeSnippetResponse 代码片段响应结构
777+
type CodeSnippetResponse struct {
778+
FilePath string `json:"filePath"`
779+
StartLine int `json:"startLine"`
780+
EndLine int `json:"endLine"`
781+
Content string `json:"content"`
782+
}
783+
784+
// CodeSnippetsBatchResponse 批量代码片段响应结构
785+
type CodeSnippetsBatchResponse struct {
786+
Code string `json:"code"`
787+
Message string `json:"message"`
788+
Success bool `json:"success"`
789+
Data struct {
790+
List []CodeSnippetResponse `json:"list"`
791+
} `json:"data"`
792+
}
793+
794+
// fetchCodeContentsBatch 批量获取代码片段Content
795+
func fetchCodeContentsBatch(ctx context.Context, cfg config.VectorStoreConf, clientId, codebasePath string, snippets []CodeSnippetRequest, authorization string) (map[string]string, error) {
796+
if len(snippets) == 0 {
797+
return nil, nil
798+
}
799+
800+
// 构建请求体
801+
request := CodeSnippetsBatchRequest{
802+
ClientId: clientId,
803+
WorkspacePath: codebasePath,
804+
CodeSnippets: snippets,
805+
}
806+
807+
jsonData, err := json.Marshal(request)
808+
if err != nil {
809+
return nil, fmt.Errorf("failed to marshal request: %w", err)
810+
}
811+
812+
// 构建API请求URL
813+
apiURL := "http://localhost:11380/codebase-indexer/api/v1/snippets/read"
814+
815+
tracer.WithTrace(ctx).Infof("fetchCodeContentsBatch: %s", apiURL)
816+
817+
// 创建HTTP请求
818+
req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewBuffer(jsonData))
819+
if err != nil {
820+
return nil, fmt.Errorf("failed to create request: %w", err)
821+
}
822+
823+
// 设置请求头
824+
req.Header.Set("Content-Type", "application/json")
825+
if authorization != "" {
826+
req.Header.Set("Authorization", authorization)
827+
}
828+
829+
// 发送HTTP POST请求
830+
resp, err := http.DefaultClient.Do(req)
831+
if err != nil {
832+
return nil, fmt.Errorf("failed to fetch code contents batch: %w", err)
833+
}
834+
defer resp.Body.Close()
835+
836+
// 检查响应状态码
837+
if resp.StatusCode != http.StatusOK {
838+
return nil, fmt.Errorf("unexpected status code: %d", resp.StatusCode)
839+
}
840+
841+
// 读取响应体
842+
body, err := io.ReadAll(resp.Body)
843+
if err != nil {
844+
return nil, fmt.Errorf("failed to read response body: %w", err)
845+
}
846+
847+
// 解析响应
848+
var batchResponse CodeSnippetsBatchResponse
849+
if err := json.Unmarshal(body, &batchResponse); err != nil {
850+
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
851+
}
852+
853+
if !batchResponse.Success {
854+
return nil, fmt.Errorf("API request failed: %s", batchResponse.Message)
855+
}
856+
857+
// 构建filePath到content的映射
858+
contentMap := make(map[string]string)
859+
for _, snippet := range batchResponse.Data.List {
860+
key := fmt.Sprintf("%s:%d-%d", snippet.FilePath, snippet.StartLine, snippet.EndLine)
861+
contentMap[key] = snippet.Content
862+
}
863+
864+
return contentMap, nil
865+
}
866+
696867
// fetchCodeContent 通过API获取代码片段的Content
697868
func fetchCodeContent(ctx context.Context, cfg config.VectorStoreConf, clientId, codebasePath, filePath string, startLine, endLine int, authorization string) (string, error) {
698869
// 构建API请求URL

‎internal/svc/service_context.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -109,8 +109,8 @@ func NewServiceContext(ctx context.Context, c config.Config) (*ServiceContext, e
109109
svcCtx.Embedder = embedder
110110
svcCtx.CodeSplitter = splitter
111111
svcCtx.DistLock = lock
112-
// 状态管理器
113-
svcCtx.StatusManager = redisstore.NewStatusManager(client)
112+
// 状态管理器 - 使用配置中的默认过期时间
113+
svcCtx.StatusManager = redisstore.NewStatusManagerWithExpiration(client, c.Redis.DefaultExpiration)
114114

115115
// 向量知识库
116116
vectorStore, err := vector.NewVectorStoreWithStatusManager(c.VectorStore, embedder, reranker, svcCtx.StatusManager, "")

0 commit comments

Comments
 (0)