RAG_class.py

#import recurive textsplitter
from sentence_transformers import SentenceTransformer
import chromadb
import uuid
import os
from dotenv import load_dotenv
load_dotenv()
from openai import OpenAI

class RAG_1177:
    def __init__(self):
        self.db_name = "RAG_1177"
        
        self.model = SentenceTransformer('KBLab/sentence-bert-swedish-cased')
        self.client = chromadb.PersistentClient(path="RAG_1177_db")
        self.db = self.client.get_or_create_collection(self.db_name)

        self.example_questions  = [
            "hur kan jag lindra min huvudvärk?",
            "Vad är symptomen för covid-19?",
            "Vilka läkemedel finns det för smärtlindring?",
            "Hur kan jag bäst lindra mina symtom på pollenallergi?",
            "Vad kan jag göra för att förbättra min sömn?",
            "vad kan jag göra om mitt barn får magsjuka?"]

    def get_embeddings(self, text):
        embeddings = self.model.encode(text)
        return (embeddings.tolist())
    
    def get_ids(self, num_ids):
        ids = [str(uuid.uuid4()) for _ in range(num_ids)]
        return ids
    
    def delete_collection(self):
        self.client.delete_collection(self.db_name)
        return
    
    def retrieve(self, query, num_results):
        query_emb = self.get_embeddings(query)
        result = self.db.query(query_embeddings=query_emb, n_results=num_results, include=['documents', 'metadatas'])
        result_urls = result['metadatas'][0]
        result_docs = result['documents'][0]

        url_list = set([item['url'] for item in result_urls])
        result_urls = "Läs mer på:\n"
        for i, url in enumerate(url_list, start=1):
            result_urls += f"{i}: {url}\n"

        return result_docs, result_urls
    
    def insert(self,docs, emb, urls, ids):
        self.db.add(documents=docs, embeddings=emb, metadatas=urls, ids=ids)
        return 

    def relevant_question(self, query):
        history = []
        system_prompt = """Du ska bedömma frågor och svara endast JA eller NEJ.

        Om frågan är relaterad till 1177.se sjukrådgivningsdatabas, hälsa, vård, välmående, sjukrådgivning ska du svara "JA".

        Om frågan inte är relaterad till ovanämnde ämnen ska du svara NEJ"""

        user_prompt = f"""Svara endast JA eller NEJ ifall följande fråga är är relevant till 1177.se, sjukrådgivning, hälsa, vård och välmående

        Frågan är: {query}"""
            
        history.append({"role": "system", "content": system_prompt})
        history.append({"role": "user", "content": user_prompt})

        client = OpenAI()
        response = client.chat.completions.create(
                model='gpt-4o',
                messages=history,
                temperature=0.4,
                stream=False
            )
        return (response.choices[0].message.content)
    

    def rag_user_prompt(self, query, num_docs):
        documents, urls = self.retrieve(query, num_docs)

        user_prompt = f"""Ge detaljerade och tydliga svar på frågan baserat på informationen ovan och ange länkarna snyggt formaterat i slutet av svaret.
        Om frågan är en uppföljningsfråga kan du strunta i den givna informationen. Svara inte på frågan om det inte är relaterat till hälsa och sjukvård.
        
        Fråga: {query}

        Information från 1177.se:
        {documents[0]}

        Referenser:
        {urls}
        """
        return user_prompt
     
    def system_prompt(self):
        prompt = """
        Välkommen till 1177 AI-chatbot!

        Jag är här för att hjälpa dig med sjukvårdsfrågor genom att ge detaljerade svar baserade på dokumentation från 1177.se. Nedan ser du de olika typerna av frågor jag kan svara på:
        
        1. **Relevanta frågor med svar från 1177:** Jag kan ge detaljerade svar och inkludera tydliga länkar när frågan kan besvaras med information från 1177.se.

        2. **Relevanta följdfrågor:** Om du har uppföljningsfrågor kan jag använda tidigare svar som kontext för att ge mer information."


        Låt mig veta hur jag kan hjälpa dig!"""
        return prompt

    def get_description(self):
        description = """<ul>
    <li><strong>Specialiserad:</strong> Designad för att endast ge svar på sjukvårdsfrågor som man ofta besöker 1177 för.</li>
    <li><strong>Svar med referenser:</strong> Alla svar är baserade på data från <a href="https://www.1177.se/" target="_blank">1177-hemsidan</a> och innehåller källhänvisningar.</li>

    <p>För frågor & feedback: hassan.saado94@gmail.com, <a href="https://www.linkedin.com/in/hassan-saado/">Linkedin</a></li></p>"""
        return description