chore: AI Chat improved

Backend/app/api/v1/endpoints/notes.py

@@ -2,13 +2,12 @@ from fastapi import APIRouter, Depends, HTTPException, status, File, UploadFile
 from sqlalchemy.ext.asyncio import AsyncSession
 from app.models import User
 from app.models.tables import PDFData
-from app.api.deps import get_db, get_current_user
+from app.api.deps import get_db, get_current_user, get_chroma_collection
 from app.schema import AI_chat_input
 from app.llm import stream_chat
 import uuid
 from fastapi.responses import StreamingResponse
 from chromadb.api.models.Collection import Collection 
-from app.api.deps import get_chroma_collection
 from pathlib import Path
 from llama_index.readers.file import PyMuPDFReader
 from llama_index.core.node_parser import SentenceSplitter
@@ -16,7 +15,7 @@ from typing import Annotated
 import shutil
 import os
 from sentence_transformers import SentenceTransformer
-
+from .quiz import search_logic
 
 router = APIRouter(prefix="/notes")
 
@@ -28,12 +27,15 @@ embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
 @router.post("/stream_chat", response_class=StreamingResponse)
 async def ai_chat(
     Input_model: AI_chat_input, 
+    collection: Collection = Depends(get_chroma_collection), 
     current_user: User = Depends(get_current_user)
 ):
     messages_dict = [msg.model_dump() for msg in Input_model.messages]
+    query = f"{Input_model.context};{Input_model.messages[-1].content}"
+    retrieved_docs: str | None = await search_logic(query, collection)
 
     return StreamingResponse(
-        stream_chat(messages_dict, Input_model.context),
+        stream_chat(messages_dict, Input_model.context, retrieved_docs),
         media_type="text/plain"
     )
 

Backend/app/api/v1/endpoints/quiz.py

@@ -9,15 +9,52 @@ from chromadb.api.models.Collection import Collection
 from app.api.deps import get_chroma_collection
 from app.llm import call_llm
 import uuid
+import logging
+
 
 router = APIRouter(prefix="/quiz")
 
+
+# 1. Set up a logger (if you haven't already globally)
+logger = logging.getLogger("uvicorn.error") # reusing uvicorn's logger ensures it shows up in your terminal
+
 async def search_logic(query: str, collection: Collection):
-    results = await collection.query(
-        query_texts=[query],
-        n_results=5
-    )
-    return ''.join(results['documents'][0])
+    # Log the incoming query
+    logger.info(f"🔍 [Search Logic] Starting search for query: '{query}'")
+
+    try:
+        results = await collection.query(
+            query_texts=[query],
+            n_results=5
+        )
+        
+        # Log the raw results to see exactly what ChromaDB returned (helps spot NoneTypes)
+        logger.info(f"📄 [Search Logic] Raw results from DB: {results}")
+
+        if results and results.get('documents') and len(results['documents']) > 0:
+            raw_docs = results['documents'][0]
+            
+            # Filter None values and Log how many were found vs valid
+            valid_docs = [str(doc) for doc in raw_docs if doc is not None]
+            
+            logger.info(f"✅ [Search Logic] Processing: Found {len(raw_docs)} items. Valid text items: {len(valid_docs)}")
+            
+            if len(raw_docs) != len(valid_docs):
+                logger.warning("⚠️ [Search Logic] Warning: Some documents contained NoneType and were skipped.")
+
+            # Join with a space (safer than empty string)
+            final_context = " ".join(valid_docs)
+            return final_context
+            
+        else:
+            logger.warning("⚠️ [Search Logic] No documents found for this query.")
+            return ""
+
+    except Exception as e:
+        # Log the full error if something crashes
+        logger.error(f"❌ [Search Logic] CRITICAL ERROR: {str(e)}")
+        # You might want to re-raise the error or return empty depending on your needs
+        return ""
 
 @router.get("/search_docs")
 async def search_documents(

Backend/app/llm.py

@@ -37,36 +37,61 @@ async def call_llm(prompt:str):
 
 
 
-async def stream_chat(messages:List[dict], context:str):
+from typing import List
+
+async def stream_chat(messages: List[dict], context: str, retrieved_docs: str | None):
     system_instruction = {
         "role": "system", 
-        "content": "You are a helpful AI assistant. Answer the user's question strictly based on the provided context."
+        "content": "You are a helpful AI assistant. Answer the user's question based on the provided context and retrieved documents."
     }
     
     conversation_history = [msg.copy() for msg in messages]
 
     if conversation_history and conversation_history[-1]['role'] == 'user':
         last_user_msg = conversation_history[-1]
-        # Rewrite the content to: Context + \n\n + Question
-        last_user_msg['content'] = (
-            f"Here is the context/notes you must use:\n"
-            f"---------------------\n"
-            f"{context}\n"
-            f"---------------------\n\n"
-            f"User Question: {last_user_msg['content']}"
-        )
+        original_question = last_user_msg['content']
+        
+        # Start constructing the augmented prompt
+        augmented_content = ""
+
+        # 1. Add Manual Context
+        if context:
+            augmented_content += (
+                f"Here is the context/notes you must use:\n"
+                f"---------------------\n"
+                f"{context}\n"
+                f"---------------------\n\n"
+            )
+
+        # 2. Add Retrieved Documents (New Logic)
+        if retrieved_docs:
+            augmented_content += (
+                f"Here is background information/retrieved documents:\n"
+                f"---------------------\n"
+                f"{retrieved_docs}\n"
+                f"---------------------\n\n"
+            )
+
+        # 3. Add the User Question
+        augmented_content += f"User Question: {original_question}"
+
+        # Update the message content
+        last_user_msg['content'] = augmented_content
+
     else:
         # Fallback: If for some reason there is no user message, add one.
+        # We combine context and docs here too just in case.
+        combined_context = f"{context}\n\n{retrieved_docs or ''}"
         conversation_history.append({
             "role": "user", 
-            "content": f"Context:\n{context}\n\nPlease analyze this."
+            "content": f"Context:\n{combined_context}\n\nPlease analyze this."
         })
 
-    # 3. Combine System + Modified User History
+    # 4. Combine System + Modified User History
     full_history = [system_instruction] + conversation_history
 
     try:
-        
+        # Ensure 'client' is initialized before this function in your code
         stream = await client.chat.completions.create(
             model="openai/gpt-oss-20b",
             messages=full_history,