Merge pull request #37 from AET-DevOps25/feature/genai-improvement

MErge walids pr Feature/genai improvement

Author: yassinsws

genAi/compose.yml

@@ -3,7 +3,7 @@ services:
     build:
       context: .
       dockerfile: Dockerfile
-    container_name: studymate-genai-dev
+    container_name: genai-genai-dev
     restart: unless-stopped
     ports:
       - "8081:8081"  # FastAPI GenAI service
@@ -20,7 +20,7 @@ services:
 
   weaviate:
     image: cr.weaviate.io/semitechnologies/weaviate:1.30.3
-    container_name: studymate-weaviate-dev
+    container_name: genai-weaviate-dev
     restart: unless-stopped
     command:
       - --host

genAi/llm.py

@@ -28,9 +28,10 @@ class StudyLLM:
         base_url="https://gpu.aet.cit.tum.de/api/"
     )
 
-    def __init__(self, doc_path: str):
+    def __init__(self):
         base_system_template = ("You are an expert on the information in the context given below.\n"
-                                     "Use the context as your primary knowledge source. If you can't fulfill your task given the context, just say that.\n"
+                                     "Use the context as your primary knowledge source.\n"
+                                    "Do not mention that you have a context, just use it to fulfill your task.\n"
                                     "context: {context}\n"
                                     "Your task is {task}"
                                     )
@@ -39,12 +40,27 @@ def __init__(self, doc_path: str):
             ('human', '{input}')
         ])
         try: 
-            self.rag_helper = RAGHelper(doc_path)
+            self.rag_helper = RAGHelper()
         except Exception as e:
             raise ValueError(f"Error initializing RAGHelper: {e}")
 
 
-    async def prompt(self, prompt: str) -> str:
+    async def load_document(self, doc_name: str, path: str, user_id: str):
+        """
+        Load a document into the RAG system.
+        
+        Args:
+            doc_name (str): The name of the document to load.
+            user_id (str): The ID of the user loading the document.
+        """
+        try:
+            await self.rag_helper.load_document(doc_name, path, user_id)
+            return f"Document {doc_name} loaded successfully for user {user_id}."
+        except Exception as e:
+            error_msg = f"Failed to load document {doc_name} for user {user_id}: {str(e)}"
+            raise ValueError(error_msg)
+    
+    async def prompt(self, prompt: str, user_id: str) -> str:
         """
         Call the LLM with a given prompt.
         
@@ -59,7 +75,7 @@ async def prompt(self, prompt: str) -> str:
             "If you're asked a question that does not relate to your context, do not answer it - instead, answer by saying you're only familiar with <the topic in your context>.\n"
             )
 
-        context = self.rag_helper.retrieve(prompt, top_k=5)
+        context = self.rag_helper.retrieve(query=prompt, user_id=user_id, top_k=5)
         chain = self.base_prompt_template | self.chat_llm
         response = await chain.ainvoke({
             'context': context,
@@ -69,7 +85,7 @@ async def prompt(self, prompt: str) -> str:
 
         return response.content
 
-    async def summarize(self):
+    async def summarize(self, document_name: str, user_id: str) -> str:
         """
         Summarize the given document using the LLM.
         
@@ -102,30 +118,36 @@ async def summarize(self):
             combine_prompt=combine_prompt
         )
 
-        result = await chain.ainvoke({"input_documents": self.rag_helper.summary_chunks})
+        doc_name = f"{user_id}_{document_name}"
+        chunks = self.rag_helper.get_generation_chunks(user_id, doc_name)
+        result = await chain.ainvoke({"input_documents": chunks})
 
         return result["output_text"]
 
-    async def generate_flashcards(self):
+    async def generate_flashcards(self, document_name: str, user_id: str):
         """
         Generate flashcards from the document using the LLM.
         
         Returns:
             list: A list of flashcard objects.
         """
         flashcard_chain = FlashcardChain(self.generation_llm)
-        cards = await flashcard_chain.invoke(self.rag_helper.summary_chunks)
+        doc_name = f"{user_id}_{document_name}"
+        chunks = self.rag_helper.get_generation_chunks(user_id, doc_name)
+        cards = await flashcard_chain.invoke(chunks)
         return cards
 
-    async def generate_quiz(self):
+    async def generate_quiz(self, document_name: str, user_id: str):
         """
         Generate a quiz from the document using the LLM.
         
         Returns:
             list: A quiz object.
         """
         quiz_chain = QuizChain(self.generation_llm)
