memory up

main.py

@@ -13,6 +13,7 @@ import json
 from prompts import *
 from typing import Literal
 from models import *
+from fastapi.middleware.cors import CORSMiddleware
 
 load_dotenv()
 
@@ -77,6 +78,7 @@ models = ["gpt-4o","gpt-4o-mini","mistral-large-latest"]
 class UserInput(BaseModel):
     prompt: str
     enterprise_id: str
+    user_id: Optional[str] = None
     stream: Optional[bool] = False
     messages: Optional[list[dict]] = []
     style_tonality: Optional[StyleWriter] = None
@@ -159,6 +161,17 @@ def get_documents(enterprise_id: str):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
     
+@app.get("/documents/memory/{enterprise_id}/{user_id}")
+def get_documents(enterprise_id: str, user_id: str):
+    try:
+        docs_names = []
+        for ids in  index.list(prefix=f"kb_{user_id}_", namespace=enterprise_id):
+            for id in ids:
+                docs_names.append(id)
+        return docs_names
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
+    
 @app.delete("/documents/{enterprise_id}/{filename_id}")
 def delete_document(enterprise_id: str, filename_id: str):
     try:
@@ -168,6 +181,24 @@ def delete_document(enterprise_id: str, filename_id: str):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
     
+@app.delete("/documents/memory/all/{enterprise_id}/{user_id}/")
+def delete_document(enterprise_id: str, user_id: str):
+    try:
+        for ids in index.list(prefix=f"kb_{user_id}_", namespace=enterprise_id):
+            index.delete(ids=ids, namespace=enterprise_id)
+        return {"message": "Document deleted", "chunks_deleted": ids}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
+    
+@app.delete("/documents/memory/{enterprise_id}/{user_id}/{info_id}")
+def delete_document(enterprise_id: str, user_id: str, info_id: str):
+    try:
+        for ids in index.list(prefix=f"{info_id}", namespace=enterprise_id):
+            index.delete(ids=ids, namespace=enterprise_id)
+        return {"message": "Document deleted", "chunks_deleted": ids}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
+    
 @app.delete("/documents/all/{enterprise_id}")
 def delete_all_documents(enterprise_id: str):
     try:
@@ -181,13 +212,13 @@ import asyncio
 
 GENERATION_TIMEOUT_SEC = 60
 
-async def stream_generator(response, prompt):
+async def stream_generator(response, prompt, info_memoire):
     async with async_timeout.timeout(GENERATION_TIMEOUT_SEC):
         try:
             async for chunk in response:
                 if isinstance(chunk, bytes):
                     chunk = chunk.decode('utf-8')  # Convert bytes to str if needed
-                yield json.dumps({"prompt": prompt, "content": chunk})
+                yield json.dumps({"prompt": prompt, "content": chunk, "info_memoire":info_memoire})
         except asyncio.TimeoutError:
             raise HTTPException(status_code=504, detail="Stream timed out")
 
@@ -204,6 +235,11 @@ def generate_answer(user_input: UserInput):
         context = get_retreive_answer(enterprise_id, prompt, index, common_namespace)
 
         #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 = ""
@@ -237,12 +273,13 @@ def generate_answer(user_input: UserInput):
                                                      index=index)
         
         if user_input.stream:
-            return StreamingResponse(stream_generator(answer,prompt_formated), media_type="application/json")
+            return StreamingResponse(stream_generator(answer,prompt_formated,info_memoire), media_type="application/json")
         
         return {
             "prompt": prompt_formated,
             "answer": answer,
             "context": context,
+            "info_memoire": infos_added_to_kb
         }
 
     except Exception as e:

rag.py

@@ -83,7 +83,7 @@ def get_vectorstore(text_chunks,filename, file_type,namespace,index,enterprise_n
         return False
     
 
-def add_to_knowledge_base(enterprise_id,information,index,enterprise_name):
+def add_to_knowledge_base(enterprise_id,information,index,enterprise_name,user_id=""):
     ''' Add information to the knowledge base 
     Args:
         enterprise_id (str): the enterprise id
@@ -96,13 +96,13 @@ def add_to_knowledge_base(enterprise_id,information,index,enterprise_name):
 
         document = Document(
             page_content=information,
-            metadata={"filename":"knowledge_base","file_type":"text", "filename_id":"knowledge_base", "entreprise_name":enterprise_name},
+            metadata={"filename":"knowledge_base","file_type":"text", "filename_id":"knowledge_base", "entreprise_name":enterprise_name, "user_id":user_id},
         )
 
-        uuid = f"knowledge_base_{uuid4()}"
+        uuid = f"kb_{user_id}_{uuid4()}"
 
         vector_store.add_documents(documents=[document], id=uuid)
-        return True
+        return uuid
     
     except Exception as e:
         print(e)
@@ -158,7 +158,7 @@ def get_retreive_answer(enterprise_id,prompt,index,common_id):
         print(e)
         return False
     
-def handle_calling_add_to_knowledge_base(query,enterprise_id = "",index = "",enterprise_name = "",llm = None):
+def handle_calling_add_to_knowledge_base(query,enterprise_id = "",index = "",enterprise_name = "",user_id = "",llm = None):
     ''' Handle the calling of the add_to_knowledge_base function
     if the user, in his query, wants to add information to the knowledge base, the function will be called
     '''
@@ -167,7 +167,7 @@ def handle_calling_add_to_knowledge_base(query,enterprise_id = "",index = "",ent
 
         Determine if the user wants to add something to the knowledge base.
 
-        - If the user wants to add something, output 'add' followed by the content to add.
+        - If the user wants to add something, extract the valuable information, reformulate and output 'add' followed by the information.
         - If the user does not want to add something, output 'no action'.
 
         Ensure your response is only 'add <content>' or 'no action'.
@@ -180,7 +180,7 @@ def handle_calling_add_to_knowledge_base(query,enterprise_id = "",index = "",ent
     prompt = PromptTemplate.from_template(template)
 
     if not llm:
-        llm = ChatOpenAI(model="gpt-4o-mini",temperature=0)
+        llm = ChatOpenAI(model="gpt-4o",temperature=0)
 
     llm_with_tool = llm.bind_tools([AddToKnowledgeBase])
 
@@ -193,10 +193,12 @@ def handle_calling_add_to_knowledge_base(query,enterprise_id = "",index = "",ent
     if response.startswith("add"):
         item = response[len("add"):].strip()
         if item:
-            add_to_knowledge_base(enterprise_id,item,index,enterprise_name)
+            item_id = add_to_knowledge_base(enterprise_id,item,index,enterprise_name,user_id)
             print("added to knowledge base")
-            return True
-        
+            
+            print(item)
+            return {"item_id":item_id,"item":item}
+    
     print(response)
     return False