Update main.py

main.py

@@ -1,4 +1,10 @@
+import re
 import os
+from enum import Enum
+from uuid import uuid4
+import base64
+import requests
+from io import BytesIO
 import time
 import asyncio
 import logging
@@ -7,9 +13,9 @@ import tiktoken
 from uuid import uuid4
 from functools import lru_cache
 from typing import Optional, List, Dict, Literal
-from fastapi import FastAPI, HTTPException, Depends, Security, BackgroundTasks, Query
+from fastapi import FastAPI, HTTPException, Depends, Security, BackgroundTasks, Query, Header
 from fastapi.security import APIKeyHeader
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel, Field
 from openai import OpenAI
 # ============================================================================
@@ -32,6 +38,7 @@ app = FastAPI()
 
 # API key configuration
 API_KEY_NAME = "X-API-Key"
+PEXELS_API_KEY = os.environ["PEXELS_API_KEY"]
 API_KEY = os.environ.get("CHAT_AUTH_KEY", "default_secret_key")
 api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
 
@@ -69,6 +76,7 @@ class LLMAgentQueryModel(BaseModel):
             }
         }
 
+
 # API key and client setup
 @lru_cache()
 def get_api_keys():
@@ -264,6 +272,136 @@ async def text_to_speech(
 
     return StreamingResponse(generate(), media_type="audio/mpeg")
 
+
+
+# ============================================================================
+# PPT AGENT
+# ============================================================================
+
+class PresentationChatModel(BaseModel):
+    prompt: str
+    conversation_id: Optional[str] = None
+    model_id: str
+    user_id: str
+    
+# Enum for output formats
+class OutputFormatEnum(str, Enum):
+    html = "html"
+    pdf = "pdf"
+    pptx = "pptx"
+
+# Class model for presentation data
+class PresentationModel(BaseModel):
+    markdown: str
+    output_format: OutputFormatEnum = OutputFormatEnum.html
+    
+def get_pexels_image(query):
+    default_img_url = "https://images.pexels.com/photos/593158/pexels-photo-593158.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=2"
+    url = f"https://api.pexels.com/v1/search?query={query}&per_page=1"
+    headers = {"Authorization": PEXELS_API_KEY}
+    try:
+        response = requests.get(url, headers=headers)
+        if response.status_code == 200:
+            data = response.json()
+            if data["total_results"] > 0:
+                return data["photos"][0]["src"]["medium"]
+        return default_img_url
+    except:
+        return default_img_url
+
+def replace_image_keywords(text):
+    def replace_match(match):
+        bg_params = match.group(1)
+        keyword = re.sub(r'[^\w\s]', ' ', match.group(2)).strip()
+        image_url = get_pexels_image(keyword)
+        return f"![bg {bg_params}]({image_url})"
+
+    pattern = r'!\[bg (.*?)\]\((.*?)\)'
+    return re.sub(pattern, replace_match, text)
+
+def convert_markdown_marp(markdown, output_format='html'):
+    API_URL = "https://pvanand-marpit-backend.hf.space/convert"
+    if output_format not in ['html', 'pdf', 'pptx']:
+        raise ValueError(f"Invalid output format. Supported formats are: html, pdf, pptx")
+
+    data = {
+        "markdown": markdown,
+        "outputFormat": output_format,
+        "options": []
+    }
+
+    try:
+        response = requests.post(API_URL, json=data, timeout=30)
+        response.raise_for_status()
+        return response.content
+    except requests.exceptions.RequestException as e:
+        logger.error(f"An error occurred while connecting to the API: {e}")
+        return None
+
+
+@app.post("/convert-md-to-presentation")
+async def create_presentation(data: PresentationModel):
+    if not data.markdown:
+        raise HTTPException(status_code=400, detail="Please provide Markdown text.")
+
+    markdown_with_images = replace_image_keywords(extract_data_from_tag(markdown, "marp_presentation"))
+    result = convert_markdown_marp(markdown_with_images, output_format)
+
+    if result:
+        if output_format == 'html':
+            return {"html": result.decode()}
+        elif output_format == 'pdf':
+            return StreamingResponse(BytesIO(result), media_type="application/pdf")
+        elif output_format == 'pptx':
+            return StreamingResponse(BytesIO(result), media_type="application/vnd.openxmlformats-officedocument.presentationml.presentation")
+    else:
+        raise HTTPException(status_code=500, detail="Failed to create presentation.")
+
+
+@app.post("/presentation-agent")
+async def presentation_chat(query: PresentationChatModel, background_tasks: BackgroundTasks, api_key: str = Depends(verify_api_key)):
+    """
+    Presentation chat endpoint that generates a presentation based on user queries.
+    Uses the llm_agent function and returns both markdown and HTML output.
+    Requires API Key authentication via X-API-Key header.
+    """
+    logger.info(f"Received presentation chat query: {query.prompt}")
+
+    # Create a new LLMAgentQueryModel with a specific system message for presentation generation
+    llm_query = LLMAgentQueryModel(
+        prompt=query.prompt,
+        conversation_id=query.conversation_id,
+        system_message="You are a presentation creation assistant. Generate a Marp-formatted presentation based on the user's request. Use the Marp format provided earlier.",
+        model_id=query.model_id,
+        user_id=query.user_id
+    )
+
+    # Use the llm_agent function to generate the presentation content
+    response_stream = await llm_agent(llm_query, background_tasks, api_key)
+
+    # Collect the entire response
+    full_response = ""
+    async for chunk in response_stream.body_iterator:
+        full_response += chunk.decode()
+
+    # Extract the Marp presentation content
+    marp_content = extract_data_from_tag(full_response, "marp_presentation")
+
+    # Replace image keywords
+    marp_content_with_images = replace_image_keywords(marp_content)
+
+    # Convert Markdown to HTML
+    html_content = convert_markdown_marp(marp_content_with_images, 'html')
+
+    if html_content:
+        return JSONResponse({
+            "markdown": marp_content_with_images,
+            "html": html_content.decode()
+        })
+    else:
+        raise HTTPException(status_code=500, detail="Failed to create presentation.")
+
+
 # ============================================================================
 # Main Execution
 # ============================================================================