style and ton

main.py

@@ -35,17 +35,24 @@ index = pc.Index(index_name)
 
 app = FastAPI()
 
+class StyleWriter(BaseModel):
+    style: str
+    tonality: str
 
 class UserInput(BaseModel):
     prompt: str
     enterprise_id: str
     stream: Optional[bool] = False
     messages: Optional[list[dict]] = []
+    style_tonality: Optional[StyleWriter] = None
+
 
 class EnterpriseData(BaseModel):
     name: str
     id: Optional[str] = None
 
+
+
 tasks = []
 
 @app.get("/")
@@ -124,7 +131,7 @@ async def stream_generator(response):
     async with async_timeout.timeout(GENERATION_TIMEOUT_SEC):
         try:
             async for chunk in response:
-                yield {"content":chunk}
+                yield "random data"
         except asyncio.TimeoutError:
             raise HTTPException(status_code=504, detail="Stream timed out")
     
@@ -139,7 +146,7 @@ def generate_answer(user_input: UserInput):
         if not context:
             context = "No context found"
 
-        answer = generate_response_via_langchain(prompt, model="gpt-4o",stream=user_input.stream,context = context , messages=user_input.messages)
+        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)
         
         if user_input.stream:
             return StreamingResponse(answer, media_type="application/json")

rag.py

@@ -82,9 +82,9 @@ def get_retreive_answer(enterprise_id,prompt,index):
         print(e)
         return False
     
-def generate_response_via_langchain(query: str, stream: bool = False, model: str = "gpt-4o-mini",context:str="",messages = []) :
+def generate_response_via_langchain(query: str, stream: bool = False, model: str = "gpt-4o-mini",context:str="",messages = [],style:str="formal",tonality:str="neutral"):
     # Define the prompt template
-    template = "Sachant le context suivant: {context}, et l'historique de la conversation: {messages}, {query}"
+    template = "En tant qu'IA experte en marketing, réponds avec un style {style} et une tonalité {tonality} dans ta communcation, sachant le context suivant: {context}, et l'historique de la conversation: {messages}, {query}"
     prompt = PromptTemplate.from_template(template)
     
     # Initialize the OpenAI LLM with the specified model
@@ -95,10 +95,10 @@ def generate_response_via_langchain(query: str, stream: bool = False, model: str
 
     if stream:
         # Return a generator that yields streamed responses
-        return llm_chain.astream({ "query": query, "context": context, "messages": messages})
+        return llm_chain.astream({ "query": query, "context": context, "messages": messages, "style": style, "tonality": tonality })
     
     # Invoke the LLM chain and return the result
-    return llm_chain.invoke({ "query": query, "context": context, "messages": messages})
+    return llm_chain.invoke({"query": query, "context": context, "messages": messages, "style": style, "tonality": tonality})