from typing import Dict
from pydantic import BaseModel
from 

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,
	)
	return {
		'records': [
			{
				'content': chunk.data.text,
				'score': chunk.score,
				'title': chunk.data.documentSet.documentSetName,
				'metadata': {
					'document_id': str(chunk.data.documentSet.documentSetId),
				},
			} for chunk in chunks
		]
	}