cyberia/research/mimi/internal/persist/embedding.sql.go

go
ฯ€ 0.0% 184 words ยท 1 min 
// Code generated by sqlc. DO NOT EDIT.
// versions:
//   sqlc v1.29.0
// source: embedding.sql

package persist

import (
	"context"

	"github.com/jackc/pgx/v5/pgtype"
	"github.com/pgvector/pgvector-go"
)

const findCosine = `-- name: FindCosine :many
SELECT
    id,
    text,
    metadata
FROM
    embedding
WHERE
    embedding <=> $1
`

type FindCosineRow struct {
	ID       pgtype.UUID
	Text     string
	Metadata []byte
}

func (q *Queries) FindCosine(ctx context.Context, embedding pgvector.Vector) ([]FindCosineRow, error) {
	rows, err := q.db.Query(ctx, findCosine, embedding)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var items []FindCosineRow
	for rows.Next() {
		var i FindCosineRow
		if err := rows.Scan(&i.ID, &i.Text, &i.Metadata); err != nil {
			return nil, err
		}
		items = append(items, i)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return items, nil
}

const saveEmbedding = `-- name: SaveEmbedding :one
INSERT INTO
    embedding (text, metadata, embedding)
VALUES
    ($1, $2, $3)
RETURNING
    id
`

type SaveEmbeddingParams struct {
	Text      string
	Metadata  []byte
	Embedding pgvector.Vector
}

func (q *Queries) SaveEmbedding(ctx context.Context, arg SaveEmbeddingParams) (pgtype.UUID, error) {
	row := q.db.QueryRow(ctx, saveEmbedding, arg.Text, arg.Metadata, arg.Embedding)
	var id pgtype.UUID
	err := row.Scan(&id)
	return id, err
}

Graph