rag_utils.py

"""
RAG utilities for retrieving relevant medical information from Pinecone.
"""
import os
from pinecone import Pinecone
from sentence_transformers import SentenceTransformer
from dotenv import load_dotenv

load_dotenv()

PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
INDEX_NAME = os.getenv("PINECONE_INDEX_NAME", "doctor-assist")

if not PINECONE_API_KEY:
    raise ValueError("PINECONE_API_KEY not found in environment variables. Please check your .env file.")

pc = Pinecone(api_key=PINECONE_API_KEY)
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')

def get_relevant_context(query, top_k=5):
    """
    Retrieve relevant medical context from Pinecone based on query.
    
    Args:
        query: The search query
        top_k: Number of relevant chunks to retrieve
    
    Returns:
        List of relevant text chunks
    """
    try:
        index = pc.Index(INDEX_NAME)
        
        query_embedding = embedding_model.encode(query).tolist()
        
        results = index.query(
            vector=query_embedding,
            top_k=top_k,
            include_metadata=True
        )
        
        contexts = []
        for match in results.matches:
            if match.metadata and 'text' in match.metadata:
                contexts.append(match.metadata['text'])
        
        return contexts
    except Exception as e:
        print(f"Error retrieving context: {e}")
        return []

def format_context_for_prompt(contexts):
    """Format retrieved contexts into a prompt-friendly string."""
    if not contexts:
        return "No relevant medical information found."
    
    formatted = "Relevant Medical Information:\n\n"
    for i, context in enumerate(contexts, 1):
        formatted += f"{i}. {context}\n\n"
    
    return formatted