File size: 1,274 Bytes
f315fdc |
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 |
# modules/retriever.py
from typing import List, Dict
from openai import OpenAI
def retrieve_videos(query: str, collection, top_k: int = 3) -> List[Dict]:
client = OpenAI()
# Create embedding for query
embedding = client.embeddings.create(
input=query,
model="text-embedding-3-small"
).data[0].embedding
# Query Chroma
results = collection.query(
query_embeddings=[embedding],
n_results=top_k,
include=["metadatas", "documents", "distances"]
)
# Build list of standardized dicts
videos = []
metadatas_list = results.get("metadatas", [[]])[0] # list of metadata dicts
documents_list = results.get("documents", [[]])[0] # list of text
distances_list = results.get("distances", [[]])[0] # optional
for idx, meta in enumerate(metadatas_list):
videos.append({
"video_id": meta.get("video_id", ""),
"video_title": meta.get("video_title", meta.get("title", documents_list[idx])),
"channel": meta.get("channel", meta.get("channel_title", "")),
"description": documents_list[idx] if idx < len(documents_list) else "",
"score": distances_list[idx] if idx < len(distances_list) else None
})
return videos
|