agent.py

from smolagents import CodeAgent, Tool
import os
import time
import requests
import re
from typing import Dict, List, Optional, Any
from smolagents.models import Model
from smolagents.default_tools import DuckDuckGoSearchTool, FinalAnswerTool
from dotenv import load_dotenv
import os

from myTools.ExtractWikipediaSection import ExtractWikipediaSection
from myTools.ExtractWebContentWithSelenium import ExtractWebContentWithSelenium
from myTools.GetPlaceholderImageTool import GetPlaceholderImageTool
from myTools.GetSVG import GetSVG
from myTools.GetSVGList import GetSVGList
from myTools.GetLogo import GetLogo

load_dotenv()
api_key = os.getenv("MistralApiKey")



class ChatMessage:
    def __init__(self, role: str, content: str):
        self.role = role
        self.content = content
        self.raw = None

    @classmethod
    def from_dict(cls, message_dict: Dict[str, str]) -> 'ChatMessage':
        return cls(role=message_dict['role'], content=message_dict['content'])

    def __repr__(self):
        return f"ChatMessage(role={self.role}, content={self.content})"



class MistralClient:
    def __init__(self, api_key: str, api_base: str = "https://api.mistral.ai"):
        self.api_key = api_key
        self.api_base = api_base

    def generate(self, model_id: str, messages: List[Dict[str, str]], **kwargs):
        url = f"{self.api_base}/v1/chat/completions"
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        data = {
            "model": model_id,
            "messages": messages,
            "temperature": 0.7,
            "max_tokens": 4096,
            **kwargs
        }
        retries = 5
        backoff = 2.0
        for attempt in range(1, retries + 1):
            response = requests.post(url, headers=headers, json=data)

            if response.status_code == 429:
                print(f"[429] Too many requests - attempt {attempt}/{retries}, please wait {backoff}s...")
                time.sleep(backoff)
                backoff *= 2
                continue

            try:
                response.raise_for_status()
                return response.json()
            except requests.exceptions.HTTPError as e:
                print(f"[ERREUR] HTTP {response.status_code}: {e}")
                raise e

        raise RuntimeError(f"Failed {retries} attempts (429 of errors)")



