from pinecone import Pinecone
from dotenv import load_dotenv
import os

load_dotenv()
class PineconeRetriever:
    """Class for initializing and querying Pinecone."""
    def __init__(self, api_key, index_name, namespace):
        self.api_key = api_key
        self.index_name = index_name
        self.namespace = namespace
        self.pc = Pinecone(api_key=self.api_key)
        self.index  = self.pc.Index(self.index_name)
    
    def retrieve_data(self, query, top_k=1):
        """Retrieve relevant data from Pinecone."""
        # Convert the query into a numerical vector that Pinecone can search with
        query_embedding = self.pc.inference.embed(
            model="multilingual-e5-large",
            inputs=[query],
            parameters={
                "input_type": "query",
                "truncate": "END"
            }
        )

        results = self.index.query(namespace=self.namespace,
            vector=query_embedding[0].values,
            top_k=top_k,
            include_values=False,
            include_metadata=True
        )

        return results

    def retrieve_context(self, query, top_k=1):
        results = self.retrieve_data(query, top_k)
        context = ""
        for match in results['matches']:
            context += match['metadata']['source_text']
        return context

if __name__ == "__main__":
    pinecone_api=os.getenv("PINECONE_API_KEY"),
    pinecone_index=os.getenv("PINECONE_INDEX"),
    pinecone_namespace=os.getenv("PINECONE_NAMESPACE")
    
    retriever = PineconeRetriever(API_KEY, INDEX_NAME, NAMESPACE)
    query_result = retriever.retrieve_context("hi")
    print(query_result)