Added the retrieval tool node

my_agent/utils/tools.py

@@ -4,17 +4,13 @@ from dotenv import load_dotenv
 load_dotenv()
 import os
 import numpy as np
+from langchain_core.tools import tool
+from .data_loader import load_influencer_data
+from .models_loader import  ST
 
-os.environ['GROQ_API_KEY']=os.getenv('GROQ_API_KEY')
 
-llm = ChatGroq(
-    model="llama3-8b-8192",
-    temperature=0,
-    max_tokens=None,
-    timeout=None,
-    max_retries=2,
+os.environ['GROQ_API_KEY']=os.getenv('GROQ_API_KEY')
 
-)
 
 class StoryFormatter(BaseModel):
     """Always use this tool to structure your response to the user."""
@@ -30,4 +26,20 @@ class BrainstromTopicFormatter(BaseModel):
   topic1:str=Field(description="First brainstorming topic of the story")
   topic2:str=Field(description="Second brainstorming topic of the story")
   topic3:str=Field(description="Third brainstorming topic of the story")
-  topic4:str=Field(description="Fourth brainstorming topic of the story")
+  topic4:str=Field(description="Fourth brainstorming topic of the story")
+
+class QueryFormatter(BaseModel):
+    idea:str = Field(description="The video idea which the user wants to create.")
+    business_details: str = Field(description="The details of the business of that user.")
+
+@tool("influencer's data-retrieval-tool", args_schema=QueryFormatter, return_direct=False,description="Retrieve influencer-related data for a given query.")
+def retrieve_tool(idea, business_details):
+
+    """This tool is responsible for the retrieval of the influencer's data using semantic search by reading the video idea and the business details of the user. """
+    embedded_query = ST.encode(str(idea)+str(business_details))  # Embed each topic
+    data = load_influencer_data()
+    scores, retrieved_examples = data.get_nearest_examples("embeddings", embedded_query, k=1)
+
+    # Construct a list of dictionaries for this topic
+    result = [{user: story} for user, story in zip(retrieved_examples['username'], retrieved_examples['agentic_story'])]
+    return result