nodes.py

"""
nodes.py — LangGraph 各節點實作
每個節點都是純函數：(state) → dict（partial state update）
"""
import os
import time
import uuid
import json
import httpx
import asyncio
from typing import Optional

# 服務 URL（由環境變數設定，在 HuggingFace Spaces 間通信）
CREW_SERVICE_URL = os.getenv("CREW_SERVICE_URL", "http://localhost:7861")

# 分類用的 LLM
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_groq import ChatGroq

def _get_classifier_llm():
    """用最快速的 LLM 做意圖分類"""
    groq_key = os.getenv("GROQ_KEY_1") or os.getenv("GROQ_KEY")
    if groq_key:
        return ChatGroq(model="llama-3.1-8b-instant", groq_api_key=groq_key, temperature=0)
    google_key = os.getenv("GOOGLE_API_KEY_1") or os.getenv("GOOGLE_API_KEY")
    if google_key:
        return ChatGoogleGenerativeAI(
            model="gemini-1.5-flash",
            google_api_key=google_key,
            temperature=0,
            convert_system_message_to_human=True,
        )
    raise RuntimeError("No LLM keys available for classifier")

CLASSIFIER_LLM = None

def _classifier():
    global CLASSIFIER_LLM
    if CLASSIFIER_LLM is None:
        CLASSIFIER_LLM = _get_classifier_llm()
    return CLASSIFIER_LLM


# ── 節點實作 ──────────────────────────────────────────────────────────────────

def classify_intent(state: dict) -> dict:
    """
    節點 1：分析用戶輸入，決定意圖和部門
    """
    user_input = state["user_input"]
    prompt = f"""
你是 AI 公司的任務分配系統。分析以下用戶輸入，回傳 JSON。

用戶輸入：{user_input}

可用意圖（按優先級判斷）：
- code: 需要【寫程式碼、開發、建立、實作】任何軟體、網站、應用、腳本、工具、API、系統。
  例如：寫網站、建立 API、開發 App、寫腳本、實作功能、建資料庫
- review: 需要 code review、審查程式碼、檢查 bug、技術審查
- strategy: 需要制定【商業策略、OKR、市場計劃】，純規劃不涉及實作
- prd: 需要撰寫產品需求文件、功能規格書
- marketing: 需要行銷文案、社群貼文、部落格、郵件
- research: 需要市場研究、競品分析、行業報告
- sales: 需要銷售文案、冷郵件、客戶開發
- chat: 一般問答

判斷規則：
- 只要提到「寫」、「建」、「開發」、「實作」、「製作」加上任何軟體/網站/工具，一律是 code
- 「網站」、「App」、「API」、「腳本」相關需求，一律是 code
- 只有純商業規劃、不涉及任何實作，才是 strategy

只回傳 JSON，不要其他文字：
{{"intent": "...", "department": "...", "summary": "30 字內的任務摘要"}}
"""
    try:
        result = _classifier().invoke(prompt)
        content = result.content if hasattr(result, "content") else str(result)
        # 清理可能的 markdown code block
        content = content.strip().strip("```json").strip("```").strip()
        parsed = json.loads(content)
        return {
            "intent": parsed.get("intent", "chat"),
            "department": parsed.get("department", "management"),
            "messages": [{"role": "system", "content": f"Intent classified: {parsed}"}],
        }
    except Exception as e:
        return {
            "intent": "chat",
            "department": "management",
            "error": f"Classification failed: {e}",
        }


def route_to_department(state: dict) -> dict:
    """節點 2：確認路由（可加業務邏輯）"""
    return {"job_id": str(uuid.uuid4())[:8], "retry_count": 0}


def _post_to_crew(endpoint: str, payload: dict, job_id: str) -> dict:
    try:
        resp = httpx.post(
            f"{CREW_SERVICE_URL}/{endpoint}",
            json={**payload, "job_id": job_id},
            timeout=30,
        )
        resp.raise_for_status()
        data = resp.json()
        returned_job_id = data.get("job_id", job_id)
        print(f"[crew] submitted to /{endpoint}, job_id: {returned_job_id}")
        return {"job_status": "queued", "job_id": returned_job_id}
    except Exception as e:
        print(f"[crew] error posting to /{endpoint}: {e}")
        return {"error": str(e), "job_status": "error"}


def call_crew_strategy(state: dict) -> dict:
    """節點：呼叫 CrewAI CEO 制定策略"""
    return _post_to_crew(
        "strategy",
        {"objective": state["user_input"]},
        state["job_id"],
    )


def call_crew_prd(state: dict) -> dict:
    """節點：呼叫 CrewAI PM 撰寫 PRD"""
    return _post_to_crew(
        "product-spec",
        {"feature_request": state["user_input"]},
        state["job_id"],
    )