class MistralApiModel(Model):
    """A class to interact with Mistral's API for language model interaction."""

    def __init__(
        self,
        model_id: str,
        api_key: Optional[str] = None,
        api_base: Optional[str] = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.model_id = model_id
        self.api_key = api_key
        self.api_base = api_base or "https://api.mistral.ai"
        self.client = MistralClient(api_key=self.api_key, api_base=self.api_base)

    def generate(
        self,
        messages: List[Dict[str, str]],
        **kwargs,
    ) -> ChatMessage:
        return self.__call__(messages=messages, **kwargs)

    def __call__(
        self,
        messages: List[Dict[str, str]],
        stop_sequences: Optional[List[str]] = None,
        **kwargs,
    ) -> ChatMessage:
        completion_kwargs = self._prepare_completion_kwargs(
            messages=messages,
            stop_sequences=stop_sequences,
            **kwargs,
        )

        response = self.client.generate(model_id=self.model_id, **completion_kwargs)
        

        message = ChatMessage.from_dict(response['choices'][0]['message'])
        message.raw = response
        usage = response.get('usage', {})
        message.token_usage = type('TokenUsage', (), {
            "input_tokens": usage.get("prompt_tokens", 0),
            "output_tokens": usage.get("completion_tokens", 0),
            "total_tokens": usage.get("total_tokens", 0)
        })()
        return message

class HektoreAgent:
    def __init__(self):
        print("Agent initialized.")

    def __call__(self, section: str, brief: str, colorOne: str, colorTwo: str, language: str, company_name: str, current_html: str = "", additional_args: dict = {}) -> str:

        print(f"Agent received brief : {brief}...")
        print(f"Additional args received : {additional_args}...")

        all_tools = [
                    GetPlaceholderImageTool(),
                    ExtractWebContentWithSelenium(),
                    DuckDuckGoSearchTool(),
                    GetSVG(),
                    GetSVGList(),
                    GetLogo(),
                    FinalAnswerTool()
                    ]
        marketing_tools = [
            ExtractWebContentWithSelenium(),
            DuckDuckGoSearchTool(),
            FinalAnswerTool()
            ]

        model = MistralApiModel(
            model_id="codestral-2501",
            api_base="https://api.mistral.ai",
            api_key=api_key
        )
        
        marketing_agent = CodeAgent(
            model=model,
            name="MarketingSEOExpert",
            description="Specialist in marketing web and SEO, generate powerful content with impact. Can't generate HTML, juste generate raw text.",
            tools=marketing_tools,
            additional_authorized_imports=[
                "geopandas",
                "plotly",
                "shapely",
                "json",
                "pandas",
                "numpy",
                "time",
                "openpyxl",
                "pdfminer",
                "pdfminer.six",
                "PyPDF2",
                "io",
                "open",
                "librosa",
                "bs4",
                "os",
                "builtins.open",
                "builtins.write",
                "PyGithub",
                "requests"
            ],
            verbosity_level=2,
            #final_answer_checks=[check_reasoning],
            max_steps=15,
)
  
        manager_agent = CodeAgent(
            model=model,
            tools=all_tools,
            #managed_agents=[marketing_agent],
            additional_authorized_imports=[
                "geopandas",
                "plotly",
                "shapely",
                "json",
                "pandas",
                "numpy",
                "time",
                "openpyxl",
                "pdfminer",
                "pdfminer.six",
                "PyPDF2",
                "io",
                "open",
                "librosa",
                "bs4",
                "os",
                "builtins.open",
                "builtins.write",
                "PyGithub",
                "requests"
            ],
            planning_interval=10,
            verbosity_level=2,
            #final_answer_checks=[check_reasoning],
            max_steps=100,
)
        
        

        prompt = f"""
        You are a professional AI website generator.
        Your job is to build a rich, modern, and impressive  HTML + TailwindCSS, using semantic HTML5 and responsive design.
        You will build the website section by section. Only output the HTML markup for the current section.
        The content text must be awesome and very very rich.
        Do not output any <html>, <head>, <body>, or <!DOCTYPE> tags.
        Only write the core section content.
        The name or the company is : {company_name}
        The website is about : {brief}
        Language of the website : {language}
        You are encouraged to use advanced design elements and modern layout ideas.
        We need a lot of text in each section. All text must be SEO compliant. We need Call to action.
        Develop this section :
        {section}
        
        
        
        **Styling and layout**:
        - Use Tailwind utility classes
        - Use Tailwind's `container`, `grid`, `flex`, `gap`, `rounded`, `shadow`, `text-*`, `bg-*` utilities
        - Add responsive breakpoints (`sm:`, `md:`, `lg:`) where needed
        - Add hover/focus/transition effects for buttons and cards
        - Add a fixed header if relevant
        - The CTA text must be impactful and very short
        - Use these 2 colors : Primary : {colorOne} and secondary : {colorTwo}
        - Don't use join in python to build multiple block in html section
        - Use TailwindCSS via CDN:  
            `<script src="https://cdn.tailwindcss.com"></script>`
        **Extras (encouraged but optional)**:
        - Use Alpine.js for interactions (mobile nav, carousel, etc.)
        - Add animation classes (`transition`, `duration-300`, `ease-in-out`)
        - Use SVG separators or background patterns
        - Include meta tags for SEO, accessibility attributes, and semantic labels
        - Use real picture, no lorem ipsum.
        
        You may use Alpine.js for dynamic UI behavior (e.g., toggles, mobile menus, tabs).  
        Load it via CDN:  
        <script src="https://cdn.jsdelivr.net/npm/alpinejs@3.x.x/dist/cdn.min.js" defer></script>
        You may use Swiper.js for testimonials or logo sliders:  
        <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/swiper@11/swiper-bundle.min.css" />  
        <script src="https://cdn.jsdelivr.net/npm/swiper@11/swiper-bundle.min.js"></script>
        You may use animate.css (Animatation) for smooth animations :  
        <link
            rel="stylesheet"
            href="https://cdnjs.cloudflare.com/ajax/libs/animate.css/4.1.1/animate.min.css"
        />
        You may use Micromodal.js (Modal) for simple modal :  
        <script src="https://unpkg.com/micromodal/dist/micromodal.min.js"></script>
        The page should look ready to present to enterprise clients. Do not oversimplify. No `Lorem ipsum`. Write realistic, sharp content. Only output the HTML for this section. Do not include <html>, <head>, or <body>. Do not include DOCTYPE.
            """

        result = manager_agent.run(task=prompt, additional_args=additional_args)
        return result