app.py

from smolagents import CodeAgent, Tool, DuckDuckGoSearchTool, load_tool, tool
from models import OpenAIModel  # Import our local model
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
import re
import os
import logging
from jinja2 import Template, StrictUndefined

from Gradio_UI import GradioUI

# Configure logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def calculate_min_price(prices: list[float])-> str: #it's import to specify the return type
    """A tool that calculates the min price from list of product prices
    Args:
        prices: list of product prices of
    """
    min_price =min(prices)
    return f"The minimum price is {min_price}"

@tool
def extract_price_from_snippet(snippet: str) -> list[str]:
    """
    A simple function to extract prices from a text snippet using regex.
    You can enhance this function for more complex price extraction.
    Args:
        snippet: text of all prices
    """
    # A basic regular expression to detect common price formats like $29.99, 29.99 USD, etc.
    price_pattern = r'\$\d+(?:,\d{3})*(?:\.\d{2})?|\d+(?:,\d{3})*(?:\.\d{2})?\s*(USD|EUR|GBP|INR|AUD|CAD)?'
    matches = re.findall(price_pattern, snippet)
    matches = [str(x) for x in matches]
    return matches


@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    try:
        # Create timezone object
        tz = pytz.timezone(timezone)
        # Get current time in that timezone
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"


final_answer = FinalAnswerTool()

# Load configuration from agent.json
with open("agent.json", 'r') as f:
    config = yaml.safe_load(f)

# Initialize OpenAI model using the configuration
model = OpenAIModel(
    model=config["model"]["data"]["model_id"],
    max_tokens=config["model"]["data"]["max_tokens"],
    temperature=config["model"]["data"]["temperature"],
    api_key=os.getenv("OPENAI_API_KEY")
)

# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

