import uuid from typing import List, Dict from qdrant_client import QdrantClient, models as qmodels from llama_index.llms.openai import OpenAI from fastembed import TextEmbedding from models import FoodItem from utils import synthesize_food_item likes = ["dosa", "fanta", "croissant", "waffles"] dislikes = ["virgin mojito"] menu = ["croissant", "mango", "jalebi"] class RecommendationEngine: def __init__( self, category: str, qdrant: QdrantClient, fastembed_model: TextEmbedding ) -> None: self.collection = f"{category}_preferences" self.qdrant = qdrant self.embedding_model = fastembed_model if self.qdrant.collection_exists(self.collection): self.counter = self.qdrant.count(self.collection, exact=True).count else: self.reset() self.counter = 0 def reset(self): self.qdrant.recreate_collection( self.collection, vectors_config=qmodels.VectorParams( size=384, distance=qmodels.Distance.COSINE ), ) def _generate_vector(self, model_json: dict): embedding_txt = "" for k, v in model_json.items(): embedding_txt += f"{k}: {v}" return list(self.embedding_model.passage_embed([embedding_txt]))[0] def _insert_preference(self, item: FoodItem, *args, **kwargs): model_json: dict = item.model_dump() embedding = self._generate_vector(model_json) model_json.update(kwargs) self.qdrant.upsert( self.collection, points=[ qmodels.PointStruct( id=self.counter, payload=model_json, vector=embedding ) ], ) self.counter += 1 def like(self, item: FoodItem): self._insert_preference(item, liked=True) def dislike(self, item: FoodItem): self._insert_preference(item, liked=False) def recommend_from_given( self, items: List[FoodItem], limit: int = 3 ) -> Dict[str, float]: liked_points, _offset = self.qdrant.scroll( self.collection, scroll_filter={"must": [{"key": "liked", "match": {"value": True}}]}, ) disliked_points, _offset = self.qdrant.scroll( self.collection, scroll_filter={"must": [{"key": "liked", "match": {"value": False}}]}, ) # Insert points in DB so they can be recommended: # A bit ugly but this is the best possible thing at the moment. query_id = str(uuid.uuid1()) for item in items: self._insert_preference(item, query_id=query_id) scored_points = self.qdrant.recommend( self.collection, positive=[p.id for p in liked_points], negative=[p.id for p in disliked_points], query_filter={"must": [{"key": "query_id", "match": {"value": query_id}}]}, with_payload=True, strategy="best_score", ) self.qdrant.delete(self.collection, [p.id for p in scored_points]) return {point.payload["name"]: point.score for point in scored_points} if __name__ == "__main__": llm = OpenAI(model="gpt-3.5-turbo") qdrant = QdrantClient() fastembed_model = TextEmbedding() rec_engine = RecommendationEngine("food", qdrant, fastembed_model) if rec_engine.counter != len(likes) + len(dislikes): rec_engine.reset() print("Filling with starter data") for food_name in likes: food_item = synthesize_food_item(food_name, llm) rec_engine.like(food_item) for food_name in dislikes: food_item = synthesize_food_item(food_name, llm) rec_engine.dislike(food_item) new_items = [synthesize_food_item(food_name, llm) for food_name in menu] recommendations = rec_engine.recommend_from_given(items=new_items) print(recommendations)