You've already forked agentic-coding-workflow
bugfix
This commit is contained in:
@@ -0,0 +1,230 @@
|
||||
package projectmemory
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"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"
|
||||
|
||||
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
|
||||
projectmemorysqlc "github.com/yan1h/agent-coding-workflow/internal/projectmemory/sqlc"
|
||||
)
|
||||
|
||||
// Repository is the projectmemory persistence contract.
|
||||
type Repository interface {
|
||||
Insert(ctx context.Context, in WriteInput, embeddingModel *string) (*Entry, error)
|
||||
SetEmbedding(ctx context.Context, id uuid.UUID, vec []float32, model string) error
|
||||
Get(ctx context.Context, id uuid.UUID) (*Entry, error)
|
||||
ListByProject(ctx context.Context, projectID uuid.UUID, kind string) ([]Entry, error)
|
||||
SearchKeyword(ctx context.Context, projectID uuid.UUID, kind, text string, topK int) ([]Entry, error)
|
||||
SearchKNN(ctx context.Context, projectID uuid.UUID, vec []float32, kind string, topK int) ([]Entry, error)
|
||||
Delete(ctx context.Context, id uuid.UUID) error
|
||||
}
|
||||
|
||||
// PgRepository is the production Repository on a pgxpool.Pool.
|
||||
type PgRepository struct {
|
||||
pool *pgxpool.Pool
|
||||
q *projectmemorysqlc.Queries
|
||||
}
|
||||
|
||||
// NewPostgresRepository builds a PgRepository.
|
||||
func NewPostgresRepository(pool *pgxpool.Pool) *PgRepository {
|
||||
return &PgRepository{pool: pool, q: projectmemorysqlc.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
|
||||
}
|
||||
|
||||
// entryFields is the common projection shared by every sqlc row type in this
|
||||
// package; a tiny adapter keeps a single conversion path.
|
||||
type entryFields struct {
|
||||
ID pgtype.UUID
|
||||
ProjectID pgtype.UUID
|
||||
WorkspaceID pgtype.UUID
|
||||
Kind string
|
||||
Title string
|
||||
Body string
|
||||
Tags []string
|
||||
Source string
|
||||
EmbeddingModel *string
|
||||
CreatedBy pgtype.UUID
|
||||
CreatedAt pgtype.Timestamptz
|
||||
UpdatedAt pgtype.Timestamptz
|
||||
}
|
||||
|
||||
func toEntry(f entryFields) Entry {
|
||||
e := Entry{
|
||||
ID: uuid.UUID(f.ID.Bytes),
|
||||
ProjectID: uuid.UUID(f.ProjectID.Bytes),
|
||||
WorkspaceID: uuidPtr(f.WorkspaceID),
|
||||
Kind: f.Kind,
|
||||
Title: f.Title,
|
||||
Body: f.Body,
|
||||
Tags: f.Tags,
|
||||
Source: f.Source,
|
||||
EmbeddingModel: f.EmbeddingModel,
|
||||
CreatedBy: uuidPtr(f.CreatedBy),
|
||||
}
|
||||
if e.Tags == nil {
|
||||
e.Tags = []string{}
|
||||
}
|
||||
if f.CreatedAt.Valid {
|
||||
e.CreatedAt = f.CreatedAt.Time
|
||||
}
|
||||
if f.UpdatedAt.Valid {
|
||||
e.UpdatedAt = f.UpdatedAt.Time
|
||||
}
|
||||
return e
|
||||
}
|
||||
|
||||
func (r *PgRepository) Insert(ctx context.Context, in WriteInput, embeddingModel *string) (*Entry, error) {
|
||||
tags := in.Tags
|
||||
if tags == nil {
|
||||
// TEXT[] NOT NULL: never write a nil slice (would violate the constraint).
|
||||
tags = []string{}
|
||||
}
|
||||
source := in.Source
|
||||
if source == "" {
|
||||
source = SourceAgent
|
||||
}
|
||||
row, err := r.q.InsertEntry(ctx, projectmemorysqlc.InsertEntryParams{
|
||||
ID: pgUUID(uuid.New()),
|
||||
ProjectID: pgUUID(in.ProjectID),
|
||||
WorkspaceID: pgUUIDPtr(in.WorkspaceID),
|
||||
Kind: in.Kind,
|
||||
Title: in.Title,
|
||||
Body: in.Body,
|
||||
Tags: tags,
|
||||
Source: source,
|
||||
EmbeddingModel: embeddingModel,
|
||||
CreatedBy: pgUUIDPtr(in.CreatedBy),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "insert project memory")
|
||||
}
|
||||
e := toEntry(entryFields(row))
|
||||
return &e, nil
|
||||
}
|
||||
|
||||
// SetEmbedding writes the vector + model with a hand-written pgx UPDATE (the
|
||||
// embedding column is excluded from sqlc). updated_at is bumped.
|
||||
func (r *PgRepository) SetEmbedding(ctx context.Context, id uuid.UUID, vec []float32, model string) error {
|
||||
_, err := r.pool.Exec(ctx,
|
||||
`UPDATE project_memory SET embedding = $2::vector, embedding_model = $3, updated_at = now() WHERE id = $1`,
|
||||
pgUUID(id), pgvector.NewVector(vec), model)
|
||||
if err != nil {
|
||||
return errs.Wrap(err, errs.CodeInternal, "set memory embedding")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *PgRepository) Get(ctx context.Context, id uuid.UUID) (*Entry, error) {
|
||||
row, err := r.q.GetEntry(ctx, pgUUID(id))
|
||||
if err != nil {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, errs.New(errs.CodeNotFound, "memory entry not found")
|
||||
}
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "get memory entry")
|
||||
}
|
||||
e := toEntry(entryFields(row))
|
||||
return &e, nil
|
||||
}
|
||||
|
||||
func (r *PgRepository) ListByProject(ctx context.Context, projectID uuid.UUID, kind string) ([]Entry, error) {
|
||||
rows, err := r.q.ListByProject(ctx, projectmemorysqlc.ListByProjectParams{
|
||||
ProjectID: pgUUID(projectID),
|
||||
Column2: kind,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "list project memory")
|
||||
}
|
||||
out := make([]Entry, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
out = append(out, toEntry(entryFields(row)))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *PgRepository) SearchKeyword(ctx context.Context, projectID uuid.UUID, kind, text string, topK int) ([]Entry, error) {
|
||||
if topK <= 0 {
|
||||
topK = 10
|
||||
}
|
||||
rows, err := r.q.SearchByKeyword(ctx, projectmemorysqlc.SearchByKeywordParams{
|
||||
ProjectID: pgUUID(projectID),
|
||||
Column2: kind,
|
||||
Column3: text,
|
||||
Limit: int32(topK),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "search memory keyword")
|
||||
}
|
||||
out := make([]Entry, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
out = append(out, toEntry(entryFields(row)))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// SearchKNN runs cosine-KNN over entries that have an embedding, scoped to the
|
||||
// project (and optional kind). Hand-written because the vector column is excluded
|
||||
// from sqlc.
|
||||
func (r *PgRepository) SearchKNN(ctx context.Context, projectID uuid.UUID, vec []float32, kind string, topK int) ([]Entry, error) {
|
||||
if topK <= 0 {
|
||||
topK = 10
|
||||
}
|
||||
q := `
|
||||
SELECT id, project_id, workspace_id, kind, title, body, tags, source,
|
||||
embedding_model, created_by, created_at, updated_at
|
||||
FROM project_memory
|
||||
WHERE project_id = $1
|
||||
AND embedding IS NOT NULL
|
||||
AND ($3::text = '' OR kind = $3)
|
||||
ORDER BY embedding <=> $2::vector
|
||||
LIMIT $4`
|
||||
rows, err := r.pool.Query(ctx, q, pgUUID(projectID), pgvector.NewVector(vec), kind, int32(topK))
|
||||
if err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "search memory knn")
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []Entry
|
||||
for rows.Next() {
|
||||
var f entryFields
|
||||
if err := rows.Scan(&f.ID, &f.ProjectID, &f.WorkspaceID, &f.Kind, &f.Title,
|
||||
&f.Body, &f.Tags, &f.Source, &f.EmbeddingModel, &f.CreatedBy,
|
||||
&f.CreatedAt, &f.UpdatedAt); err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "scan memory knn row")
|
||||
}
|
||||
out = append(out, toEntry(f))
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "memory knn rows")
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *PgRepository) Delete(ctx context.Context, id uuid.UUID) error {
|
||||
if err := r.q.DeleteEntry(ctx, pgUUID(id)); err != nil {
|
||||
return errs.Wrap(err, errs.CodeInternal, "delete memory entry")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user