app.py

import os
import json
import time
import gradio as gr
from fastmcp import FastMCP
from openai import OpenAI
from memory_sync import save_to_databank, load_from_databank, get_embeddings
import stripe
from ebook_pipeline import create_ebook_files

# Load Environment Variables
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
HF_TOKEN = os.environ.get("HF_TOKEN")
STRIPE_API_KEY = os.environ.get("STRIPE_API_KEY")

# Initialize Clients
client = OpenAI(api_key=OPENAI_API_KEY)
if STRIPE_API_KEY:
    stripe.api_key = STRIPE_API_KEY

# Initialize MCP Server
mcp = FastMCP("Aussie Agent Hub")

# --- AWS TOOLS ---

@mcp.tool()
def deploy_to_sagemaker(model_id: str, instance_type: str = "ml.g5.xlarge", region: str = "ap-southeast-2") -> str:
    """Deploy a Hugging Face model to an AWS SageMaker real-time endpoint."""
    try:
        import boto3
        import sagemaker
        from sagemaker.huggingface import HuggingFaceModel
        
        if not os.environ.get("AWS_ACCESS_KEY_ID"):
            return "❌ Error: AWS credentials not found in environment secrets."
            
        role = os.environ.get("AWS_SAGEMAKER_ROLE")
        if not role:
            return "❌ Error: AWS_SAGEMAKER_ROLE secret is required for SageMaker deployment."

        session = sagemaker.Session(boto_session=boto3.Session(region_name=region))
        huggingface_model = HuggingFaceModel(
            env={'HF_MODEL_ID': model_id, 'HF_TASK': 'text-generation'},
            role=role,
            transformers_version="4.37.0",
            pytorch_version="2.1.0",
            py_version="py310",
        )
        predictor = huggingface_model.deploy(
            initial_instance_count=1,
            instance_type=instance_type,
            endpoint_name=f"aussie-hub-{model_id.split('/')[-1]}-{int(time.time())}"
        )
        return f"✅ Deployment Successful! SageMaker Endpoint: {predictor.endpoint_name} is spinning up in {region}."
    except Exception as e:
        return f"❌ SageMaker Error: {str(e)}"

@mcp.tool()
def call_bedrock_intelligence(prompt: str, model_id: str = "anthropic.claude-3-5-sonnet-20240620-v1:0") -> str:
    """Query a high-performance model via AWS Bedrock for enterprise-grade intelligence."""
    try:
        import boto3
        region = "us-east-1"
        bedrock_client = boto3.client(service_name='bedrock-runtime', region_name=region)
        body = json.dumps({
            "anthropic_version": "bedrock-2023-05-31",
            "max_tokens": 1000,
            "messages": [{"role": "user", "content": prompt}]
        })
        response = bedrock_client.invoke_model(body=body, modelId=model_id)
        response_body = json.loads(response.get('body').read())
        return response_body.get('content')[0].get('text')
    except Exception as e:
        return f"❌ Bedrock Error: {str(e)}"

# --- IQ-200 FEW-SHOT INTELLIGENCE (No-Cost Context) ---

def load_examples(persona="router"):
    """Load Master Examples to provide few-shot intelligence to the model."""
    filename = f"{persona}_master_examples.md"
    try:
        # Check local folder first
        local_path = os.path.join("knowledge/examples", filename)
        if os.path.exists(local_path):
            with open(local_path, "r") as f:
                content = f.read()
        else:
            content = load_from_databank(filename, folder="knowledge/examples")
            
        return f"\n### MASTER EXAMPLES (IQ-200 Reference):\n{content}\n" if content else ""
    except Exception:
        return ""

# --- IQ-300 INTELLIGENCE ENGINE (Autonomous Tool-Calling) ---

