You've already forked agentic-coding-workflow
231 lines
6.9 KiB
Go
231 lines
6.9 KiB
Go
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
|
|
}
|