54 lines
1.6 KiB
Go
54 lines
1.6 KiB
Go
|
package teachingassets
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"fmt"
|
||
|
|
|
||
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
||
|
|
)
|
||
|
|
|
||
|
|
type PostgresStore struct {
|
||
|
|
pool *pgxpool.Pool
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewPostgresStore(pool *pgxpool.Pool) *PostgresStore {
|
||
|
|
return &PostgresStore{pool: pool}
|
||
|
|
}
|
||
|
|
|
||
|
|
func toJSON(v any) string {
|
||
|
|
b, _ := json.Marshal(v)
|
||
|
|
return string(b)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *PostgresStore) SavePrompt(prompt PromptCandidate) (PromptCandidate, error) {
|
||
|
|
_, err := s.pool.Exec(context.Background(),
|
||
|
|
`INSERT INTO teaching_asset_prompts (id, concept_id, asset_type, prompt, model_key, review_state, requires_model_id_verification, source_evidence, created_at)
|
||
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8::jsonb, $9)`,
|
||
|
|
prompt.ID, prompt.Concept.ID, string(prompt.AssetType), prompt.Prompt,
|
||
|
|
prompt.ModelKey, string(prompt.ReviewState), prompt.RequiresModelIDVerification,
|
||
|
|
toJSON(prompt.SourceEvidence), prompt.CreatedAt,
|
||
|
|
)
|
||
|
|
if err != nil {
|
||
|
|
return PromptCandidate{}, fmt.Errorf("insert prompt: %w", err)
|
||
|
|
}
|
||
|
|
return prompt, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *PostgresStore) Snapshot() Snapshot {
|
||
|
|
rows, _ := s.pool.Query(context.Background(),
|
||
|
|
`SELECT id, concept_id, asset_type, prompt, model_key, review_state, requires_model_id_verification, source_evidence, created_at FROM teaching_asset_prompts`)
|
||
|
|
defer rows.Close()
|
||
|
|
|
||
|
|
var prompts []PromptCandidate
|
||
|
|
for rows.Next() {
|
||
|
|
var p PromptCandidate
|
||
|
|
var evidenceJSON string
|
||
|
|
rows.Scan(&p.ID, &p.Concept.ID, &p.AssetType, &p.Prompt, &p.ModelKey,
|
||
|
|
&p.ReviewState, &p.RequiresModelIDVerification, &evidenceJSON, &p.CreatedAt)
|
||
|
|
json.Unmarshal([]byte(evidenceJSON), &p.SourceEvidence)
|
||
|
|
prompts = append(prompts, p)
|
||
|
|
}
|
||
|
|
return Snapshot{Prompts: prompts}
|
||
|
|
}
|