src/agents/CleanReponseAgent.py

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage, HumanMessage

from src.core.settings import get_settings

load_dotenv()
_settings = get_settings()
client = ChatOpenAI(model=_settings.models.formatter_model)


def ai_format_response(raw_text: str):
    system = f"""
You are a UI formatting assistant for a Streamlit chat application.

STRICT RULES:

1. FORMAT:
- Use HTML tags ONLY (no Markdown headings)
- Use:
  <h2> for main title
  <h3> for sections
  <ul><li> for bullet points
  <br> for spacing

2. STRUCTURE:
- Start with:
  <h2>Title</h2>

- For each source:
  <h3>Source Name</h3>
  <a href="URL" target="_blank">Visit Source</a>

  <ul>
    <li>Point 1</li>
    <li>Point 2</li>
  </ul>

3. FOLLOW-UP QUESTIONS:
<h3>Follow-up Questions</h3>
<ul>
  <li>Question 1</li>
</ul>

4. OUTPUT:
- Must render correctly in Streamlit
- No Markdown
- No explanations
- Return only formatted HTML

5. ADDITIONAL REQUIREMENTS:
- Group by source
- Show title + clickable URL
- Provide 3-5 bullet point summary per source
- Keep it concise and readable
- Remove noise like "Read now", "Login", etc.
- Add headings
- Add section 'Followup Questions' and provide 3-5 questions related to the content
- Remove ```html or any kind of code block
Content:
{raw_text}
"""
    response = client.invoke([
        SystemMessage(content=system),
        HumanMessage(content=raw_text)
    ])

    return response.content