added introductory prompt

backend.py

@@ -5,8 +5,7 @@ from interface import GemmaLLMInterface
 from llama_index.core.node_parser import SentenceSplitter
 from llama_index.embeddings.instructor import InstructorEmbedding
 import gradio as gr
-from llama_index.core import ChatPromptTemplate
-from llama_index.core import Settings, VectorStoreIndex, SimpleDirectoryReader, PromptTemplate, load_index_from_storage
+from llama_index.core import Settings, VectorStoreIndex, SimpleDirectoryReader, PromptTemplate, load_index_from_storage, StorageContext
 from llama_index.core.node_parser import SentenceSplitter
 from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
@@ -46,11 +45,15 @@ documents_paths = {
     'payment': 'data/payment'
 }
 
-session_state = {"documents_loaded": False, 
+session_state = {"index": False,
+                 "documents_loaded": False, 
                  "document_db": None, 
                  "original_message": None, 
                  "clarification": False}
 
+PERSIST_DIR = "./db"
+os.makedirs(PERSIST_DIR, exist_ok=True)
+
 ############################---------------------------------
 
 # Get the parser
@@ -66,47 +69,21 @@ def build_index(path: str):
     # Build the vector store index from the nodes
     index = VectorStoreIndex(nodes)
     
+    storage_context = StorageContext.from_defaults()
+    index.storage_context.persist(persist_dir=PERSIST_DIR)
+    
     return index
 
 
-# Global variables
-global_index = None
-global_session_state = {}
-
-def initialize_global_state():
-    global global_index, global_session_state
-    global_index = None
-    global_session_state = {
-        "documents_loaded": False,
-        "document_db": None,
-        "original_message": None,
-        "clarification": False,
-        "conversation": []
-    }
-
-# Call this at the start of your script
-initialize_global_state()
 
 @spaces.GPU(duration=30)
 def handle_query(query_str: str, 
-                 chat_history: list[tuple[str, str]], 
-                 session: dict[str, Any]) -> Iterator[str]:
-    global global_index, global_session_state
-    
-    # Update global session state with any new information from the passed session
-    global_session_state.update(session)
+                 chat_history: list[tuple[str, str]]) -> Iterator[str]:
     
-    # Update conversation history
-    for user, assistant in chat_history:
-        global_session_state["conversation"].extend([
-            ChatMessage(role=MessageRole.USER, content=user),
-            ChatMessage(role=MessageRole.ASSISTANT, content=assistant),
-        ])
+    global conversation
     
-    # Add current query to conversation
-    global_session_state["conversation"].append(ChatMessage(role=MessageRole.USER, content=query_str))
     
-    if global_index is None:
+    if not session_state["index"]:
         matched_path = None
         words = query_str.lower()
         for key, path in documents_paths.items():
@@ -114,39 +91,73 @@ def handle_query(query_str: str,
                 matched_path = path
                 break
         if matched_path:
-            global_index = build_index(matched_path)
-            global_session_state["documents_loaded"] = True
-        else:
-            global_index = build_index("data/chiarimento.txt")
-            global_session_state["clarification"] = True
+            index = build_index(matched_path)
+            session_state["index"] = True
+            
+        else: ## CHIEDI CHIARIMENTO
+            conversation: List[ChatMessage] = []
+            for user, assistant in chat_history:
+                conversation.extend(
+                    [
+                        ChatMessage(role=MessageRole.USER, content=user),
+                        
+                        ChatMessage(role=MessageRole.ASSISTANT, content=assistant),
+                    ]
+                )
+            index = build_index("data/chiarimento.txt")
+            
+                 
+    
+    else:
+      
+        
+        # The index is already built, no need to rebuild it.
+        conversation: List[ChatMessage] = []
+        for user, assistant in chat_history:
+            conversation.extend(
+                [
+                    ChatMessage(role=MessageRole.USER, content=user),
+                    
+                    ChatMessage(role=MessageRole.ASSISTANT, content=assistant),
+                ]
+            )
+
+        #conversation.append( ChatMessage(role=MessageRole.USER, content=query_str))
+        #pass    
 
     try:
+         
+        storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
+        index = load_index_from_storage(storage_context)  
         memory = ChatMemoryBuffer.from_defaults(token_limit=None)
         
-        chat_engine = global_index.as_chat_engine(
-            chat_mode="condense_plus_context",
-            memory=memory,
-            similarity_top_k=4,
-            response_mode="tree_summarize",
-            context_prompt = (
-                "Sei un assistente Q&A italiano di nome Odi, che risponde solo alle domande o richieste pertinenti in modo preciso."
-                " Quando un utente ti chiede informazioni su di te o sul tuo creatore puoi dire che sei un assistente ricercatore creato dagli Osservatori Digitali e fornire gli argomenti di cui sei esperto."
-                " Ecco i documenti rilevanti per il contesto:\n"
-                "{context_str}"
-                "\nIstruzione: Usa la cronologia delle chat precedenti, o il contesto sopra, per interagire e aiutare l'utente a rispondere alla sua domanda."
-            ),
-            verbose=False,
+        chat_engine = index.as_chat_engine(
+        chat_mode="condense_plus_context",
+        memory=memory,
+        similarity_top_k=4, 
+        response_mode="tree_summarize", #Good for summarization purposes
+
+        context_prompt = (
+        "Sei un assistente Q&A italiano di nome Odi, che risponde solo alle domande o richieste pertinenti in modo preciso."
+        " Quando un utente ti chiede informazioni su di te o sul tuo creatore puoi dire che sei un assistente ricercatore creato dagli Osservatori Digitali e fornire gli argomenti di cui sei esperto."
+        " Ecco i documenti rilevanti per il contesto:\n"
+        "{context_str}"
+        "\nIstruzione: Usa la cronologia delle chat precedenti, o il contesto sopra, per interagire e aiutare l'utente a rispondere alla sua domanda."
+          ),
+        verbose=False,
         )
         
-        response = chat_engine.stream_chat(query_str, global_session_state["conversation"])
         
         outputs = []
+        response = chat_engine.stream_chat(query_str, conversation)
+        #response = chat_engine.chat(query_str)
         for token in response.response_gen:
+          #if not token.startswith("system:") and not token.startswith("user:"):
+          
             outputs.append(token)
+            #print(f"Generated token: {token}")
             yield "".join(outputs)
-        
-        # Update the session with any changes
-        session.update(global_session_state)
+          
         
     except Exception as e:
         yield f"Error processing query: {str(e)}"