retrieval updated in ideation

src/genai/ideation_agent/utils/tools.py

@@ -7,7 +7,7 @@ import ast
 import faiss
 import tiktoken
 from src.genai.utils.models_loader import embedding_model
-from src.genai.utils.load_embeddings import caption_embeddings , caption_index , caption_df
+from src.genai.utils.load_embeddings import  caption_index , caption_df, ideas_index , ideas_df
 from src.genai.utils.utils import clean_text
 
 @tool("influencers_data_retrieval_tool", args_schema=QueryFormatter, return_direct=False,description="Retrieve influencer-related data for a given query.")
@@ -55,3 +55,30 @@ def retrieve_tool(business_details):
     return encoding.decode(trimmed_response)
 
 
+@tool("imdb_movies_ideas_retrieval_tool", args_schema=QueryFormatter, return_direct=False,description="Retrieve imdb movies-related idea for a given query.")
+def retrieve_tool(business_details):
+    '''
+    Always invoke this tool.
+    Retrieve the ideas of imdb_movies by semantic search of **business details**.
+    '''
+    query_embedding = np.array(embedding_model.embed_query(str(business_details))).reshape(1, -1).astype('float32')
+    faiss.normalize_L2(query_embedding)
+
+    top_k = 5
+    distances, indices = ideas_index.search(query_embedding, top_k)
+
+    outer_list = []
+    for rank, (idx, sim) in enumerate(indices[0], 1):
+        row = ideas_df.iloc[idx]
+        res = {
+            'rank': rank,
+            'idea': row['idea'],
+        }
+
+        inner_list = [
+            f"[{res['rank']}]. The retrieved idea  is: **{res['idea']}\n**",
+        ]
+        outer_list.append(inner_list)
+
+    cleaned_response = clean_text(str(outer_list))
+    return str(cleaned_response)