with open("prompts.yaml", 'r') as stream:
    yaml_content = yaml.safe_load(stream)
    if not isinstance(yaml_content, dict):
        raise ValueError("YAML content must be a dictionary")
    if 'prompt_templates' not in yaml_content:
        raise ValueError("YAML must contain 'prompt_templates' key")
    prompt_templates = yaml_content['prompt_templates']
    if not isinstance(prompt_templates, dict):
        raise ValueError("prompt_templates must be a dictionary")
    
    # Process templates to ensure they are valid Jinja2 templates
    def process_template(value, is_system_prompt=False):
        if isinstance(value, str):
            logger.debug(f"Processing template string: {value[:100]}...")
            try:
                # Log the exact content before processing
                logger.debug(f"Raw template content: {value}")
                
                # For system prompt, only ensure proper spacing in template expressions
                if is_system_prompt:
                    value = re.sub(r'{{([^{}]+)}}', lambda m: '{{ ' + m.group(1).strip() + ' }}', value)
                    value = re.sub(r'{%([^{}]+)%}', lambda m: '{% ' + m.group(1).strip() + ' %}', value)
                    value = value.replace('\n', '\\n')
                else:
                    # For other templates, wrap non-template content
                    if '{{' not in value and '{%' not in value:
                        value = f'{{{{ "{value}" }}}}'
                    else:
                        # Ensure proper spacing in template expressions
                        value = re.sub(r'{{([^{}]+)}}', lambda m: '{{ ' + m.group(1).strip() + ' }}', value)
                        value = re.sub(r'{%([^{}]+)%}', lambda m: '{% ' + m.group(1).strip() + ' %}', value)
                    value = value.replace('\n', '\\n')
                
                logger.debug(f"Processed template content: {value}")
            except Exception as e:
                logger.error(f"Error processing template string: {str(e)}")
                raise
            return value
        elif isinstance(value, dict):
            logger.debug(f"Processing template dictionary with keys: {list(value.keys())}")
            try:
                return {k: process_template(v, is_system_prompt=(k == 'system_prompt')) for k, v in value.items()}
            except Exception as e:
                logger.error(f"Error processing template dictionary: {str(e)}")
                raise
        return value
    
    logger.debug("Starting template processing...")
    try:
        # Log the original templates
        logger.debug("Original templates:")
        for key, value in prompt_templates.items():
            if isinstance(value, dict):
                logger.debug(f"{key}: {list(value.keys())}")
                for subkey, subvalue in value.items():
                    logger.debug(f"{key}.{subkey}: {str(subvalue)[:200]}...")
            else:
                logger.debug(f"{key}: {str(value)[:200]}...")

        prompt_templates = process_template(prompt_templates)
        logger.debug("Template processing completed")
        
        # Log the final processed templates
        logger.debug("Final processed templates:")
        for key, value in prompt_templates.items():
            if isinstance(value, dict):
                logger.debug(f"{key}: {list(value.keys())}")
                for subkey, subvalue in value.items():
                    logger.debug(f"{key}.{subkey}: {str(subvalue)[:200]}...")
            else:
                logger.debug(f"{key}: {str(value)[:200]}...")
    except Exception as e:
        logger.error(f"Error during template processing: {str(e)}")
        raise

    # Validate templates before creating agent
    logger.debug("Validating templates...")
    try:
        for key, value in prompt_templates.items():
            if isinstance(value, dict):
                for subkey, subvalue in value.items():
                    if isinstance(subvalue, str):
                        logger.debug(f"Validating template: {key}.{subkey}")
                        logger.debug(f"Template content: {subvalue[:200]}...")
                        # Test template compilation
                        template = Template(subvalue, undefined=StrictUndefined)
                        # Test template rendering with dummy data
                        try:
                            rendered = template.render(tools=[], task="test", name="test", final_answer="test", remaining_steps=1, answer_facts="test")
                            logger.debug(f"Template render test successful for {key}.{subkey}")
                            logger.debug(f"Rendered content: {rendered[:200]}...")
                        except Exception as e:
                            logger.error(f"Template render test failed for {key}.{subkey}: {str(e)}")
                            raise
            elif isinstance(value, str):
                logger.debug(f"Validating template: {key}")
                logger.debug(f"Template content: {value[:200]}...")
                template = Template(value, undefined=StrictUndefined)
                try:
                    rendered = template.render(tools=[], task="test", name="test", final_answer="test", remaining_steps=1, answer_facts="test")
                    logger.debug(f"Template render test successful for {key}")
                    logger.debug(f"Rendered content: {rendered[:200]}...")
                except Exception as e:
                    logger.error(f"Template render test failed for {key}: {str(e)}")
                    raise
        logger.debug("Template validation completed successfully")
    except Exception as e:
        logger.error(f"Template validation failed: {str(e)}")
        raise

    # Pre-process the system prompt template
    if 'system_prompt' in prompt_templates and 'text' in prompt_templates['system_prompt']:
        system_prompt = prompt_templates['system_prompt']['text']
        logger.debug(f"Original system prompt: {system_prompt[:200]}...")
        # Only process template variables, leave other text unchanged
        system_prompt = re.sub(r'{{([^{}]+)}}', lambda m: '{{ ' + m.group(1).strip() + ' }}', system_prompt)
        system_prompt = re.sub(r'{%([^{}]+)%}', lambda m: '{% ' + m.group(1).strip() + ' %}', system_prompt)
        prompt_templates['system_prompt']['text'] = system_prompt
        logger.debug(f"Pre-processed system prompt: {system_prompt[:200]}...")

    # Log the final templates before creating the agent
    logger.debug("Final templates before creating agent:")
    for key, value in prompt_templates.items():
        if isinstance(value, dict):
            logger.debug(f"{key}:")
            for subkey, subvalue in value.items():
                logger.debug(f"  {subkey}: {str(subvalue)[:200]}...")
        else:
            logger.debug(f"{key}: {str(value)[:200]}...")

    # Create a copy of the templates for the agent
    agent_templates = {}
    for key, value in prompt_templates.items():
        if isinstance(value, dict):
            agent_templates[key] = {}
            for subkey, subvalue in value.items():
                if isinstance(subvalue, str):
                    # Ensure proper spacing in template expressions
                    content = subvalue
                    # Add spaces around template expressions if missing
                    content = re.sub(r'{{([^{}]+)}}', r'{{ \1 }}', content)
                    content = re.sub(r'{%([^{}]+)%}', r'{% \1 %}', content)
                    # Wrap non-template content in a simple template expression
                    if '{{' not in content and '{%' not in content:
                        escaped_content = content.replace('"', '\\"')
                        content = f'{{{{ "{escaped_content}" }}}}'
                    agent_templates[key][subkey] = content
        elif isinstance(value, str):
            # Ensure proper spacing in template expressions
            content = value
            # Add spaces around template expressions if missing
            content = re.sub(r'{{([^{}]+)}}', r'{{ \1 }}', content)
            content = re.sub(r'{%([^{}]+)%}', r'{% \1 %}', content)
            # Wrap non-template content in a simple template expression
            if '{{' not in content and '{%' not in content:
                escaped_content = content.replace('"', '\\"')
                content = f'{{{{ "{escaped_content}" }}}}'
            agent_templates[key] = content
        else:
            agent_templates[key] = value

    logger.debug("Templates for agent:")
    for key, value in agent_templates.items():
        if isinstance(value, dict):
            logger.debug(f"{key}:")
            for subkey, subvalue in value.items():
                logger.debug(f"  {subkey}: {str(subvalue)[:200]}...")
        else:
            logger.debug(f"{key}: {str(value)[:200]}...")
    
agent = CodeAgent(
    model=model,
    tools=[final_answer, DuckDuckGoSearchTool(), calculate_min_price, extract_price_from_snippet, get_current_time_in_timezone],
    max_steps=15,  # Increased max steps for more complex reasoning
    verbosity_level=2,  # Increased verbosity for better debugging
    grammar=None,
    planning_interval=1,  # Added planning interval to ensure proper planning
    name="question_answering_agent",
    description="An agent specialized in answering various types of questions using available tools. The agent must use the final_answer tool to submit its answer.",
    prompt_templates=agent_templates
)

# Configure Gradio UI with sharing enabled
GradioUI(agent).launch(share=True)