trpc-agent-go/examples/knowledge/sqlitevec_store_impl.go at main · coder-mtj/trpc-agent-go · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
//
// Tencent is pleased to support the open source community by making trpc-agent-go available.
//
// Copyright (C) 2025 Tencent.  All rights reserved.
//
// trpc-agent-go is licensed under the Apache License Version 2.0.
//
//

//go:build cgo && sqliteveccgo

package util

import (
	"trpc.group/trpc-go/trpc-agent-go/knowledge/vectorstore"
	"trpc.group/trpc-go/trpc-agent-go/knowledge/vectorstore/sqlitevec"
)

const (
	sqliteVecDSNEnvKey           = "SQLITEVEC_DSN"
	defaultSQLiteVecDSN          = "file:knowledge.sqlite?_busy_timeout=5000"
	sqliteVecTableEnvKey         = "SQLITEVEC_TABLE"
	defaultSQLiteVecTable        = "knowledge_documents"
	sqliteVecMetadataTableEnvKey = "SQLITEVEC_METADATA_TABLE"
	defaultSQLiteVecMetadata     = "knowledge_document_meta"
)

func newSQLiteVecStore() (vectorstore.VectorStore, error) {
	dsn := GetEnvOrDefault(sqliteVecDSNEnvKey, defaultSQLiteVecDSN)
	table := GetEnvOrDefault(
		sqliteVecTableEnvKey,
		defaultSQLiteVecTable,
	)
	metaTable := GetEnvOrDefault(
		sqliteVecMetadataTableEnvKey,
		defaultSQLiteVecMetadata,
	)

	return sqlitevec.New(
		sqlitevec.WithDSN(dsn),
		sqlitevec.WithTableName(table),
		sqlitevec.WithMetadataTableName(metaTable),
	)
}