add authoriztion

Author: xixingde

internal/handler/semantic.go

@@ -1,12 +1,12 @@
 package handler
 
-
 import (
-	"github.com/zgsm-ai/codebase-indexer/internal/response"
-	"net/http"
 	"fmt"
+	"net/http"
+
 	"github.com/zeromicro/go-zero/rest/httpx"
 	"github.com/zgsm-ai/codebase-indexer/internal/logic"
+	"github.com/zgsm-ai/codebase-indexer/internal/response"
 	"github.com/zgsm-ai/codebase-indexer/internal/svc"
 	"github.com/zgsm-ai/codebase-indexer/internal/types"
 )
@@ -20,12 +20,21 @@ func semanticSearchHandler(svcCtx *svc.ServiceContext) http.HandlerFunc {
 			return
 		}
 
+		// 从请求头获取 Authorization
+		authorization := r.Header.Get("Authorization")
+
+		// 验证 Authorization 头是否存在
+		if authorization == "" {
+			response.Error(w, response.NewAuthError("missing Authorization header"))
+			return
+		}
+
 		// 打印请求参数用于调试
-		fmt.Printf("Semantic search request received - ClientId: %s, CodebasePath: %s, Query: %s, TopK: %d, ScoreThreshold: %f\n", 
-		req.ClientId, req.CodebasePath, req.Query, req.TopK, req.ScoreThreshold)
+		fmt.Printf("Semantic search request received - ClientId: %s, CodebasePath: %s, Query: %s, TopK: %d, ScoreThreshold: %f\n",
+			req.ClientId, req.CodebasePath, req.Query, req.TopK, req.ScoreThreshold)
 
 		l := logic.NewSemanticSearchLogic(r.Context(), svcCtx)
-		resp, err := l.SemanticSearch(&req)
+		resp, err := l.SemanticSearch(&req, authorization)
 		if err != nil {
 			response.Error(w, err)
 		} else {

internal/logic/semantic.go

@@ -37,7 +37,7 @@ func NewSemanticSearchLogic(ctx context.Context, svcCtx *svc.ServiceContext) *Se
 	}
 }
 
-func (l *SemanticLogic) SemanticSearch(req *types.SemanticSearchRequest) (resp *types.SemanticSearchResponseData, err error) {
+func (l *SemanticLogic) SemanticSearch(req *types.SemanticSearchRequest, authorization string) (resp *types.SemanticSearchResponseData, err error) {
 	topK := req.TopK
 	if topK < minPositive {
 		topK = defaultTopK
@@ -61,8 +61,13 @@ func (l *SemanticLogic) SemanticSearch(req *types.SemanticSearchRequest) (resp *
 	ctx := context.WithValue(l.ctx, tracer.Key, tracer.RequestTraceId(int(codebase.ID)))
 
 	documents, err := l.svcCtx.VectorStore.Query(ctx, req.Query, topK,
-		vector.Options{CodebaseId: codebase.ID, ClientId: clientId,
-			CodebasePath: codebase.Path, CodebaseName: codebase.Name})
+		vector.Options{
+			CodebaseId:    codebase.ID,
+			ClientId:      clientId,
+			CodebasePath:  codebase.Path,
+			CodebaseName:  codebase.Name,
+			Authorization: authorization,
+		})
 	if err != nil {
 		return nil, err
 	}

internal/store/vector/vector_store.go

@@ -25,13 +25,14 @@ type Store interface {
 const vectorWeaviate = "weaviate"
 
 type Options struct {
-	CodebaseId   int32
-	SyncId       int32
-	RequestId    string
-	CodebasePath string
-	CodebaseName string
-	TotalFiles   int
-	ClientId     string
+	CodebaseId    int32
+	SyncId        int32
+	RequestId     string
+	CodebasePath  string
+	CodebaseName  string
+	TotalFiles    int
+	ClientId      string
+	Authorization string
 }
 
 func NewVectorStore(cfg config.VectorStoreConf, embedder Embedder, reranker Reranker) (Store, error) {

internal/store/vector/weaviate_wrapper.go

@@ -248,15 +248,15 @@ func (r *weaviateWrapper) SimilaritySearch(ctx context.Context, query string, nu
 		return nil, fmt.Errorf("query weaviate failed: %w", err)
 	}
 
-	items, err := r.unmarshalSimilarSearchResponse(res, options.CodebasePath, options.ClientId)
+	items, err := r.unmarshalSimilarSearchResponse(res, options.CodebasePath, options.ClientId, options.Authorization)
 	if err != nil {
 		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
 	}
 
 	return items, nil
 }
 
-func (r *weaviateWrapper) unmarshalSimilarSearchResponse(res *models.GraphQLResponse, codebasePath, clientId string) ([]*types.SemanticFileItem, error) {
+func (r *weaviateWrapper) unmarshalSimilarSearchResponse(res *models.GraphQLResponse, codebasePath, clientId string, authorization string) ([]*types.SemanticFileItem, error) {
 	// Get the data for our class
 	data, ok := res.Data["Get"].(map[string]interface{})
 	if !ok {
@@ -299,7 +299,7 @@ func (r *weaviateWrapper) unmarshalSimilarSearchResponse(res *models.GraphQLResp
 
 			// 通过fetchCodeContent接口获取代码片段
 			if codebasePath != "" {
-				fetchedContent, err := fetchCodeContent(context.Background(), r.cfg, clientId, codebasePath, filePath, startLine, endLine)
+				fetchedContent, err := fetchCodeContent(context.Background(), r.cfg, clientId, codebasePath, filePath, startLine, endLine, authorization)
 				if err == nil && fetchedContent != "" {
 					content = fetchedContent
 				}
@@ -694,7 +694,7 @@ func (r *weaviateWrapper) Query(ctx context.Context, query string, topK int, opt
 }
 
 // fetchCodeContent 通过API获取代码片段的Content
-func fetchCodeContent(ctx context.Context, cfg config.VectorStoreConf, clientId, codebasePath, filePath string, startLine, endLine int) (string, error) {
+func fetchCodeContent(ctx context.Context, cfg config.VectorStoreConf, clientId, codebasePath, filePath string, startLine, endLine int, authorization string) (string, error) {
 	// 构建API请求URL
 	baseURL := cfg.BaseURL
 
@@ -730,8 +730,19 @@ func fetchCodeContent(ctx context.Context, cfg config.VectorStoreConf, clientId,
 
 	tracer.WithTrace(ctx).Infof("fetchCodeContent %s: ", requestURL)
 
+	// 创建HTTP请求
+	req, err := http.NewRequestWithContext(ctx, "GET", requestURL, nil)
+	if err != nil {
+		return "", fmt.Errorf("failed to create request: %w", err)
+	}
+
+	// 添加Authorization头
+	if authorization != "" {
+		req.Header.Set("Authorization", authorization)
+	}
+
 	// 发送HTTP GET请求
-	resp, err := http.Get(requestURL)
+	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
 		return "", fmt.Errorf("failed to fetch code content: %w", err)
 	}