def call_crew_marketing(state: dict) -> dict:
    """節點：呼叫 CrewAI 市場部"""
    # 試著從輸入推斷 content_type
    user_input = state["user_input"].lower()
    content_type = "blog"
    for ct in ["twitter", "linkedin", "email", "newsletter", "blog"]:
        if ct in user_input:
            content_type = ct
            break
    return _post_to_crew(
        "marketing",
        {"topic": state["user_input"], "content_type": content_type},
        state["job_id"],
    )


def call_crew_research(state: dict) -> dict:
    """節點：呼叫 CrewAI 分析師做市場研究"""
    return _post_to_crew(
        "research",
        {"query": state["user_input"]},
        state["job_id"],
    )


def call_code_generation(state: dict) -> dict:
    """直接用 Groq 生成程式碼，不依賴 OpenHands"""
    from langchain_groq import ChatGroq

    groq_key = os.getenv("GROQ_KEY_1") or os.getenv("GROQ_KEY_2")
    if not groq_key:
        return {"job_status": "error", "error": "No Groq key available"}

    llm = ChatGroq(
        model="llama-3.3-70b-versatile",
        groq_api_key=groq_key,
        temperature=0,
    )

    intent = state.get("intent", "code")
    user_input = state["user_input"]

    if intent == "review":
        prompt = f"""You are a senior software engineer. Review the following code.

Code:
{user_input}

Output in English only:
1. Issue list (severity: High/Medium/Low)
2. Improvement suggestions
3. Fixed complete code
4. Security considerations

Use Markdown format."""

    else:
        prompt = f"""You are a senior software engineer. Complete the following task.

Task: {user_input}

STRICT RULES:
1. ALL code comments, variable names, function names, and output text must be in ENGLISH only. No Chinese characters anywhere in the code.
2. ALL content inside the code (text, labels, placeholders, error messages) must be in ENGLISH.
3. Use ONLY standard library or the most common packages. Do NOT use any framework that requires separate installation (no Astro, no Next.js, no Nuxt, no SvelteKit) unless the user explicitly requested it by name.
4. If building a website, use plain HTML + CSS + JavaScript in a single file, or use a minimal setup (e.g. Vite + vanilla JS). Output must be immediately runnable with minimal setup.
5. If a framework is absolutely necessary, provide the COMPLETE setup commands and explain every step.

Output format:
## Setup
(exact commands to run, copy-paste ready)

## File Structure
(list of files to create)

## Code
(complete code for each file, with English comments)

## How to Run
(exact run command)

Use Markdown format. Output in English only."""

    try:
        result = llm.invoke(prompt)
        return {
            "job_status": "done",
            "openhands_result": result.content,
        }
    except Exception as e:
        return {"job_status": "error", "error": str(e)}


def wait_for_job(state: dict) -> dict:
    job_id = state.get("job_id")
    retry = state.get("retry_count", 0)

    MAX_RETRIES = 100
    POLL_INTERVAL = 5

    if not job_id or job_id == "None":
        return {"job_status": "error", "error": "job_id is None, task was not submitted correctly"}

    if retry >= MAX_RETRIES:
        return {"job_status": "error", "error": "Timeout: job took too long"}

    time.sleep(POLL_INTERVAL)

    try:
        resp = httpx.get(
            f"{CREW_SERVICE_URL}/jobs/{job_id}",
            timeout=10,
        )
        resp.raise_for_status()
        data = resp.json()
        status = data.get("status", "running")
        print(f"[wait_job] CrewAI job {job_id} status: {status}")

        if status == "done":
            return {"job_status": "done", "crew_result": data.get("result", "")}
        elif status == "error":
            return {
                "job_status": "error",
                "error": data.get("error", "CrewAI unknown error"),
            }
        return {"job_status": "running", "retry_count": retry + 1}

    except httpx.HTTPStatusError as e:
        print(f"[wait_job] HTTP error: {e.response.status_code} - {e.response.text}")
        return {"job_status": "running", "retry_count": retry + 1}
    except Exception as e:
        print(f"[wait_job] Exception: {type(e).__name__}: {e}")
        return {"job_status": "running", "retry_count": retry + 1}

def format_response(state: dict) -> dict:
    crew_result = state.get("crew_result") or ""
    openhands_result = state.get("openhands_result") or ""
    
    if crew_result:
        final = crew_result
    elif openhands_result:
        final = openhands_result
    else:
        final = "任務已完成，但沒有收到回應內容。"

    return {
        "final_response": final,
        "messages": [{"role": "assistant", "content": final}],
    }


def handle_error(state: dict) -> dict:
    """節點：錯誤處理"""
    error = state.get("error", "未知錯誤")
    response = (
        f"⚠️ 任務執行遇到問題：\n\n```\n{error}\n```\n\n"
        "請稍後重試，或簡化您的請求。"
    )
    return {
        "final_response": response,
        "messages": [{"role": "assistant", "content": response}],
    }