app.py

# -----------------------------------------------------------------------------
# Copyright (c) 2025 Vikash Kumar
#
# This source code is the property of Vikash Kumar and may not be reproduced,
# distributed, or used without explicit permission.
#
# Author: Vikash Kumar
# GitHub: https://github.com/vikukumar
# LinkedIn: https://linkedin.com/in/vikash-edude
# -----------------------------------------------------------------------------

from dotenv import load_dotenv
from fastapi.staticfiles import StaticFiles
load_dotenv()
from os import getenv

OPENROUTER_API_KEY = getenv('OPENROUTER_API_KEY',None)
HF_API_KEY = getenv('HF_API_KEY',None)
STABILITY_API_KEY = getenv('STABILITY_API_KEY',None)

if OPENROUTER_API_KEY is None and HF_API_KEY is None and STABILITY_API_KEY is None:
    raise Exception("API Key Missing please Add into .env file")

from pydantic import BaseModel, HttpUrl, Base64Str, StrictStr,Field, StrictBool
from typing import Any

class Generate(BaseModel):
    rid:StrictStr=Field(description="Request Id to track error or progress")
    url:HttpUrl=Field(default= "https://vikshro.in",description="Website Url including http/https so that AI Agent Can Understand the Business")
    ratio:StrictStr=Field(default="4:5",description="Image Aspect Ratio default is 4:5")
    product:StrictBool=Field(default=False,description="Web Url is Product Url or not?")


class GenerateResponse(BaseModel):
    prompt:StrictStr = Field(description="Image Generation Prompt")
    image:Any = Field(default='',description="Image base64 encoded format")
    message:StrictStr = Field(None, description="Error If occur in Image Generation")
    content:str= Field(None,description="Social Media Contents")
    product:str=Field(None,description="Real Product Image")


class ErrorResponse(BaseModel):
    status_code:int = Field(description="Error Status Code")
    message:StrictStr  = Field(description="Error Message")


from fastapi import FastAPI, status, Request, HTTPException
from fastapi.exceptions import ValidationException, RequestValidationError, ResponseValidationError, StarletteHTTPException
from fastapi.responses import FileResponse, JSONResponse, Response

Errors= {
    status.HTTP_404_NOT_FOUND:{"model":ErrorResponse,"description":"Not Found Error"},
    status.HTTP_400_BAD_REQUEST:{"model":ErrorResponse,"description":"Bad Request"},
    status.HTTP_500_INTERNAL_SERVER_ERROR:{"model":ErrorResponse,"description":"Internal Server Error"},
    status.HTTP_422_UNPROCESSABLE_ENTITY:{"model":ErrorResponse,"description":"Unprocessable Entity"}
}


APP = FastAPI(title="🤖 SMIGAI",
              version="1.0.0",
              summary="🤖 SMIGAI: AI Agent for Social Media Content & Image Generation",
              description="""This AI Agent generates social media content and visuals for a brand or company using their URL or domain.
              \n* Generate engaging social media posts tailored to a company's brand or website
              \n* Automatically create AI-generated images relevant to the brand
              \n* Simple web interface to test and run content/image generation
              \n* Plug-and-play with OpenRouter, Hugging Face, or Stability AI APIs""",
              servers=[
                  {'url':'http://localhost:5000',
                   'description':'Localhost Development Server'
                  },
                  {'url':'https://smigai.vikshro.in',
                   'description':'Production Server'
                  },
                ],
              terms_of_service='https://vikshro.in/Terms',
              contact={
                  "name": "Vikash Kumar",
                  "url": "https://github.com/vikukumar",
                  "email": "vikashkumar0037@gmail.com",
              },
              openapi_url="/opa",
              license_info={
                  "name": "VIKSHRO 1.0",
                  "url": "https://vikshro.in/Privacy",
                },
              docs_url=None, redoc_url=None,
              responses=Errors)
APP.openapi_version = "3.1.1"

@APP.exception_handler(StarletteHTTPException)
@APP.exception_handler(HTTPException)
async def default_error(req, err):
    return JSONResponse(content=ErrorResponse(status_code=err.status_code,message=str(err)).model_dump(),status_code=err.status_code)

@APP.exception_handler(Exception)
async def exceptions(req, err:Exception):
    return JSONResponse(content=ErrorResponse(status_code=500,message=str(err)).model_dump(),status_code=500)

@APP.exception_handler(ValidationException)
@APP.exception_handler(RequestValidationError)
@APP.exception_handler(ResponseValidationError)
async def default_val_error(req,err):
    return JSONResponse(content=ErrorResponse(status_code=422,message=err._errors[0]['msg']).model_dump(),status_code=422)

from lib import Agent

SESSIONS={}

from fastapi.openapi.docs import (
    get_redoc_html,
    get_swagger_ui_html,
    get_swagger_ui_oauth2_redirect_html,
)

@APP.get("/docs", include_in_schema=False)
async def custom_swagger_ui_html(req:Request):
    return get_swagger_ui_html(
        openapi_url=APP.openapi_url,
        title=APP.title + " - Swagger UI",
        oauth2_redirect_url=APP.swagger_ui_oauth2_redirect_url,
        swagger_js_url="https://unpkg.com/swagger-ui-dist@5/swagger-ui-bundle.js",
        swagger_css_url="https://unpkg.com/swagger-ui-dist@5/swagger-ui.css",
        swagger_favicon_url="/favicon.png",
    )

@APP.get(APP.swagger_ui_oauth2_redirect_url, include_in_schema=False)
async def swagger_ui_redirect():
    return get_swagger_ui_oauth2_redirect_html()


@APP.get("/redoc", include_in_schema=False)
async def redoc_html():
    return get_redoc_html(
        openapi_url=APP.openapi_url,
        title=APP.title + " - ReDoc",
        redoc_js_url="https://unpkg.com/redoc@2/bundles/redoc.standalone.js",
        redoc_favicon_url="/favicon.png",
    )


""" Main App Functionalities """

@APP.post('/api/v1/generate',response_model=GenerateResponse,tags=['SMIGAI'],name='Generate',summary='Generate Image based on Company Website',operation_id='post_v1_generate',status_code=200)
async def Post_V1_Generate(req:Request,gen:Generate):
    agent = Agent(gen.url,referer=req.headers.get('referer'),title='SMIGAI',product=gen.product)
    SESSIONS[gen.rid] = agent
    try:
        image = agent.generate(ratio=gen.ratio)
        return GenerateResponse(prompt=agent.img_prompt,image=image,message='Image generated Successfully',content=agent.content,product=agent.prodbas64).model_dump()
    except Exception as e:
        print(e)
        return JSONResponse(content=GenerateResponse(prompt=agent.img_prompt,image='',message=str(e),content='',product=agent.prodbas64).model_dump(),status_code=400)
    

APP.mount("/", StaticFiles(directory="public", html=True), name="public")

# -----------------------------------------------------------------------------
# Copyright (c) 2025 Vikash Kumar
#
# This source code is the property of Vikash Kumar and may not be reproduced,
# distributed, or used without explicit permission.
#
# Author: Vikash Kumar
# GitHub: https://github.com/vikukumar
# LinkedIn: https://linkedin.com/in/vikash-edude
# -----------------------------------------------------------------------------