Test v2

main.py

@@ -226,7 +226,6 @@ async def stream_generator(response, prompt, info_memoire):
             raise HTTPException(status_code=504, detail="Stream timed out")
 
     
-
 @app.post("/generate-answer/")
 def generate_answer(user_input: UserInput):
     try:
@@ -289,6 +288,80 @@ def generate_answer(user_input: UserInput):
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
+
+
+
+async def stream_generator2(response, prompt, info_memoire):
+    async with async_timeout.timeout(GENERATION_TIMEOUT_SEC):
+        try:
+            async for chunk in response:
+                if isinstance(chunk, bytes):
+                    yield chunk.decode('utf-8')  # Convert bytes to str if needed
+        except asyncio.TimeoutError:
+            raise HTTPException(status_code=504, detail="Stream timed out")
+            
+@app.post("/v2/generate-answer/")
+def generate_answer(user_input: UserInput):
+    try:
+        print(user_input)
+        
+        prompt = user_input.prompt
+        enterprise_id = user_input.enterprise_id
+
+        template_prompt = base_template
+
+        context = get_retreive_answer(enterprise_id, prompt, index, common_namespace, user_id=user_input.user_id)
+
+        #final_prompt_simplified = prompt_formatting(prompt,template,context)
+        infos_added_to_kb = handle_calling_add_to_knowledge_base(prompt,enterprise_id,index,getattr(user_input,"marque",""),user_id=getattr(user_input,"user_id",""))
+        if infos_added_to_kb:
+            prompt = prompt + "l'information a été ajoutée à la base de connaissance: " + infos_added_to_kb['item']
+        else:
+            infos_added_to_kb = {}
+
+        if not context:
+            context = ""
+
+        if user_input.style_tonality is None:
+            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",""),
+                                                     enterprise_id=enterprise_id,
+                                                     index=index)
+        else:
+            prompt_formated = prompt_reformatting(template_prompt,
+                                                  context,
+                                                  prompt,
+                                                  style=getattr(user_input.style_tonality,"style","neutral"),
+                                                  tonality=getattr(user_input.style_tonality,"tonality","formal"),
+                                                  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=getattr(user_input.style_tonality,"style","neutral"),
+                                                     tonality=getattr(user_input.style_tonality,"tonality","formal"),
+                                                     template=template_prompt,
+                                                     enterprise_name=getattr(user_input,"marque",""),
+                                                     enterprise_id=enterprise_id,
+                                                     index=index)
+        
+        if user_input.stream:
+            return StreamingResponse(stream_generator2(answer,prompt_formated,infos_added_to_kb), media_type="application/json")
+        
+        return {
+            "prompt": prompt_formated,
+            "answer": answer,
+            "context": context,
+            "info_memoire": infos_added_to_kb
+        }
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
     
 @app.get("/models")
 def get_models():