from transformers import pipeline
from langchain.llms import CTransformers
from langchain.prompts import ChatPromptTemplate

# Text classification model for routing input
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

# Load LLM models using CTransformers
general_llm = CTransformers(
    model="./llama-2-7b.Q4_K_M.gguf", 
    model_type="llama",
    config={'max_new_tokens': 512, 'temperature': 0.7}
)

medical_llm = CTransformers(
    model="./BioMistral-7B.Q4_K_M.gguf", 
    model_type="llama",
    config={'max_new_tokens': 512, 'temperature': 0.7}
)

# Prompt template for generating responses
template = """
You are a versatile AI assistant that can provide both medical advice and help users with general concerns...
"""

# Compile the prompt template using LangChain
prompt = ChatPromptTemplate.from_template(template)

def route_llm(user_input):
    """
    Routes user input to the appropriate LLM (medical or general).

    Parameters:
        user_input (str): The user's input message.

    Returns:
        CTransformers: The selected LLM model (general or medical).
    """
    result = classifier(user_input, ['medical', 'general'])
    label = result['labels'][0] if 'labels' in result else 'general'
    return medical_llm if label == 'medical' else general_llm