AzenKain
a8f0597e59
All checks were successful
Build and Release / release (push) Successful in 2m13s
139 lines
3.5 KiB
Go
139 lines
3.5 KiB
Go
// Code generated by sqlc. DO NOT EDIT.
|
|
// versions:
|
|
// sqlc v1.30.0
|
|
// source: rag.sql
|
|
|
|
package sqlc
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/jackc/pgx/v5/pgtype"
|
|
"github.com/pgvector/pgvector-go"
|
|
)
|
|
|
|
const createRagChunk = `-- name: CreateRagChunk :one
|
|
INSERT INTO rag_chunks (
|
|
id, source_type, source_id, project_id, chunk_index, content, embedding
|
|
) VALUES (
|
|
COALESCE($7::uuid, uuidv7()),
|
|
$1, $2, $3, $4, $5, $6
|
|
)
|
|
RETURNING id, source_type, source_id, project_id, chunk_index, content, embedding, created_at, updated_at
|
|
`
|
|
|
|
type CreateRagChunkParams struct {
|
|
SourceType string `json:"source_type"`
|
|
SourceID pgtype.UUID `json:"source_id"`
|
|
ProjectID pgtype.UUID `json:"project_id"`
|
|
ChunkIndex int32 `json:"chunk_index"`
|
|
Content string `json:"content"`
|
|
Embedding pgvector.Vector `json:"embedding"`
|
|
ID pgtype.UUID `json:"id"`
|
|
}
|
|
|
|
func (q *Queries) CreateRagChunk(ctx context.Context, arg CreateRagChunkParams) (RagChunk, error) {
|
|
row := q.db.QueryRow(ctx, createRagChunk,
|
|
arg.SourceType,
|
|
arg.SourceID,
|
|
arg.ProjectID,
|
|
arg.ChunkIndex,
|
|
arg.Content,
|
|
arg.Embedding,
|
|
arg.ID,
|
|
)
|
|
var i RagChunk
|
|
err := row.Scan(
|
|
&i.ID,
|
|
&i.SourceType,
|
|
&i.SourceID,
|
|
&i.ProjectID,
|
|
&i.ChunkIndex,
|
|
&i.Content,
|
|
&i.Embedding,
|
|
&i.CreatedAt,
|
|
&i.UpdatedAt,
|
|
)
|
|
return i, err
|
|
}
|
|
|
|
const deleteRagChunksBySourceIDs = `-- name: DeleteRagChunksBySourceIDs :exec
|
|
DELETE FROM rag_chunks
|
|
WHERE source_type = $1 AND source_id = ANY($2::uuid[])
|
|
`
|
|
|
|
type DeleteRagChunksBySourceIDsParams struct {
|
|
SourceType string `json:"source_type"`
|
|
Column2 []pgtype.UUID `json:"column_2"`
|
|
}
|
|
|
|
func (q *Queries) DeleteRagChunksBySourceIDs(ctx context.Context, arg DeleteRagChunksBySourceIDsParams) error {
|
|
_, err := q.db.Exec(ctx, deleteRagChunksBySourceIDs, arg.SourceType, arg.Column2)
|
|
return err
|
|
}
|
|
|
|
const searchRagChunks = `-- name: SearchRagChunks :many
|
|
SELECT
|
|
id, source_type, source_id, project_id, chunk_index, content,
|
|
(1 - (embedding <=> $1))::float8 AS similarity
|
|
FROM rag_chunks
|
|
WHERE 1=1
|
|
AND ($2::uuid IS NULL OR project_id = $2::uuid)
|
|
AND ($3::varchar IS NULL OR source_type = $3::varchar)
|
|
AND (1 - (embedding <=> $1))::float8 >= $4::float8
|
|
ORDER BY embedding <=> $1
|
|
LIMIT $5
|
|
`
|
|
|
|
type SearchRagChunksParams struct {
|
|
Embedding pgvector.Vector `json:"embedding"`
|
|
ProjectID pgtype.UUID `json:"project_id"`
|
|
SourceType pgtype.Text `json:"source_type"`
|
|
MatchThreshold float64 `json:"match_threshold"`
|
|
MatchCount int32 `json:"match_count"`
|
|
}
|
|
|
|
type SearchRagChunksRow struct {
|
|
ID pgtype.UUID `json:"id"`
|
|
SourceType string `json:"source_type"`
|
|
SourceID pgtype.UUID `json:"source_id"`
|
|
ProjectID pgtype.UUID `json:"project_id"`
|
|
ChunkIndex int32 `json:"chunk_index"`
|
|
Content string `json:"content"`
|
|
Similarity float64 `json:"similarity"`
|
|
}
|
|
|
|
func (q *Queries) SearchRagChunks(ctx context.Context, arg SearchRagChunksParams) ([]SearchRagChunksRow, error) {
|
|
rows, err := q.db.Query(ctx, searchRagChunks,
|
|
arg.Embedding,
|
|
arg.ProjectID,
|
|
arg.SourceType,
|
|
arg.MatchThreshold,
|
|
arg.MatchCount,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
items := []SearchRagChunksRow{}
|
|
for rows.Next() {
|
|
var i SearchRagChunksRow
|
|
if err := rows.Scan(
|
|
&i.ID,
|
|
&i.SourceType,
|
|
&i.SourceID,
|
|
&i.ProjectID,
|
|
&i.ChunkIndex,
|
|
&i.Content,
|
|
&i.Similarity,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
items = append(items, i)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return items, nil
|
|
}
|