Update app.py

app.py

@@ -4,7 +4,7 @@ from fastapi.responses import HTMLResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 from huggingface_hub import InferenceClient
-import uvicorn
+import re
 
 # Initialize FastAPI app
 app = FastAPI()
@@ -22,6 +22,25 @@ class InfographicRequest(BaseModel):
 # Load prompt template from environment variable
 PROMPT_TEMPLATE = os.getenv("PROMPT_TEMPLATE")
 
+
+async def extract_code_blocks(markdown_text):
+    """
+    Extracts code blocks from the given Markdown text.
+
+    Args:
+        markdown_text (str): The Markdown content as a string.
+
+    Returns:
+        list: A list of code blocks extracted from the Markdown.
+    """
+    # Regex to match code blocks (fenced with triple backticks)
+    code_block_pattern = re.compile(r'```.*?\n(.*?)```', re.DOTALL)
+
+    # Find all code blocks
+    code_blocks = code_block_pattern.findall(markdown_text)
+    
+    return code_blocks
+
 # Route to serve the HTML template
 @app.get("/", response_class=HTMLResponse)
 async def serve_frontend():
@@ -35,7 +54,6 @@ async def generate_infographic(request: InfographicRequest):
     prompt = PROMPT_TEMPLATE.format(description=description)
 
     try:
-        # Query Hugging Face model
         messages = [{"role": "user", "content": prompt}]
         stream = client.chat.completions.create(
             model="Qwen/Qwen2.5-Coder-32B-Instruct",
@@ -46,13 +64,16 @@ async def generate_infographic(request: InfographicRequest):
             stream=True,
         )
 
-        # Collect the HTML content from the stream
-        generated_html = ""
+        
+        generated_text = ""
         for chunk in stream:
-            generated_html += chunk.choices[0].delta.content
-
-        # Return the generated HTML content
-        return JSONResponse(content={"html": generated_html})
+            generated_text += chunk.choices[0].delta.content
+        
+        code_blocks=extract_code_blocks(generated_text)
+        if code_blocks:
+            return JSONResponse(content={"html": code_blocks[0]})
+        else:
+            return JSONResponse(content={"error": "No generation"},status_code=500)
 
     except Exception as e:
         return JSONResponse(content={"error": str(e)}, status_code=500)