changing base prompt

main.py

@@ -10,6 +10,8 @@ from rag import *
 from fastapi.responses import StreamingResponse
 import json
 from prompts import *
+from typing import Literal
+from models import *
 
 load_dotenv()
 
@@ -42,12 +44,16 @@ class StyleWriter(BaseModel):
     style: str
     tonality: str
 
+models = ["gpt-4o","gpt-4o-mini","mistral-large-latest"]
+
 class UserInput(BaseModel):
     prompt: str
     enterprise_id: str
     stream: Optional[bool] = False
     messages: Optional[list[dict]] = []
     style_tonality: Optional[StyleWriter] = None
+    marque: Optional[str] = None
+    model: Literal["gpt-4o","gpt-4o-mini","mistral-large-latest"] = "gpt-4o"
 
 
 class EnterpriseData(BaseModel):
@@ -175,11 +181,11 @@ def generate_answer(user_input: UserInput):
             context = ""
 
         if user_input.style_tonality is None:
-            prompt_formated = prompt_reformatting(template_prompt,context,prompt)
-            answer = generate_response_via_langchain(prompt, model="gpt-4o",stream=user_input.stream,context = context , messages=user_input.messages,template=template_prompt)
+            prompt_formated = prompt_reformatting(template_prompt,context,prompt,enterprise_name=getattr(user_input,"marque",""))
+            answer = generate_response_via_langchain(prompt, model=getattr(user_input,"model","gpt-4o"),stream=user_input.stream,context = context , messages=user_input.messages,template=template_prompt,enterprise_name=getattr(user_input,"marque",""))
         else:
-            prompt_formated = prompt_reformatting(template_prompt,context,prompt,style=user_input.style_tonality.style,tonality=user_input.style_tonality.tonality)
-            answer = generate_response_via_langchain(prompt, model="gpt-4o",stream=user_input.stream,context = context , messages=user_input.messages,style=user_input.style_tonality.style,tonality=user_input.style_tonality.tonality,template=template_prompt)
+            prompt_formated = prompt_reformatting(template_prompt,context,prompt,style=user_input.style_tonality.style,tonality=user_input.style_tonality.tonality,enterprise_name=getattr(user_input,"marque",""))
+            answer = generate_response_via_langchain(prompt,model=getattr(user_input,"model","gpt-4o"),stream=user_input.stream,context = context , messages=user_input.messages,style=user_input.style_tonality.style,tonality=user_input.style_tonality.tonality,template=template_prompt,enterprise_name=getattr(user_input,"marque",""))
         
         if user_input.stream:
             return StreamingResponse(stream_generator(answer,prompt_formated), media_type="application/json")
@@ -192,6 +198,10 @@ def generate_answer(user_input: UserInput):
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
+    
+@app.get("/models")
+def get_models():
+    return {"models": models}
 
 
     

models.py

@@ -0,0 +1 @@
+models = ["gpt-4o","gpt-4o-mini","mistral-large-latest"]

rag.py

@@ -7,6 +7,7 @@ from langchain_core.documents import Document
 from langchain_openai import ChatOpenAI
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import PromptTemplate
+from langchain_mistralai import ChatMistralAI
 from uuid import uuid4
 
 import unicodedata
@@ -105,9 +106,16 @@ def generate_response_via_langchain(query: str, stream: bool = False, model: str
 
 
     prompt = PromptTemplate.from_template(template)
+
+    print(f"model: {model}")
+    print(f"marque: {enterprise_name}")
     
     # Initialize the OpenAI LLM with the specified model
-    llm = ChatOpenAI(model=model,temperature=0)
+    if model.startswith("gpt"):
+        llm = ChatOpenAI(model=model,temperature=0)
+    if model.startswith("mistral"):
+        llm = ChatMistralAI(model=model,temperature=0)
+
     
     # Create an LLM chain with the prompt and the LLM
     llm_chain = prompt | llm | StrOutputParser()

requirements.txt

@@ -13,4 +13,4 @@ langchain
 langchain-openai
 langchain-community
 langchain-pinecone
-
+langchain_mistralai