from __future__ import annotations
import typing as t

from langchain.chains import LLMChain
from langchain.llms import OpenLLM
from langchain.prompts import PromptTemplate
from pydantic import BaseModel

import bentoml
from bentoml.io import JSON, Text

prompt = PromptTemplate(
    input_variables=["name", "industry", "work_years", "education", "skills"],
    template="""
You are a self-introduction creator with a strong background in interviewing and persuasive writing.
The content you create lets the interviewer grasp the candidate's unique value proposition swiftly,
aligning it with the job's requirements and the company's ethos. You are given the following information
to craft a compelling self-introduction for a software engineer role.

Name: {name}
Industry: {industry}
Experience: {work_years}
Education: {education}
Skills: {skills}
Self introduction:
    """,
)

# Function to generate an OpenLLM instance
def gen_llm(model_name: str, model_id: str | None = None) -> OpenLLM:
    lc_llm = OpenLLM(model_name=model_name, model_id=model_id, embedded=False)
    # Download the model for the LLM instance
    lc_llm.runner.download_model()
    return lc_llm
   
llm = gen_llm("dolly-v2", model_id="databricks/dolly-v2-12b")
chain = LLMChain(llm=llm, prompt=prompt)

# Define a BentoML Service with the LLM Runner
svc = bentoml.Service("self-introduction-service", runners=[llm.runner])

class Query(BaseModel):
    name: str
    industry: str
    work_years: str
    education: str
    skills: t.List[str]
    llm_config: t.Dict[str, t.Any]
    
SAMPLE_INPUT = Query(
    name="John Smith",
    industry="AI",
    work_years="5 years",
    education="MSc, Computer Science, Stanford University",
    skills=["Kubernetes", "BentoML", "AWS", "Python", "Go", "Java"],
    llm_config=llm.runner.config.model_dump(),
)

# Define the API for the Service with JSON as input and plain text as output
@svc.api(input=JSON.from_sample(sample=SAMPLE_INPUT), output=Text())
def generate(query: Query):
    return chain.run(
    {"name": query.name, "industry": query.industry, "work_years": query.work_years, "education": query.education,
     "skills": ", ".join(query.skills)}
)