-        quiz = await quiz_chain.invoke(self.rag_helper.summary_chunks)
+        doc_name = f"{user_id}_{document_name}"
+        chunks = self.rag_helper.get_generation_chunks(user_id, doc_name)
+        quiz = await quiz_chain.invoke(chunks)
         return quiz
 
     def cleanup(self):

genAi/main.py

@@ -3,7 +3,7 @@
 from fastapi import FastAPI
 from fastapi.responses import JSONResponse
 from helpers import save_document
-from request_models import CreateSessionRequest, PromptRequest, SummaryRequest, QuizRequest, FlashcardRequest
+from request_models import LoadDocumentRequest, PromptRequest, SummaryRequest, QuizRequest, FlashcardRequest
 from llm import StudyLLM
 from prometheus_fastapi_instrumentator import Instrumentator
 
@@ -12,14 +12,13 @@
 logger = logging.getLogger(__name__)
 
 
-llm_instances: dict[str, StudyLLM] = {}
+llm_instance: StudyLLM = StudyLLM()
 
 @asynccontextmanager
 async def lifespan(_):
     yield
     # Shutdown: cleanup
-    for llm in llm_instances.values():
-        llm.cleanup()
+    llm_instance.cleanup()
 
 app = FastAPI(
     title="tutor",
@@ -48,8 +47,8 @@ async def lifespan(_):
     ).instrument(app).expose(app)
 
 
-# llm_instances["dummy"] = StudyLLM("./documents/example/W07_Microservices_and_Scalable_Architectures.pdf") # TODO: remove
-# llm_instances["dummy2"] = StudyLLM("./documents/example/dummy_knowledge.txt") # TODO: remove
+# StudyLLM(user_id='dummy', doc_path="./documents/example/dummy_knowledge.txt") # TODO: remove
+# llm_instances["dummy2"] = StudyLLM(user_id='dummy', doc_path="./documents/example/W07_Microservices_and_Scalable_Architectures.pdf") # TODO: remove
 
 # Auxiliary Endpoints
 @app.get("/health")
@@ -62,44 +61,34 @@ async def health_check():
 
 
 # AI Tasks Endpoints
-@app.post("/session/load")
-async def load_session(data: CreateSessionRequest):
+@app.post("/document")
+async def load_document(data: LoadDocumentRequest):
     """
-    Create a new session with the LLM for a given document URL.
+    Load a new document in the LLM instance.
     """
     try:
-        if data.session_id in llm_instances:
-            logger.info(f"Session {data.session_id} already exists")
-            return {"message": "Session already loaded."}
-        
-        logger.info(f"Creating new session {data.session_id} for document {data.document_name}")
-        doc_name = f"{data.session_id}_{data.document_name}"
+        logger.info(f"Loading new document {data.document_name} for user {data.user_id}")
+        doc_name = f"{data.user_id}_{data.document_name}"
         path = save_document(doc_name, data.document_base64)
-        llm_instances[data.session_id] = StudyLLM(path)
-        logger.info(f"Session {data.session_id} created successfully")
-        return {"message": "Session created successfully."}
+        await llm_instance.load_document(doc_name, path, data.user_id)
+        logger.info(f"Document {doc_name} created successfully")
+        return {"message": "Document loaded successfully."}
     except Exception as e:
-        error_msg = f"Failed to create session {data.session_id}: {str(e)}"
+        error_msg = f"Failed to load document {doc_name}: {str(e)}"
         logger.error(error_msg)
         return {"error": error_msg}
 
-
 @app.post("/chat")
 async def receive_prompt(data: PromptRequest):
     """
     Receive a prompt and return a response from the LLM.
     """
     try:
-        if data.session_id not in llm_instances:
-            error_msg = f"Session {data.session_id} not found. Please ensure the document was processed successfully."
-            logger.error(error_msg)
-            return JSONResponse(status_code=404, content={"response": f"ERROR: {error_msg}"})
-        
-        logger.info(f"Processing chat request for session {data.session_id}")
-        response = await llm_instances[data.session_id].prompt(data.message)
+        logger.info(f"Processing chat request for user {data.user_id}")
+        response = await llm_instance.prompt(data.message, data.user_id)
         return {"response": response}
     except Exception as e:
-        error_msg = f"Chat error for session {data.session_id}: {str(e)}"
+        error_msg = f"Chat error for user {data.user_id}: {str(e)}"
         logger.error(error_msg)
         return {"response": f"ERROR: {error_msg}"}
 
