from fastapi import FastAPI
from langserve import add_routes
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_huggingface import HuggingFacePipeline
from fastapi.middleware.cors import CORSMiddleware
import torch

app = FastAPI(
    title="Phi-3 Code Explainer API",
    description="Explains Python code using the Phi-3 Mini model and LangChain.",
    version="1.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

model_id = "microsoft/phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, torch_dtype=torch.float16, device_map="auto"
)

generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=256,
    return_full_text=False,
    do_sample=False
)

llm = HuggingFacePipeline(pipeline=generator)

prompt = ChatPromptTemplate.from_template(
    "### Instruction:\nExplain this Python code step-by-step:\nUse code with caution\n{code}\n\n### Explanation:"
)

parser = StrOutputParser()
chain = prompt | llm | parser

add_routes(app, chain, path="/explain")