57 lines
1.6 KiB
Python
57 lines
1.6 KiB
Python
import os
|
|
from typing import List, Optional
|
|
|
|
from ppt_generator.models.query_and_prompt_models import (
|
|
IconCategoryEnum,
|
|
IconQueryCollectionWithData,
|
|
)
|
|
from langchain_core.vectorstores import InMemoryVectorStore
|
|
|
|
|
|
async def get_icon(
|
|
vector_store: InMemoryVectorStore,
|
|
input: IconQueryCollectionWithData,
|
|
output_path: str,
|
|
) -> str:
|
|
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
|
|
|
query = input.icon_query.queries[0]
|
|
results = vector_store.similarity_search(query=query, k=1)
|
|
icon_name = results[0].page_content
|
|
|
|
with open(output_path, "wb") as f_a:
|
|
try:
|
|
with open(f"assets/icons/bold/{icon_name}.png", "rb") as f_b:
|
|
f_a.write(f_b.read())
|
|
except Exception as e:
|
|
print("Error finding icon: ", e)
|
|
with open(f"assets/icons/placeholder.png", "rb") as f_b:
|
|
f_a.write(f_b.read())
|
|
|
|
|
|
async def get_icons(
|
|
vector_store: InMemoryVectorStore,
|
|
query: str,
|
|
page: int,
|
|
limit: int,
|
|
category: Optional[IconCategoryEnum],
|
|
temp_dir: str,
|
|
) -> List[str]:
|
|
|
|
results = await vector_store.asimilarity_search(query=query, k=limit)
|
|
|
|
icon_names = [result.page_content for result in results]
|
|
|
|
icon_paths = [f"assets/icons/bold/{each}-bold.png" for each in icon_names]
|
|
|
|
icon_temp_paths = []
|
|
|
|
for each in icon_paths:
|
|
icon_temp_path = os.path.join(temp_dir, os.path.basename(each))
|
|
icon_temp_paths.append(icon_temp_path)
|
|
|
|
with open(icon_temp_path, "wb") as f_a:
|
|
with open(each, "rb") as f_b:
|
|
f_a.write(f_b.read())
|
|
|
|
return icon_temp_paths
|