You've already forked agentic-coding-workflow
bugfix
This commit is contained in:
@@ -0,0 +1,294 @@
|
||||
package codeindex
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgtype"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
"github.com/pgvector/pgvector-go"
|
||||
|
||||
codeindexsqlc "github.com/yan1h/agent-coding-workflow/internal/codeindex/sqlc"
|
||||
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
|
||||
)
|
||||
|
||||
// ChunkRow is a fully-populated chunk ready for bulk insert (embedding included).
|
||||
type ChunkRow struct {
|
||||
FilePath string
|
||||
Lang string
|
||||
StartLine int
|
||||
EndLine int
|
||||
Content string
|
||||
ContentSHA []byte
|
||||
TokenCount int
|
||||
Embedding []float32
|
||||
}
|
||||
|
||||
// Repository is the codeindex persistence contract.
|
||||
type Repository interface {
|
||||
InsertRun(ctx context.Context, r *Run) (*Run, error)
|
||||
GetRun(ctx context.Context, id uuid.UUID) (*Run, error)
|
||||
GetRunByCommit(ctx context.Context, wsID uuid.UUID, commitSHA, model string) (*Run, error)
|
||||
GetLatestCompletedRun(ctx context.Context, wsID uuid.UUID) (*Run, error)
|
||||
MarkRunning(ctx context.Context, id uuid.UUID) error
|
||||
MarkCompleted(ctx context.Context, id uuid.UUID, files, chunks int) error
|
||||
MarkFailed(ctx context.Context, id uuid.UUID, errMsg string) error
|
||||
MarkStale(ctx context.Context, wsID uuid.UUID) error
|
||||
MarkStuckStale(ctx context.Context, runningBefore time.Time) ([]uuid.UUID, error)
|
||||
// ReplaceChunks deletes any existing chunks for run then bulk-inserts rows.
|
||||
ReplaceChunks(ctx context.Context, runID, wsID uuid.UUID, commitSHA string, rows []ChunkRow) error
|
||||
// SearchKNN returns the top-k chunks of the latest completed run for wsID
|
||||
// ordered by cosine similarity to queryVec, with optional path/lang filters.
|
||||
SearchKNN(ctx context.Context, wsID uuid.UUID, queryVec []float32, topK int, pathPrefix, lang string) ([]Snippet, error)
|
||||
}
|
||||
|
||||
// PgRepository is the production Repository on a pgxpool.Pool.
|
||||
type PgRepository struct {
|
||||
pool *pgxpool.Pool
|
||||
q *codeindexsqlc.Queries
|
||||
}
|
||||
|
||||
// NewPostgresRepository builds a PgRepository.
|
||||
func NewPostgresRepository(pool *pgxpool.Pool) *PgRepository {
|
||||
return &PgRepository{pool: pool, q: codeindexsqlc.New(pool)}
|
||||
}
|
||||
|
||||
var _ Repository = (*PgRepository)(nil)
|
||||
|
||||
func pgUUID(u uuid.UUID) pgtype.UUID { return pgtype.UUID{Bytes: u, Valid: true} }
|
||||
|
||||
func pgUUIDPtr(u *uuid.UUID) pgtype.UUID {
|
||||
if u == nil {
|
||||
return pgtype.UUID{}
|
||||
}
|
||||
return pgtype.UUID{Bytes: *u, Valid: true}
|
||||
}
|
||||
|
||||
func uuidPtr(p pgtype.UUID) *uuid.UUID {
|
||||
if !p.Valid {
|
||||
return nil
|
||||
}
|
||||
id := uuid.UUID(p.Bytes)
|
||||
return &id
|
||||
}
|
||||
|
||||
func runFromRow(r codeindexsqlc.CodeIndexRun) *Run {
|
||||
out := &Run{
|
||||
ID: uuid.UUID(r.ID.Bytes),
|
||||
WorkspaceID: uuid.UUID(r.WorkspaceID.Bytes),
|
||||
WorktreeID: uuidPtr(r.WorktreeID),
|
||||
Branch: r.Branch,
|
||||
CommitSHA: r.CommitSha,
|
||||
Status: r.Status,
|
||||
EmbeddingEndpointID: uuidPtr(r.EmbeddingEndpointID),
|
||||
EmbeddingModel: r.EmbeddingModel,
|
||||
Dims: int(r.Dims),
|
||||
FilesIndexed: int(r.FilesIndexed),
|
||||
ChunksIndexed: int(r.ChunksIndexed),
|
||||
Error: r.Error,
|
||||
}
|
||||
if r.CreatedAt.Valid {
|
||||
out.CreatedAt = r.CreatedAt.Time
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (r *PgRepository) InsertRun(ctx context.Context, run *Run) (*Run, error) {
|
||||
row, err := r.q.InsertRun(ctx, codeindexsqlc.InsertRunParams{
|
||||
ID: pgUUID(run.ID),
|
||||
WorkspaceID: pgUUID(run.WorkspaceID),
|
||||
WorktreeID: pgUUIDPtr(run.WorktreeID),
|
||||
Branch: run.Branch,
|
||||
CommitSha: run.CommitSHA,
|
||||
Status: StatusPending,
|
||||
EmbeddingEndpointID: pgUUIDPtr(run.EmbeddingEndpointID),
|
||||
EmbeddingModel: run.EmbeddingModel,
|
||||
Dims: int32(run.Dims),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "insert code index run")
|
||||
}
|
||||
return runFromRow(row), nil
|
||||
}
|
||||
|
||||
func (r *PgRepository) GetRun(ctx context.Context, id uuid.UUID) (*Run, error) {
|
||||
row, err := r.q.GetRun(ctx, pgUUID(id))
|
||||
if err != nil {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, errs.New(errs.CodeNotFound, "code index run not found")
|
||||
}
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "get code index run")
|
||||
}
|
||||
return runFromRow(row), nil
|
||||
}
|
||||
|
||||
func (r *PgRepository) GetRunByCommit(ctx context.Context, wsID uuid.UUID, commitSHA, model string) (*Run, error) {
|
||||
row, err := r.q.GetRunByCommit(ctx, codeindexsqlc.GetRunByCommitParams{
|
||||
WorkspaceID: pgUUID(wsID),
|
||||
CommitSha: commitSHA,
|
||||
EmbeddingModel: model,
|
||||
})
|
||||
if err != nil {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "get run by commit")
|
||||
}
|
||||
return runFromRow(row), nil
|
||||
}
|
||||
|
||||
func (r *PgRepository) GetLatestCompletedRun(ctx context.Context, wsID uuid.UUID) (*Run, error) {
|
||||
row, err := r.q.GetLatestCompletedRun(ctx, pgUUID(wsID))
|
||||
if err != nil {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "get latest completed run")
|
||||
}
|
||||
return runFromRow(row), nil
|
||||
}
|
||||
|
||||
func (r *PgRepository) MarkRunning(ctx context.Context, id uuid.UUID) error {
|
||||
if err := r.q.MarkRunRunning(ctx, pgUUID(id)); err != nil {
|
||||
return errs.Wrap(err, errs.CodeInternal, "mark run running")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *PgRepository) MarkCompleted(ctx context.Context, id uuid.UUID, files, chunks int) error {
|
||||
if err := r.q.MarkRunCompleted(ctx, codeindexsqlc.MarkRunCompletedParams{
|
||||
ID: pgUUID(id),
|
||||
FilesIndexed: int32(files),
|
||||
ChunksIndexed: int32(chunks),
|
||||
}); err != nil {
|
||||
return errs.Wrap(err, errs.CodeInternal, "mark run completed")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *PgRepository) MarkFailed(ctx context.Context, id uuid.UUID, errMsg string) error {
|
||||
msg := errMsg
|
||||
if err := r.q.MarkRunFailed(ctx, codeindexsqlc.MarkRunFailedParams{
|
||||
ID: pgUUID(id),
|
||||
Error: &msg,
|
||||
}); err != nil {
|
||||
return errs.Wrap(err, errs.CodeInternal, "mark run failed")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *PgRepository) MarkStale(ctx context.Context, wsID uuid.UUID) error {
|
||||
if err := r.q.MarkRunsStale(ctx, pgUUID(wsID)); err != nil {
|
||||
return errs.Wrap(err, errs.CodeInternal, "mark runs stale")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *PgRepository) MarkStuckStale(ctx context.Context, runningBefore time.Time) ([]uuid.UUID, error) {
|
||||
rows, err := r.q.MarkStuckRunsStale(ctx, pgtype.Timestamptz{Time: runningBefore, Valid: true})
|
||||
if err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "mark stuck runs stale")
|
||||
}
|
||||
out := make([]uuid.UUID, 0, len(rows))
|
||||
for _, p := range rows {
|
||||
out = append(out, uuid.UUID(p.Bytes))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// ReplaceChunks deletes prior chunks for the run and bulk-inserts the given rows
|
||||
// in a single transaction. The embedding column uses the pgvector text literal.
|
||||
func (r *PgRepository) ReplaceChunks(ctx context.Context, runID, wsID uuid.UUID, commitSHA string, rows []ChunkRow) error {
|
||||
tx, err := r.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return errs.Wrap(err, errs.CodeInternal, "begin chunk tx")
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
|
||||
if _, err := tx.Exec(ctx, `DELETE FROM code_chunks WHERE run_id = $1`, pgUUID(runID)); err != nil {
|
||||
return errs.Wrap(err, errs.CodeInternal, "delete prior chunks")
|
||||
}
|
||||
|
||||
const insertSQL = `INSERT INTO code_chunks (
|
||||
id, run_id, workspace_id, commit_sha, file_path, lang,
|
||||
start_line, end_line, content, content_sha, token_count, embedding
|
||||
) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12::vector)`
|
||||
|
||||
batch := &pgx.Batch{}
|
||||
for _, row := range rows {
|
||||
var lang *string
|
||||
if row.Lang != "" {
|
||||
l := row.Lang
|
||||
lang = &l
|
||||
}
|
||||
var tok *int32
|
||||
if row.TokenCount > 0 {
|
||||
t := int32(row.TokenCount)
|
||||
tok = &t
|
||||
}
|
||||
batch.Queue(insertSQL,
|
||||
pgUUID(uuid.New()), pgUUID(runID), pgUUID(wsID), commitSHA,
|
||||
row.FilePath, lang, int32(row.StartLine), int32(row.EndLine),
|
||||
row.Content, row.ContentSHA, tok, pgvector.NewVector(row.Embedding),
|
||||
)
|
||||
}
|
||||
br := tx.SendBatch(ctx, batch)
|
||||
for range rows {
|
||||
if _, err := br.Exec(); err != nil {
|
||||
_ = br.Close()
|
||||
return errs.Wrap(err, errs.CodeInternal, "insert chunk")
|
||||
}
|
||||
}
|
||||
if err := br.Close(); err != nil {
|
||||
return errs.Wrap(err, errs.CodeInternal, "close chunk batch")
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return errs.Wrap(err, errs.CodeInternal, "commit chunk tx")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// SearchKNN runs cosine-distance KNN scoped to the latest completed run for the
|
||||
// workspace. Score is 1 - cosine_distance (so higher is more similar). Optional
|
||||
// path_prefix / lang filters narrow the candidate set.
|
||||
func (r *PgRepository) SearchKNN(ctx context.Context, wsID uuid.UUID, queryVec []float32, topK int, pathPrefix, lang string) ([]Snippet, error) {
|
||||
if topK <= 0 {
|
||||
topK = 10
|
||||
}
|
||||
q := `
|
||||
WITH latest AS (
|
||||
SELECT id FROM code_index_runs
|
||||
WHERE workspace_id = $1 AND status = 'completed'
|
||||
ORDER BY created_at DESC LIMIT 1
|
||||
)
|
||||
SELECT c.file_path, c.start_line, c.end_line, c.content, c.commit_sha,
|
||||
1 - (c.embedding <=> $2::vector) AS score
|
||||
FROM code_chunks c
|
||||
JOIN latest ON c.run_id = latest.id
|
||||
WHERE ($3::text = '' OR c.file_path LIKE $3 || '%')
|
||||
AND ($4::text = '' OR c.lang = $4)
|
||||
ORDER BY c.embedding <=> $2::vector
|
||||
LIMIT $5`
|
||||
rows, err := r.pool.Query(ctx, q, pgUUID(wsID), pgvector.NewVector(queryVec), pathPrefix, lang, int32(topK))
|
||||
if err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "knn search")
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []Snippet
|
||||
for rows.Next() {
|
||||
var s Snippet
|
||||
var start, end int32
|
||||
if err := rows.Scan(&s.FilePath, &start, &end, &s.Content, &s.CommitSHA, &s.Score); err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "scan knn row")
|
||||
}
|
||||
s.StartLine, s.EndLine = int(start), int(end)
|
||||
out = append(out, s)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "knn rows")
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
Reference in New Issue
Block a user