import yaml
from smolagents import CodeAgent, HfApiModel
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Simplified system prompt without any code-related instructions
system_prompt = (
    "You are a health and lifestyle advisor specializing in the early detection and prevention of hypertension. "
    "Provide direct and concise lifestyle advice based on the user's details. "
    "Output a clear answer as plain text with no extra commentary."
)

def clean_response(text):
    """
    Simple cleanup function that removes extra whitespace and ensures proper formatting.
    """
    # Remove extra whitespace
    text = ' '.join(text.split())
    # Split into paragraphs for readability
    paragraphs = text.split('\n\n')
    cleaned_paragraphs = [p.strip() for p in paragraphs if p.strip()]
    return '\n\n'.join(cleaned_paragraphs)

# Load prompt templates from YAML
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

# Initialize the model with simplified settings
model = HfApiModel(
    max_tokens=1024,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct'
)

# Initialize CodeAgent with minimal configuration
agent = CodeAgent(
    model=model,
    tools=[FinalAnswerTool()],
    max_steps=1,
    verbosity_level=0,
    name="Health Advisor",
    description=system_prompt,
    prompt_templates=prompt_templates
)

def run_agent(user_input):
    """
    Runs the agent and returns a clean, formatted response.
    """
    try:
        response = agent.run(user_input)
        return clean_response(response)
    except Exception as e:
        return f"I apologize, but I couldn't process your request. Please try again."

# Launch the Gradio UI
GradioUI(agent).launch()