malepati/custom_template_working
0
1import asyncio2import json3import chromadb4from chromadb.config import Settings5from chromadb.utils.embedding_functions import ONNXMiniLM_L6_V26 7 8class IconFinderService:9 def __init__(self):10 self.collection_name = "icons"11 self.client = chromadb.PersistentClient(12 path="chroma", settings=Settings(anonymized_telemetry=False)13 )14 print("Initializing icons collection...")15 self._initialize_icons_collection()16 print("Icons collection initialized.")17 18 def _initialize_icons_collection(self):19 self.embedding_function = ONNXMiniLM_L6_V2()20 self.embedding_function.DOWNLOAD_PATH = "chroma/models"21 self.embedding_function._download_model_if_not_exists()22 try:23 self.collection = self.client.get_collection(24 self.collection_name, embedding_function=self.embedding_function25 )26 except Exception:27 with open("assets/icons.json", "r") as f:28 icons = json.load(f)29 30 documents = []31 ids = []32 33 for i, each in enumerate(icons["icons"]):34 if each["name"].split("-")[-1] == "bold":35 doc_text = f"{each['name']} {each['tags']}"36 documents.append(doc_text)37 ids.append(each["name"])38 39 if documents:40 self.collection = self.client.create_collection(41 name=self.collection_name,42 embedding_function=self.embedding_function,43 metadata={"hnsw:space": "cosine"},44 )45 self.collection.add(documents=documents, ids=ids)46 47 async def search_icons(self, query: str, k: int = 1):48 result = await asyncio.to_thread(49 self.collection.query,50 query_texts=[query],51 n_results=k,52 )53 return [f"/static/icons/bold/{each}.svg" for each in result["ids"][0]]54 55 56try:57 ICON_FINDER_SERVICE = IconFinderService()58except Exception as e:59 print(f"Warning: IconFinderService initialization failed: {e}")60 print("Icon search functionality will be disabled.")61 ICON_FINDER_SERVICE = None62 