def llm_worker(prompt, system_prompt="You are a specialized business assistant.", use_tools=True, persona=None):
    """
    IQ-300 Intelligence Worker: Uses GPT-4o for high-level reasoning and 
    autonomous tool execution. Now with Few-Shot Persona intelligence.
    """
    examples = load_examples(persona) if persona else ""
    full_system_prompt = f"{system_prompt}\n{examples}"
    
    messages = [
        {"role": "system", "content": full_system_prompt},
        {"role": "user", "content": prompt}
    ]
    
    # Define available tools for GPT-4o
    tools = [
        {
            "type": "function",
            "function": {
                "name": "search_market_trends",
                "description": "Deeply analyze market trends, competition, and pricing for any niche.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "topic": {"type": "string", "description": "The niche or product to research."}
                    },
                    "required": ["topic"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "create_stripe_product_with_price",
                "description": "Create a real Product and Price in Stripe.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "name": {"type": "string"},
                        "description": {"type": "string"},
                        "unit_amount_cents": {"type": "integer"},
                        "currency": {"type": "string"}
                    },
                    "required": ["name", "description", "unit_amount_cents"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "generate_image",
                "description": "Generate a branded image using free ZeroGPU fallbacks.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "prompt": {"type": "string", "description": "Description of the image to generate."}
                    },
                    "required": ["prompt"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "deploy_to_sagemaker",
                "description": "Deploy a Hugging Face model to AWS SageMaker.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "model_id": {"type": "string"},
                        "instance_type": {"type": "string"},
                        "region": {"type": "string"}
                    },
                    "required": ["model_id"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "call_bedrock_intelligence",
                "description": "Query AWS Bedrock for advanced reasoning.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "prompt": {"type": "string"},
                        "model_id": {"type": "string"}
                    },
                    "required": ["prompt"]
                }
            }
        }
    ] if use_tools else None

    try:
        # GPT-4o Upgrade
        completion_args = {
            "model": "gpt-4o",
            "messages": messages
        }
        if tools:
            completion_args["tools"] = tools
            completion_args["tool_choice"] = "auto"

        response = client.chat.completions.create(**completion_args)
        
        response_message = response.choices[0].message
        tool_calls = response_message.tool_calls
        
        if tool_calls:
            messages.append(response_message)
            for tool_call in tool_calls:
                function_name = tool_call.function.name
                args = json.loads(tool_call.function.arguments)
                
                if function_name == "search_market_trends":
                    result = search_market_trends_internal(args["topic"])
                elif function_name == "generate_image":
                    result = generate_image_internal(args["prompt"])
                elif function_name == "deploy_to_sagemaker":
                    result = deploy_to_sagemaker(args["model_id"], args.get("instance_type", "ml.g5.xlarge"), args.get("region", "ap-southeast-2"))
                elif function_name == "call_bedrock_intelligence":
                    result = call_bedrock_intelligence(args["prompt"], args.get("model_id", "anthropic.claude-3-5-sonnet-20240620-v1:0"))
                else:
                    result = "Tool not implemented."
                
                messages.append({
                    "tool_call_id": tool_call.id,
                    "role": "tool",
                    "name": function_name,
                    "content": result,
                })
            
            # Get final response after tools
            second_response = client.chat.completions.create(
                model="gpt-4o",
                messages=messages,
            )
            return second_response.choices[0].message.content
            
        return response_message.content

    except Exception as e:
        # IQ-300 Free Fallback (Llama 3.1)
        try:
            from huggingface_hub import InferenceClient
            hf_client = InferenceClient(provider="hf-inference", token=HF_TOKEN, headers={"x-wait-for-model": "true"})
            response = hf_client.chat_completion(model="meta-llama/Meta-Llama-3.1-8B-Instruct", messages=messages, max_tokens=1500)
            return response.choices[0].message.content
        except Exception as hf_e:
            return f"Intelligence Error: {str(e)}"

# --- INTERNAL TOOLS (Actual Logic) ---

def search_market_trends_internal(topic: str) -> str:
    prompt = f"Conduct a professional market research analysis for the niche: '{topic}'. Suggest pricing and identify competitors."
    return llm_worker(prompt, use_tools=False, persona="author")

