import gradio as gr
from llama_cpp import Llama
import threading

MODELS = {
    "knowledge": ("bartowski/Qwen2.5-14B-Instruct-GGUF", "Qwen2.5-14B-Instruct-Q4_K_M.gguf"),
    "coder": ("bartowski/DeepSeek-Coder-V2-Lite-Instruct-GGUF", "DeepSeek-Coder-V2-Lite-Instruct-Q4_K_M.gguf"),
    "reasoning": ("bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF", "DeepSeek-R1-Distill-Qwen-7B-Q4_K_M.gguf"),
    "deep": ("bartowski/DeepSeek-R1-Distill-Qwen-14B-GGUF", "DeepSeek-R1-Distill-Qwen-14B-Q4_K_M.gguf"),
}

llms = {}

def get_llm(name):
    if name not in llms:
        repo, filename = MODELS[name]
        llms[name] = Llama.from_pretrained(repo_id=repo, filename=filename, n_ctx=2048, verbose=False)
    return llms[name]

def ask(q):
    results = {}
    def query(name):
        try:
            llm = get_llm(name)
            prompt = f"Expert. Answer precisely.\n\nQuestion: {q}\n\nAnswer:"
            out = llm(prompt, max_tokens=512, temperature=0.7)
            results[name] = out["choices"][0]["text"]
        except:
            pass
    
    threads = [threading.Thread(target=query, args=(n,)) for n in MODELS]
    for t in threads: t.start()
    for t in threads: t.join()
    
    valid = {k: v for k, v in results.items() if v}
    if not valid:
        return "エラー: 応答なし"
    
    parts = "\n\n".join([f"=== {k} ===\n{v[:300]}" for k, v in valid.items()])
    fusion_prompt = f"Synthesize into ONE ultimate answer.\n\nQUESTION: {q}\n\nEXPERT OUTPUTS:\n{parts}\n\nULTIMATE ANSWER:"
    llm = get_llm("knowledge")
    final = llm(fusion_prompt, max_tokens=1024, temperature=0.7)
    return final["choices"][0]["text"]

demo = gr.Interface(
    fn=ask,
    inputs=gr.Textbox(label="質問"),
    outputs=gr.Textbox(label="FUSION-1"),
    title="FUSION-1",
    description="4AI融合エンジン"
)

demo.launch()