chain_recommendations.py

import json
from typing import Dict
from langchain import PromptTemplate, LLMChain
from models import chat_model

# Enhanced Prompt Template
recommend_prompt_template = PromptTemplate(
    input_variables=["problems"],
    template=(
        "You are a helpful wellness recommendation system. Given the following problem severity percentages:\n"
        "{problems}\n\n"
        "Based on these strict rules:\n"
        "- If sleep_problem > 70: Recommend Sleep Improvement Package\n"
        "- If stress_problem > 70: Recommend Stress Reduction Package\n"
        "- If exercise_problem > 70: Recommend Exercise Enhancement Package\n"
        "- If all problems are between 30 and 70: Recommend Balanced Wellness Package\n"
        "- If no severe problems: Recommend General Wellness Package\n\n"
        "List the recommended wellness packages separated by commas. "
        "Do not include any additional text or explanation."
    )
)

recommend_chain = LLMChain(llm=chat_model, prompt=recommend_prompt_template)

def generate_recommendations(problems: Dict[str, float]) -> str:
    """
    Generates wellness package recommendations based on problem severity percentages.
    The function accepts a dictionary of problem severities and returns a 
    comma-separated string of recommended packages based on predefined rules.
    """
    recommendations = recommend_chain.run(problems=json.dumps(problems))
    return recommendations.strip()