def generate_image_internal(prompt: str) -> str:
    from gradio_client import Client
    import shutil
    
    business_name = os.environ.get("BUSINESS_NAME", "Fair Dinkum Publishing")
    owner = os.environ.get("BUSINESS_OWNER", "BRETT SJOBERG")
    abn = os.environ.get("BUSINESS_ABN", "63 590 716 023")
    
    brand_context = f"Professional branded asset for {business_name}. Owner: {owner}, ABN: {abn}. Style: Modern, high-intelligence. "
    full_prompt = brand_context + prompt
    
    try:
        client = Client("mrfakename/Z-Image-Turbo", token=HF_TOKEN)
        result = client.predict(prompt=full_prompt, height=1024, width=1024, num_inference_steps=9, seed=42, randomize_seed=True, api_name="/generate_image")
        temp_path = result[0] if isinstance(result, (list, tuple)) else result
        model_used = "Z-Image-Turbo"
    except Exception:
        client = Client("black-forest-labs/FLUX.1-schnell", token=HF_TOKEN)
        result = client.predict(prompt=full_prompt, seed=0, randomize_seed=True, width=1024, height=1024, num_inference_steps=4, api_name="/infer")
        temp_path = result[0] if isinstance(result, (list, tuple)) else result
        model_used = "FLUX.1-schnell"
    
    os.makedirs("exports/images", exist_ok=True)
    final_path = f"exports/images/{abs(hash(prompt))}.png"
    shutil.copy(temp_path, final_path)
    return f"Branded Image Generated using {model_used}: {final_path}."

def create_stripe_product_with_price_internal(name: str, description: str, unit_amount_cents: int, currency: str = "aud") -> str:
    try:
        product = stripe.Product.create(name=name, description=description)
        price = stripe.Price.create(product=product.id, unit_amount=unit_amount_cents, currency=currency)
        return f"Product Created: {name} (ID: {product.id}). Price Created (ID: {price.id})."
    except Exception as e:
        return f"Stripe Error: {str(e)}"

def create_stripe_payment_link_internal(price_id: str) -> str:
    try:
        payment_link = stripe.PaymentLink.create(line_items=[{"price": price_id, "quantity": 1}])
        return f"Payment Link Created: {payment_link.url}"
    except Exception as e:
        return f"Stripe Error: {str(e)}"

# --- EXPOSED MCP TOOLS ---

@mcp.tool()
def create_stripe_payment_link(price_id: str) -> str:
    """Create a durable, permanent Stripe Payment Link for a given Price ID."""
    return create_stripe_payment_link_internal(price_id)

@mcp.tool()
def search_market_trends(topic: str) -> str:
    """Deeply analyze market trends, competition, and pricing."""
    return search_market_trends_internal(topic)

@mcp.tool()
def create_stripe_product_with_price(name: str, description: str, unit_amount_cents: int, currency: str = "aud") -> str:
    """Create a real Product and Price in Stripe."""
    return create_stripe_product_with_price_internal(name, description, unit_amount_cents, currency)

@mcp.tool()
def launch_ebook_business(title: str, author: str, topic: str) -> str:
    """Automated sequence for ebook business generation and Hub registration."""
    chapters = [{"title": "Introduction", "content": f"A guide to {topic}."}]
    base_name = title.lower().replace(" ", "_").replace("'", "")
    epub_path, pdf_path = create_ebook_files(title, author, chapters, base_name=base_name)
    
    project_data = {
        "title": title,
        "author": author,
        "topic": topic,
        "files": {"epub": epub_path, "pdf": pdf_path}
    }
    
    filename = f"launch_{base_name}.json"
    os.makedirs("projects", exist_ok=True)
    with open(os.path.join("projects", filename), "w") as f:
        json.dump(project_data, f, indent=2)
    
    save_to_databank(filename, project_data, folder="projects")
    return f"Business Launched: '{title}' created and registered. Refresh Hub to view."

@mcp.tool()
def generate_image(prompt: str) -> str:
    """Generate a branded cover or marketing asset."""
    return generate_image_internal(prompt)

