Update app.py

app.py

@@ -16,27 +16,6 @@ client = Client("on1onmangoes/CNIHUB101324v10", hf_token=HF_TOKEN)
 # Update the conversation history within the function.
 # Return the updated history along with any other required outputs.
 
-def format_answer_with_documents(answer):
-    """
-    This function formats the assistant's answer and separates the personal details from the relevant documents.
-    """
-    # Extract the personal details and documents from the answer
-    personal_details = answer[0]
-    relevant_documents = answer[1:]  # Assuming documents are included after the answer
-
-    # Format the personal details
-    formatted_answer = f"Based on the documents provided, the personal details are as follows:\n\n{personal_details}\n\n"
-
-    # If there are relevant documents, format them as a separate section
-    if relevant_documents:
-        formatted_answer += "Relevant Documents:\n"
-        for idx, document in enumerate(relevant_documents, start=1):
-            formatted_answer += f"{idx}. {document.metadata['heading']} (Page {int(document.metadata['page_number'])})\n"
-            formatted_answer += f"   Source: {document.metadata['source']}\n"
-            formatted_answer += f"   Snippet: {document.page_content[:200]}...\n\n"  # Showing a snippet of content
-
-    return formatted_answer
-
 def stream_chat_with_rag(
     message: str,
     history: list,
@@ -66,13 +45,13 @@ def stream_chat_with_rag(
     # Call the API with the user's message
     question = message
     answer = client.predict(question=question, api_name="/answer_with_rag")
-    
-    # Debugging: Print the response
-    print("The Answer in stream_chat_with_rag:")
+
+    # Debugging: Print the raw response
+    print("Raw answer from API:")
     print(answer)
 
-    # Format the assistant's answer to separate the relevant documents
-    formatted_answer = format_answer_with_documents(answer)
+    # Format the assistant's answer and the relevant documents separately
+    formatted_answer = format_answer_string(answer)
 
     # Update the conversation history with the new message and answer
     history.append((message, formatted_answer))
@@ -81,6 +60,35 @@ def stream_chat_with_rag(
     return formatted_answer
 
 
+def format_answer_string(answer: str):
+    """
+    This function takes the full string returned by the API (which includes both the assistant's answer and documents)
+    and splits it into separate sections: the assistant's answer and the relevant documents.
+    """
+    # Step 1: Split the answer by the point where document context starts
+    split_marker = "Extracted Documents with Descriptive Context:"
+    if split_marker in answer:
+        assistant_answer, documents_section = answer.split(split_marker, 1)
+    else:
+        # If no documents found, return the whole answer as the assistant's response
+        return f"Assistant: {answer.strip()}"
+    
+    # Step 2: Clean and format the assistant's answer
+    formatted_answer = f"Assistant: {assistant_answer.strip()}\n\n"
+
+    # Step 3: Format the documents section
+    formatted_answer += "Relevant Documents:\n"
+    document_entries = documents_section.split("Document ")
+    
+    # Step 4: Reformat each document entry
+    for entry in document_entries[1:]:  # Skip the first empty split
+        doc_number, rest = entry.split(":", 1)
+        formatted_answer += f"{doc_number}. {rest.strip()}\n\n"
+
+    return formatted_answer
+
+
+
 # this version works just issue with formatting 
 # def stream_chat_with_rag(
 #     message: str,