from tools.code_generator import generate_code
from tools.web_search import search_web
from tools.rag_engine import answer_from_docs
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Default general model (TinyLlama)
tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
model.eval()

def route_query(prompt):
    prompt_lower = prompt.lower()

    if "code:" in prompt_lower:
        return generate_code(prompt)
    elif "search:" in prompt_lower:
        return search_web(prompt)
    elif "doc:" in prompt_lower:
        return answer_from_docs(prompt)
    else:
        inputs = tokenizer(prompt, return_tensors="pt")
        outputs = model.generate(**inputs, max_new_tokens=200)
        return tokenizer.decode(outputs[0], skip_special_tokens=True)