knowledge-base-API/enpoints/retrieval.py

from typing import Dict
from pydantic import BaseModel

from main import app
from services import collection_view

class RetrievalSetting(BaseModel):
	top_k: int
	score_threshold: float
	embedding_instruct: str | None = None
	reranker_instruct: str | None = None

class Retrieval(BaseModel):
	knowledge_id: str
	query: str
	retrieval_setting: RetrievalSetting
	metadata_condition: Dict | None = None

@app.post('/retrieval')
def retrieval(req: Retrieval):
	chunks = collection_view.search(
		req.query,
		expand_chunk=[1, 1],
		limit=req.retrieval_setting.top_k,
	)
	chunks = [vars(i) for i in chunks]
	return {
		'records': [
			{
				'content': chunk['data']['text'],
				'score': chunk['score'],
				'title': chunk['documentSet']['documentSetName'],
				'metadata': {
					'document_id': str(chunk['documentSet']['documentSetId']),
				},
			} for chunk in chunks
		]
	}