@@ -109,16 +98,11 @@ async def generate_summary(data: SummaryRequest):
     Receive a summary request and return a summary from the LLM.
     """
     try:
-        if data.session_id not in llm_instances:
-            error_msg = f"Session {data.session_id} not found. Please ensure the document was processed successfully."
-            logger.error(error_msg)
-            return {"response": f"ERROR: {error_msg}"}
-        
-        logger.info(f"Generating summary for session {data.session_id}")
-        response = await llm_instances[data.session_id].summarize()
+        logger.info(f"Generating summary for user {data.user_id}, document {data.document_name}")
+        response = await llm_instance.summarize(data.document_name, data.user_id)
         return {"response": response}
     except Exception as e:
-        error_msg = f"Summary generation error for session {data.session_id}: {str(e)}"
+        error_msg = f"Summary generation error for user {data.user_id}: {str(e)}"
         logger.error(error_msg)
         return {"response": f"ERROR: {error_msg}"}
 
@@ -128,17 +112,12 @@ async def generate_flashcards(data: FlashcardRequest):
     Receive a flashcard request and return flashcard objects from the LLM.
     """
     try:
-        if data.session_id not in llm_instances:
-            error_msg = f"Session {data.session_id} not found. Please ensure the document was processed successfully."
-            logger.error(error_msg)
-            return {"response": {"flashcards": [], "error": error_msg}}
-        
-        logger.info(f"Generating flashcards for session {data.session_id}")
-        response = await llm_instances[data.session_id].generate_flashcards()
-        logger.info(f"Flashcards generated successfully for session {data.session_id}")
+        logger.info(f"Generating flashcards for user {data.user_id}, document {data.document_name}")
+        response = await llm_instance.generate_flashcards(data.document_name, data.user_id)
+        logger.info(f"Flashcards generated successfully for user {data.user_id}")
         return {"response": response}
     except Exception as e:
-        error_msg = f"Flashcard generation error for session {data.session_id}: {str(e)}"
+        error_msg = f"Flashcard generation error for user {data.user_id}: {str(e)}"
         logger.error(error_msg)
         return {"response": {"flashcards": [], "error": error_msg}}
 
@@ -148,44 +127,11 @@ async def generate_quiz(data: QuizRequest):
     Receive a quiz request and return a quiz object from the LLM.
     """
     try:
-        if data.session_id not in llm_instances:
-            error_msg = f"Session {data.session_id} not found. Please ensure the document was processed successfully."
-            logger.error(error_msg)
-            return {"response": {"questions": [], "error": error_msg}}
-        
-        logger.info(f"Generating quiz for session {data.session_id}")
-        response = await llm_instances[data.session_id].generate_quiz()
-        logger.info(f"Quiz generated successfully for session {data.session_id}")
+        logger.info(f"Generating quiz for user {data.user_id}, document {data.document_name}")
+        response = await llm_instance.generate_quiz(data.document_name, data.user_id)
+        logger.info(f"Quiz generated successfully for user {data.user_id}")
         return {"response": response}
     except Exception as e:
-        error_msg = f"Quiz generation error for session {data.session_id}: {str(e)}"
+        error_msg = f"Quiz generation error for user {data.user_id}: {str(e)}"
         logger.error(error_msg)
         return {"response": {"questions": [], "error": error_msg}}
-
-@app.post("/process")
-async def process_document(data: SummaryRequest):
-    """Compatibility endpoint for Kotlin genai-service (/process).
-    It creates a session (if not present) and immediately returns QUEUED.
-    (Actual processing e.g. summary generation can be triggered asynchronously.)"""
-    try:
-        session_id = data.session_id or data.document_id
-        if session_id not in llm_instances:
-            logger.info(f"/process received – creating session {session_id}")
-            # save document and create LLM instance
-            doc_name = f"{session_id}_{data.document_name or 'doc.pdf'}"
-            path = save_document(doc_name, data.document_base64 or "")
-            llm_instances[session_id] = StudyLLM(path)
-        return {
-            "requestId": session_id,
-            "status": "QUEUED",
-            "message": "Document queued for processing",
-            "estimatedTime": None
-        }
-    except Exception as e:
-        logger.error(f"/process error: {str(e)}")
-        return {
-            "requestId": None,
-            "status": "FAILED",
-            "message": f"Failed to process document: {str(e)}",
-            "estimatedTime": None
-        }