services/qa_service/qna.py

import json

from services.qa_service.utils import format_prompt

class QAService:
    def __init__(self, conf, pinecone, model_pipeline, question, goals):
        self.conf = conf
        self.pc = pinecone['connection']
        self.pc_index = self.pc.Index(self.conf['embeddings']['index_name'])
        self.embedder = pinecone['embedder']
        self.model_pipeline = model_pipeline
        self.question = question
        self.goals = goals
    
    def __enter__(self):
        print("Start Q&A Service")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        print("Exiting Q&A Service")

    def parse_results(self, result):

        parsed = []
        for i in result:

            level_1 = i['matches']
            level_2 = level_1['metadata']['_node_content']
            collect = json.loads(level_2)["text"]
            parsed.append(collect)
            
        return collect
    
    def retrieve_context(self):
        """Pass embedded question into pinecone"""
        embedded_query = self.embedder.get_text_embedding(self.question)
        
        result = self.pc_index.query(
            vector=embedded_query,
            top_k=5,
            include_values=False,
            include_metadata=True
        )

        #output = self.parse_results(result)
        output = result
        return output
    
    def run(self):
        """Query pinecone outputs and infer results"""
        print(self.retrieve_context())
        context = '\n'.join(self.retrieve_context())
        prompt = format_prompt(self.question, context)
        output = self.model_pipeline.infer(prompt)
        
        return output, context