main.py

import os
from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from typing import List
from typing_extensions import TypedDict
from typing import Annotated
from langgraph.graph.message import AnyMessage, add_messages
from langchain_core.messages import HumanMessage, AIMessage
from langgraph.graph import END, StateGraph, START
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from database import create_table, return_title, return_exercise

create_table()

class Request(BaseModel):
    enonce : str
    code : str

app = FastAPI()

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

load_dotenv()
os.environ["GROQ_API_KEY"] = os.getenv('GROQ_API_KEY')
os.environ["TOKENIZERS_PARALLELISM"] = "false"

llm = ChatGroq(model="llama-3.1-70b-versatile", temperature=0.5)
#llm = ChatGroq(model="mixtral-8x7b-32768", temperature=0.5)
exercise = ""
system = """
Tu es un expert Python.
Tu dois aider un élève à résoudre un exercice de programmation Python.
Tu ne dois jamais donner la correction de l'exercice (même partiellement) à l'élève, juste lui donner des inddications lui permettant de résoudre lui même l'exercice
Tu dois t'adresser directement à l'élève.
L'élève ne peut pas te poser des questions, il peut juste te proposer son code.
Tu ne dois pas proposer à l'élève de te poser des questions
Il est inutile de proposer à l'élève de tester son code avec les exemples proposés.
Tu ne dois pas proposer aux élèves des modifications du programme qui sorte du cadre de l'exercice. Par exemple, pour l'exercice qui demande d'écrire une fonction moyenne, si dans l'énoncé il est précisé que l'on a un tableau non vide d'entier en paramètre, il est inutile de dire à l'élève que son programme doit gérer les tableaux vides. 
Tu dois t'exprimer en français
"""

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "Voici l'exercice proposé à l'élève: \n {enonce} \n\n Voici le programme proposé par l'élève : {code}"),
    ]
)

chain = prompt | llm | StrOutputParser()

class GraphState(TypedDict):
    enonce : str
    messages: Annotated[list[AnyMessage], add_messages]
    

def chatbot(state : GraphState):
    response = chain.invoke({'enonce': state['enonce'].content, 'code' : state['messages'][-1].content})
    return {"messages": [AIMessage(content=response)]}

workflow = StateGraph(GraphState)
workflow.add_node('chatbot', chatbot)

workflow.add_edge(START, 'chatbot')
workflow.add_edge('chatbot', END)

app_chatbot = workflow.compile()

@app.post('/request')
def request(req: Request):
    rep = app_chatbot.invoke({"enonce" : HumanMessage(content=req.enonce),"messages": [HumanMessage(content=req.code)]}, stream_mode="values")
    return {"response":rep['messages'][-1].content}


@app.get('/title')
def get_title():
    tab_title = return_title()
    return {"title":tab_title}

@app.get('/exercise/{id}')
def get_exercise(id : int):
    ex = return_exercise(id)
    return {"title" : ex[1].replace("\n",""), "enonce" : ex[2], "test": ex[3]}