next

src/genai/analytics_chatbot/utils/nodes.py

@@ -38,17 +38,17 @@ class RetrievePossibleEndpoints:
     def __init__(self):
         self.df = api_knowledge_df
         self.index = api_index
-        # self.results = []
-        self.results = ['/api/v1/compare/', '/api/v1/engagement/basic-metrics', '/api/v1/content/hashtags-analysis', '/api/v1/audience/emoji-count', '/api/v1/engagement/temporal_analysis']
+        self.results = []
+        # self.results = ['/api/v1/compare/', '/api/v1/engagement/basic-metrics', '/api/v1/content/hashtags-analysis', '/api/v1/audience/emoji-count', '/api/v1/engagement/temporal_analysis']
 
     def run(self,state:State):
         print('Gone to retrieve possible endpoints')
-        # query_embedding = np.array(embedding_model.embed_query(state['latest_message'])).reshape(1, -1).astype('float32')
-        # distances, indices = self.index.search(query_embedding, 5)
-        # for idx in indices[0]:
-        #     row = self.df.iloc[idx]
-        #     print('Endpoint:',row['endpoint'])
-        #     self.results.append(row['endpoint'])
+        query_embedding = np.array(embedding_model.embed_query(state['latest_message'])).reshape(1, -1).astype('float32')
+        distances, indices = self.index.search(query_embedding, 5)
+        for idx in indices[0]:
+            row = self.df.iloc[idx]
+            print('Endpoint:',row['endpoint'])
+            self.results.append(row['endpoint'])
         print('The possible endpoints are:', self.results)
         return {
             "possible_endpoints": self.results,