83 lines
1.8 KiB
Go
83 lines
1.8 KiB
Go
package llm
|
|
|
|
import (
|
|
"ai-code-assistant/pkg/database"
|
|
"context"
|
|
"github.com/tmc/langchaingo/callbacks"
|
|
"github.com/tmc/langchaingo/schema"
|
|
"github.com/tmc/langchaingo/vectorstores"
|
|
"strconv"
|
|
)
|
|
|
|
type RelevantDocs struct {
|
|
CallbacksHandler callbacks.Handler
|
|
db *database.Database
|
|
llm *LLM
|
|
repoID string
|
|
size int
|
|
}
|
|
|
|
type FileChunkID struct {
|
|
Name string
|
|
ChunkID int
|
|
Start uint64
|
|
End uint64
|
|
Score float32
|
|
Doc *schema.Document
|
|
}
|
|
|
|
func NewGetRelevantDocs(db *database.Database, llm *LLM, repoID string, size int) *RelevantDocs {
|
|
return &RelevantDocs{
|
|
db: db,
|
|
llm: llm,
|
|
repoID: repoID,
|
|
size: size,
|
|
}
|
|
}
|
|
|
|
func (rd *RelevantDocs) GetRelevantFileChunks(ctx context.Context, query string) ([]*FileChunkID, error) {
|
|
vectorStore, closeFunc, err := rd.db.GetVectorStore(ctx, rd.llm.Embedder())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer closeFunc()
|
|
|
|
retr := vectorstores.ToRetriever(vectorStore, rd.size, vectorstores.WithFilters(map[string]any{"type": "file_chunk", "repo_id": rd.repoID}))
|
|
retr.CallbacksHandler = rd.CallbacksHandler
|
|
|
|
docs, err := retr.GetRelevantDocuments(ctx, query)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var chunks []*FileChunkID
|
|
|
|
for _, doc := range docs {
|
|
chunk := &FileChunkID{
|
|
Score: doc.Score,
|
|
}
|
|
|
|
if filePath, ok := doc.Metadata["file_path"].(string); ok {
|
|
chunk.Name = filePath
|
|
}
|
|
|
|
if chunkID, ok := doc.Metadata["chunk_id"].(string); ok {
|
|
id, _ := strconv.ParseInt(chunkID, 10, 64)
|
|
chunk.ChunkID = int(id)
|
|
}
|
|
|
|
if start, ok := doc.Metadata["start"].(string); ok {
|
|
chunk.Start, _ = strconv.ParseUint(start, 10, 64)
|
|
}
|
|
|
|
if end, ok := doc.Metadata["end"].(string); ok {
|
|
chunk.End, _ = strconv.ParseUint(end, 10, 64)
|
|
}
|
|
|
|
chunk.Doc = &doc
|
|
chunks = append(chunks, chunk)
|
|
}
|
|
|
|
return chunks, nil
|
|
}
|