File size: 4,854 Bytes
c737bb6 e6cc6f7 010d51d e6cc6f7 010d51d c737bb6 010d51d e6cc6f7 010d51d e6cc6f7 010d51d e6cc6f7 010d51d c737bb6 010d51d c737bb6 010d51d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
from datetime import datetime
import logging
from typing import List
class ChatManager:
def __init__(self, documentManager, llmManager):
"""
Initialize the ChatManager.
"""
self.doc_manager = documentManager
self.llm_manager = llmManager
logging.info("ChatManager initialized")
def generate_chat_response(self, query: str, selected_docs: List[str], history: List[dict]) -> List[dict]:
"""
Generate a chat response based on the user's query and selected documents.
Args:
query (str): The user's query.
selected_docs (List[str]): List of selected document filenames from the dropdown.
history (List[dict]): The chat history as a list of {'role': str, 'content': str} dictionaries.
Returns:
List[dict]: Updated chat history with the new response in 'messages' format.
"""
timestamp = datetime.now().strftime("%H:%M:%S")
logging.info(f"Generating chat response for query: {query} at {timestamp}")
# Handle empty query
if not query:
logging.warning("Empty query received")
return history + [{"role": "assistant", "content": "Please enter a query."}]
# Handle no selected documents
if not selected_docs:
logging.warning("No documents selected")
return history + [{"role": "assistant", "content": "Please select at least one document."}]
# Retrieve the top 5 chunks based on the query and selected documents
try:
top_k_results = self.doc_manager.retrieve_top_k(query, selected_docs, k=5)
except Exception as e:
logging.error(f"Error retrieving chunks: {str(e)}")
return history + [
{"role": "user", "content": f"{query}"},
{"role": "assistant", "content": f"Error retrieving chunks: {str(e)}"}
]
if not top_k_results:
logging.info("No relevant chunks found")
return history + [
{"role": "user", "content": f"{query}"},
{"role": "assistant", "content": "No relevant information found in the selected documents."}
]
# Send the top K results to the LLM to generate a response
try:
llm_response, source_docs = self.llm_manager.generate_response(query, top_k_results)
except Exception as e:
logging.error(f"Error generating LLM response: {str(e)}")
return history + [
{"role": "user", "content": f"{query}"},
{"role": "assistant", "content": f"Error generating response: {str(e)}"}
]
# Format the response
response = llm_response
# Uncomment to include source docs in response (optional)
# for i, doc in enumerate(source_docs, 1):
# doc_id = doc.metadata.get('doc_id', 'Unknown')
# filename = next((name for name, d_id in self.doc_manager.document_ids.items() if d_id == doc_id), 'Unknown')
# response += f"\n{i}. {filename}: {doc.page_content[:100]}..."
logging.info("Chat response generated successfully")
# Return updated history with new user query and LLM response
return history + [
{"role": "user", "content": f"{query}"},
{"role": "assistant", "content": response}
]
def generate_summary(self, chunks: any, summary_type: str = "medium") -> str:
"""
Generate a summary of the selected documents.
Args:
selected_docs (List[str]): List of selected document filenames.
summary_type (str): Type of summary ("small", "medium", "detailed").
k (int): Number of chunks to retrieve from DocumentManager.
include_toc (bool): Whether to include the table of contents (if available).
Returns:
str: Generated summary.
Raises:
ValueError: If summary_type is invalid or DocumentManager/LLM is not available.
"""
if summary_type not in ["small", "medium", "detailed"]:
raise ValueError("summary_type must be 'small', 'medium', or 'detailed'")
if not chunks:
logging.warning("No documents selected for summarization")
return "Please select at least one document."
llm_summary_response = self.llm_manager.generate_summary_v0(chunks = chunks)
#logging.info(f" Summary response {llm_summary_response}")
return llm_summary_response
def generate_sample_questions(self, chunks: any):
questions = self.llm_manager.generate_questions(chunks = chunks)
return questions
|