knowledge-base-API/services.py

from qcloud_cos import CosConfig
from qcloud_cos import CosS3Client
import tcvectordb
from tcvectordb.model.ai_database import AIDatabase
from tcvectordb.model.collection_view import CollectionView, Embedding, SplitterProcess, ParsingProcess
from tcvectordb.model.index import Index
from tcvectordb.exceptions import DescribeCollectionException

from settings import settings

cos = CosS3Client(
	CosConfig(
		Region=settings.region,
		SecretId=settings.secret_id,
		SecretKey=settings.secret_key,
	),
)

vdb = tcvectordb.RPCVectorDBClient(
	url=settings.VDB_config.url,
	username=settings.VDB_config.username,
	key=settings.VDB_config.key,
)


def create_ai_database_if_not_exists(database_name: str) -> AIDatabase:
	if vdb.exists_db(database_name):
		return vdb.database(database_name)
	else:
		return vdb.create_ai_database(database_name)

def create_collection_view_if_not_exists(database: AIDatabase, collection_name: str) -> CollectionView:
	try:
		return database.collection_view(collection_name)
	except DescribeCollectionException:
		return database.create_collection_view(
			collection_name,
			embedding=Embedding(
				language='MULTI',
				enable_words_embedding=True,	
			),
			splitter_process=SplitterProcess(
				append_title_to_chunk=False,
				append_keywords_to_chunk=True,
				chunk_splitter=None,
			),
			parsing_process=ParsingProcess('VisionModelParsing'),
			index=Index(),
		)

def create_db_collection_view_if_not_exists(database: str, collection: str) -> CollectionView:
	db = create_ai_database_if_not_exists(database)
	return create_collection_view_if_not_exists(db, collection)

collection_view = create_db_collection_view_if_not_exists(
	settings.VDB_config.database,
	settings.VDB_config.collection,
)