from fastapi import FastAPI, Request
from pydantic import BaseModel
from llama_cpp import Llama

model_path = "mistral-7b-v0.1.Q4_K_M.gguf"

llm = Llama(
    model_path = model_path,
    n_ctx = 4096,
    n_threads = 4,
)

app = FastAPI()

class DSAPrompt(BaseModel):
    prompt: str

@app.post("/solve_dsa_problem")
async def solve_dsa_problem(item: DSAPrompt):
    system_prompt = '''
                    You are a Data Structures and Algorithm problem solver. You are given the following problem and you need to solve it.
                    Give a detailed explanation of the solution approach and the code in C++, Java, or Python.
                    If the input is not a DSA problem, politely refuse their request and reinsist to provide a DSA problem.
                    '''
    
    res = llm.create_chat_completion(
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": item.prompt}
        ],
        temperature = 0.9,
        max_tokens = 2048,
    )

    return res