src/genai/analytics_chatbot/utils/tools.py

import numpy as np
from src.genai.utils.data_loader import api_knowledge_df, api_index
from src.genai.utils.models_loader import embedding_model

class APIKnowledgeRetrieveTool:
    def __init__(self):
        self.df = api_knowledge_df
        self.index = api_index

    def retrieve(self,query):
        query_embedding = np.array(embedding_model.embed_query(query)).reshape(1, -1).astype('float32')
        distances, indices = self.index.search(query_embedding, 1)
        row=self.df.iloc[indices[0]]
        data = {'endpoint':row['endpoint'],
                'method':row['method'],
                'parameters':row['parameters']}
        return data