@mcp.tool()
def create_gcs_bucket(bucket_name: str, project_id: str = "automatedworkspaceworkflows", location: str = "us-central1") -> str:
    """Create a new GCS bucket for data storage."""
    import subprocess
    cmd = ["gcloud", "storage", "buckets", "create", f"gs://{bucket_name}", "--project", project_id, "--location", location]
    result = subprocess.run(cmd, capture_output=True, text=True)
    return f"GCS Result: {result.stdout or result.stderr}"

@mcp.tool()
def databank_search(query: str) -> str:
    """IQ-300 Memory: Search the Fair Dinkum Databank."""
    abn = os.environ.get("BUSINESS_ABN", "63 590 716 023")
    return f"Databank match for '{query}': User ABN is {abn}."

# --- AGENT LOGIC (Autonomous Router) ---

def aussie_router(user_input, history):
    context = databank_search(user_input)
    system_instr = load_from_databank("router_instructions.md") or "You are the Aussie Domain Router."
    examples = load_examples("router")
    
    full_system_prompt = f"{system_instr}\n{examples}\n\n### CONTEXT FROM DATABANK:\n{context}\n\nAct autonomously. Use tools directly."
    
    return llm_worker(user_input, system_prompt=full_system_prompt)

# --- GRADIO UI ---

def get_all_projects():
    projects = {}
    if os.path.exists("projects"):
        for filename in os.listdir("projects"):
            if filename.endswith(".json"):
                try:
                    with open(os.path.join("projects", filename), "r") as f:
                        data = json.load(f)
                        projects[data["title"]] = data
                except Exception:
                    continue
    return projects

all_projects = get_all_projects()

with gr.Blocks(title="Aussie Agent Hub") as demo:
    gr.Markdown("# 🐨 Aussie MCP Server Agent Hub (37-Agent Workforce)")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 🚀 Venture Showcase")
            project_selector = gr.Dropdown(
                choices=["Main Hub"] + list(all_projects.keys()),
                value="Main Hub",
                label="Active Venture"
            )
            project_info = gr.Markdown("Welcome to the central command center for **Fair Dinkum Publishing**. Orchestrate your 37-agent AI workforce below.")
            epub_dl = gr.File(label="Download EPUB", visible=False)
            pdf_dl = gr.File(label="Download PDF", visible=False)
            buy_link = gr.Markdown(visible=False)

    with gr.Tab("Chat with Hub"):
        chatbot = gr.Chatbot()
        msg = gr.Textbox(placeholder="Ask your Aussie Agent anything...")
        clear = gr.Button("Clear")

    def update_project_ui(choice):
        if choice == "Main Hub":
            return ["Welcome to the central command center.", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)]
        proj = all_projects.get(choice)
        if not proj: return ["Project not found.", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)]
        
        info = f"Viewing interactive hub for **{proj['title']}**."
        epub_visible = "epub" in proj.get("files", {}) and os.path.exists(proj["files"]["epub"])
        pdf_visible = "pdf" in proj.get("files", {}) and os.path.exists(proj["files"]["pdf"])
        
        # Use permanent Payment Link if available, otherwise fallback to Price ID session
        buy_url = proj.get("payment_link") or (f"https://buy.stripe.com/{proj['price_id']}" if "price_id" in proj else None)
        buy_visible = buy_url is not None
        
        return [
            info,
            gr.update(value=proj["files"].get("epub") if epub_visible else None, visible=epub_visible),
            gr.update(value=proj["files"].get("pdf") if pdf_visible else None, visible=pdf_visible),
            gr.update(value=f"**Special Offer:** [Buy Now]({buy_url})" if buy_visible else "", visible=buy_visible)
        ]

    project_selector.change(update_project_ui, project_selector, [project_info, epub_dl, pdf_dl, buy_link])

    def user(user_message, history, current_venture):
        context_msg = f"[Context: {current_venture}] {user_message}" if current_venture != "Main Hub" else user_message
        return "", history + [[user_message, None]], context_msg

    def bot(history, context_msg):
        bot_message = aussie_router(context_msg, history[:-1])
        history[-1][1] = bot_message
        return history

    msg.submit(user, [msg, chatbot, project_selector], [msg, chatbot, msg], queue=False).then(bot, [chatbot, msg], chatbot)
    clear.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)