Commit 83fe60f6 authored by minhnguyengp1's avatar minhnguyengp1
Browse files

baseline done! v2

parent 86b7f2a0
......@@ -32,23 +32,38 @@ def create_agent_graph():
new_state = state.copy()
rag_agent = ResearchAssistantRAG(config)
messages: List[BaseMessage] = state["messages"]
query: Union[str, Dict] = state["current_input"]
rag_context_limit: int = config.rag.context_limit
print(f"DEBUG messages: {messages}")
buffer_memory = WindowBufferMemory(context_limit=rag_context_limit)
chat_history: List[Dict[str, str]] = buffer_memory.build_context(messages)
print("DEBUG chat_history input:", chat_history)
print("DEBUG type of chat_history:", type(chat_history))
if chat_history:
for i, msg in enumerate(chat_history):
print(f"DEBUG chat_history[{i}]: {msg} (type: {type(msg)})")
response: Dict[str, Any] = rag_agent.process_query(query, chat_history=chat_history)
# messages: List[BaseMessage] = state["messages"]
# query: Union[str, Dict] = state["current_input"]
# rag_context_limit: int = config.rag.context_limit
#
# print(f"DEBUG messages: {messages}")
#
# buffer_memory = WindowBufferMemory(context_limit=rag_context_limit)
# chat_history: List[Dict[str, str]] = buffer_memory.build_context(messages)
#
# print("DEBUG chat_history input:", chat_history)
# print("DEBUG type of chat_history:", type(chat_history))
messages = state["messages"]
query = state["current_input"]
rag_context_limit = config.rag.context_limit
print(f"checkdd")
recent_context = ""
for msg in messages[-rag_context_limit:]: # limit controlled from config
if isinstance(msg, HumanMessage):
# print("######### DEBUG 1:", msg)
recent_context += f"User: {msg.content}\n"
elif isinstance(msg, AIMessage):
# print("######### DEBUG 2:", msg)
recent_context += f"Assistant: {msg.content}\n"
# if chat_history:
# for i, msg in enumerate(chat_history):
# print(f"DEBUG chat_history[{i}]: {msg} (type: {type(msg)})")
print("Before processing query:", recent_context)
response: Dict[str, Any] = rag_agent.process_query(query, chat_history=recent_context)
### CHECKPOINT
print("Final response from RAG agent:", response)
......
......@@ -5,118 +5,160 @@ import uuid
from pathlib import Path
import json
from .content_processor import ContentProcessor
from .doc_parser import DocParser
from .query_expander import QueryExpander
from .vector_store import QdrantRetriever
from .document_processor import DocumentProcessor
from .query_processor import QueryProcessor
from .reranker import Reranker
from .response_generator import ResponseGenerator
from .data_ingestion import DataIngestion
from .multi_query_generator import MultiQueryGenerator
from .query_rewriter import QueryRewriter
from .vectorstore_qdrant import VectorStore
class ResearchAssistantRAG:
def __init__(self, config):
self.config = config
self._initialize()
# self._initialize()
self.logger = logging.getLogger(f"{self.__module__}")
self.logger.info("Initializing Research Assistant RAG system")
self.doc_parser = DocParser()
self.content_processor = ContentProcessor(config)
self.parsed_content_dir = self.config.rag.parsed_content_dir
self.vector_store = VectorStore(config)
self.reranker = Reranker(config)
self.response_generator = ResponseGenerator(config)
self.parsed_content_dir = self.config.rag.parsed_content_dir
def ingest_file(self, document_path: str) -> Dict[str, Any]:
"""
Ingest a single file into the RAG system.
Args:
document_path: Path to the file to ingest
Returns:
Dictionary with ingestion results
"""
start_time = time.time()
self.logger.info(f"Ingesting file: {document_path}")
def _initialize(self):
try:
self.logger = logging.getLogger(f"{self.__module__}")
self.logger.info("Initializing Research Assistant RAG system")
self.llm = self.config.rag.llm
self.logger.info(f"Using LLM: {type(self.llm).__name__}")
self.embedding_model = self.config.rag.embedding_model
self.logger.info(f"Using embedding model: {type(self.embedding_model).__name__}")
self.sparse_embedder = self.config.rag.sparse_embedder
self.query_expander = QueryExpander(self.config)
self.query_rewriter = QueryRewriter(self.llm)
self.query_processor = QueryProcessor(self.config, self.embedding_model, self.sparse_embedder)
self.document_processor = DocumentProcessor(self.config, self.embedding_model)
self.query_generator = MultiQueryGenerator(llm=self.llm, max_queries=self.config.rag.max_queries or 5)
self.retriever = QdrantRetriever(self.config)
self.response_generator = ResponseGenerator(self.config, self.llm)
total_docs = self.retriever.count_documents()
self.logger.info(f"Vector store contains {total_docs} documents")
if total_docs == 0:
self.logger.warning("No documents in vector store. Results may be limited.")
self.top_k = getattr(self.config.rag, "top_k", 5)
self.similarity_threshold = getattr(self.config.rag, "similarity_threshold", 0.0)
self.data_ingestion = DataIngestion()
# Step 1: Parse document
self.logger.info("1. Parsing document and extracting images...")
parsed_document, images = self.doc_parser.parse_document(document_path, self.parsed_content_dir)
self.logger.info(f" Parsed document and extracted {len(images)} images")
# Step 2: Summarize images
self.logger.info("2. Summarizing images...")
image_summaries = self.content_processor.summarize_images(images)
self.logger.info(f" Generated {len(image_summaries)} image summaries")
# Step 3: Format document with image summaries
self.logger.info("3. Formatting document with image summaries...")
formatted_document = self.content_processor.format_document_with_images(parsed_document, image_summaries)
# Save parsed_document to data/parsed
parsed_dir = Path("data/parsed")
parsed_dir.mkdir(exist_ok=True, parents=True)
parsed_path = parsed_dir / (Path(document_path).stem + "_parsed.json")
with open(parsed_path, "w", encoding="utf-8") as f:
import json
json.dump(formatted_document, f, ensure_ascii=False, indent=2, default=str)
# Step 4: Chunk document into semantic sections
self.logger.info("4. Chunking document into semantic sections...")
document_chunks: List[str] = self.content_processor.chunk_document(formatted_document)
self.logger.info(f" Document split into {len(document_chunks)} chunks")
with open("document_chunks.json", "w", encoding="utf-8") as f:
json.dump(document_chunks, f, ensure_ascii=False, indent=2)
# Step 5: Create vector store and document store
self.logger.info("5. Creating vector store knowledge base...")
self.vector_store.create_vectorstore(
document_chunks=document_chunks,
document_path=document_path
)
self.logger.info("Research Assistant RAG system successfully initialized")
return {
"success": True,
"documents_ingested": 1,
"chunks_processed": len(document_chunks),
"processing_time": time.time() - start_time
}
except Exception as e:
self.logger.error(f"Initialization error: {e}")
import traceback
self.logger.error(traceback.format_exc())
raise
self.logger.error(f"Error ingesting file: {e}")
return {
"success": False,
"error": str(e),
"processing_time": time.time() - start_time
}
def process_query(self, query: str, chat_history: Optional[List[Dict[str, str]]] = None) -> Dict[str, Any]:
self.logger.info(f"RAG Agent processing query: {query}")
"""
Process a query with the RAG system.
# Process query and return result, passing chat_history
result = self.query(query, chat_history)
Args:
query: The query string
chat_history: Optional chat history for context
return result
def query(self, query: str, chat_history: Optional[List[Dict[str, str]]] = None) -> Dict[str, Any]:
exec_count = 1
self.logger.info(f"Processing query: {query}")
Returns:
Response dictionary
"""
start_time = time.time()
self.logger.info(f"RAG Agent processing query: {query}")
# Process query and return result, passing chat_history
try:
query_embedding: Dict[str, Any]
filters: Dict[str, Any]
if chat_history:
query = self.query_rewriter.rewrite(query, chat_history)
query_embedding, filters = self.query_processor.process_query(query)
### DEBUG
# dense_vector = query_embedding.get("dense")
# sparse_vector = query_embedding.get("sparse")
#
# if dense_vector:
# print("DEBUG dense_vector[:5]:", dense_vector[:5])
# print("DEBUG dense_vector length:", len(dense_vector))
#
# if sparse_vector:
# print("DEBUG sparse_vector type:", type(sparse_vector))
# print("DEBUG sparse_vector indices[:5]:", getattr(sparse_vector, "indices", [])[:5])
# print("DEBUG sparse_vector values[:5]:", getattr(sparse_vector, "values", [])[:5])
#
# print("DEBUG has as_object:", hasattr(sparse_vector, "as_object"))
###
# TODO: Enhance retrieval by incorporating chat history context into the query embedding or retrieval strategy
search_results = self._run_vector_search(query_embedding, filters, query)
if self.similarity_threshold > 0:
search_results = [item for item in search_results if item.get('score', 0) >= self.similarity_threshold]
self.logger.info(f"{len(search_results)} results passed similarity threshold")
print(f"DEBUG chat_history at {exec_count}th time: \n{chat_history}")
exec_count += 1
# Step 1: Expand query - REWRITE QUERY
# self.logger.info(f"1. Expanding query: '{query}'")
# expansion_result = self.query_expander.expand_query(query)
# expanded_query = expansion_result["expanded_query"]
# self.logger.info(f" Original: '{query}'")
# self.logger.info(f" Expanded: '{expanded_query}'")
# query = expanded_query
# Step 2: Retrieval
self.logger.info(f"2. Retrieving relevant documents for the query: '{query}'")
vectorstore, docstore = self.vector_store.load_vectorstore()
retrieved_documents = self.vector_store.retrieve_relevant_chunks(
query=query,
vectorstore=vectorstore,
docstore=docstore,
)
self.logger.info(f" Retrieved {len(retrieved_documents)} relevant document chunks")
print(f"type of retrieved_documents = {type(retrieved_documents)}")
print(f"retrieved_documents = {retrieved_documents}")
# Step 3: Rerank the retrieved documents if we have a reranker and enough documents
self.logger.info(f"3. Reranking the retrieved documents")
if self.reranker and len(retrieved_documents) > 1:
# reranked_documents, reranked_top_k_picture_paths = self.reranker.rerank(query, retrieved_documents,
# self.parsed_content_dir)
reranked_documents = self.reranker.rerank(query, retrieved_documents,
self.parsed_content_dir)
self.logger.info(f" Reranked retrieved documents and chose top {len(reranked_documents)}")
# self.logger.info(f" Found {len(reranked_top_k_picture_paths)} referenced images")
else:
self.logger.info(f" Could not rerank the retrieved documents, falling back to original scores")
reranked_documents = retrieved_documents
reranked_top_k_picture_paths = []
# Step 4: Generate response
self.logger.info("4. Generating response...")
response = self.response_generator.generate_response(
query=query,
retrieved_docs=search_results,
retrieved_docs=reranked_documents,
# picture_paths=reranked_top_k_picture_paths,
chat_history=chat_history
)
# Add timing information
processing_time = time.time() - start_time
response["processing_time"] = processing_time
......@@ -126,7 +168,7 @@ class ResearchAssistantRAG:
self.logger.error(f"Error processing query: {e}")
import traceback
self.logger.error(traceback.format_exc())
# Return error response
return {
"response": f"I encountered an error while processing your query: {str(e)}",
"sources": [],
......@@ -134,83 +176,6 @@ class ResearchAssistantRAG:
"processing_time": time.time() - start_time
}
def ingest_documents(self, documents: List[Dict[str, Any]]) -> Dict[str, Any]:
self.logger.info(f"Ingesting {len(documents)} documents")
start_time = time.time()
###
for i, doc in enumerate(documents):
print(f"Document {i} keys: {list(doc.keys())}")
###
try:
processed_dir = Path("data/processed")
processed_dir.mkdir(exist_ok=True, parents=True)
all_chunks: List[Dict[str, Any]] = []
chunk_ids: List[str] = []
for document in documents:
content = document.get("content", "")
metadata = document.get("metadata", {})
elements = document.get("elements", [])
if "id" not in metadata:
metadata["id"] = str(uuid.uuid4())
chunks_from_file = self.document_processor.process_document(content, metadata, elements)
if chunks_from_file:
all_chunks.extend(chunks_from_file)
for chunk in chunks_from_file:
chunk_ids.append(chunk["id"])
doc_path = processed_dir / f"{chunk['id']}.json"
with open(doc_path, 'w', encoding='utf-8') as f:
json_safe_chunk = self._ensure_json_serializable(chunk)
json.dump(json_safe_chunk, f, indent=2)
if not all_chunks:
return {
"success": False,
"error": "No documents were successfully processed",
"processing_time": time.time() - start_time
}
# Here don’t recompute the embedding again, simply take it from all_chunks
upsert_batch: List[Dict[str, Any]] = []
for chunk in all_chunks:
chunk_record = {
"id": chunk["id"],
"content": chunk["content"],
"embedding": chunk["embedding"],
"metadata": chunk["metadata"]
}
upsert_batch.append(chunk_record)
self.retriever.upsert_chunks(upsert_batch)
return {
"success": True,
"documents_ingested": len(documents),
"chunks_processed": len(all_chunks),
# todo: change to "chunk_ids": chunk_ids
"document_ids": chunk_ids,
"processing_time": time.time() - start_time
}
except Exception as e:
self.logger.error(f"Error ingesting documents: {e}")
import traceback
self.logger.error(traceback.format_exc())
return {
"success": False,
"error": str(e),
"documents_processed": 0,
"chunks_created": 0,
"chunks_inserted": 0,
"processing_time": time.time() - start_time
}
def _ensure_json_serializable(self, obj):
if isinstance(obj, dict):
......@@ -223,113 +188,8 @@ class ResearchAssistantRAG:
return str(obj)
def ingest_file(self, file_path: str) -> Dict[str, Any]:
start_time = time.time()
self.logger.info(f"Ingesting file: {file_path}")
try:
ingestion_result = self.data_ingestion.ingest_file(file_path)
if not ingestion_result["success"]:
return {
"success": False,
"error": ingestion_result.get("error", "Unknown error during file ingestion"),
"processing_time": time.time() - start_time
}
documents = []
if "document" in ingestion_result:
documents = [ingestion_result["document"]]
elif "documents" in ingestion_result:
documents = ingestion_result["documents"]
if documents:
return self.ingest_documents(documents)
else:
return {
"success": False,
"error": "No valid documents found in file",
"processing_time": time.time() - start_time
}
except Exception as e:
self.logger.error(f"Error ingesting file: {e}")
return {
"success": False,
"error": str(e),
"processing_time": time.time() - start_time
}
def ingest_directory(self, directory_path: str, file_extension: Optional[str] = None) -> Dict[str, Any]:
pass
def _run_vector_search(self, query_embedding: Dict[str, Any], filters: Dict[str, Any], query: str):
self.logger.info(f"Running vector search with filters: {filters}")
search_results = self.retriever.retrieve(
query_embedding=query_embedding,
filters=filters,
)
self.logger.info(f"Retrieved {len(search_results)} results from vector search")
return search_results
# def query(self, query: str, chat_history: Optional[List[Dict[str, str]]] = None) -> Dict[str, Any]:
# exec_count = 1
# self.logger.info(f"Processing query: {query}")
# start_time = time.time()
#
# try:
# # query_embedding: Dict[str, Any]
# # filters: Dict[str, Any]
# # query_embedding, filters = self.query_processor.process_query(query)
#
# # # TODO: Enhance retrieval by incorporating chat history context into the query embedding or retrieval strategy
# # search_results = self._run_vector_search(query_embedding, filters, query)
#
# # Step 1: Generate multiple sub-queries from the user query and chat history
# sub_queries = self.query_generator.generate_queries(
# user_question=query,
# chat_history=chat_history,
# )
#
# print("sub_queries: ", sub_queries)
#
# all_results = []
# for sq in sub_queries:
# sq_embedding, sq_filters = self.query_processor.process_query(sq)
# results = self._run_vector_search(sq_embedding, sq_filters, sq)
# all_results.extend(results)
#
# # 🧼 Lọc trùng theo nội dung
# seen = set()
# deduped_results = []
# for item in all_results:
# content = item.get("content")
# if content and content not in seen:
# deduped_results.append(item)
# seen.add(content)
#
# search_results = deduped_results
#
# print(f"search_results: {search_results}")
#
# if self.similarity_threshold > 0:
# search_results = [item for item in search_results if item.get('score', 0) >= self.similarity_threshold]
# self.logger.info(f"{len(search_results)} results passed similarity threshold")
#
# print(f"DEBUG chat_history at {exec_count}th time: \n{chat_history}")
# exec_count += 1
#
# response = self.response_generator.generate_response(
# query=query,
# retrieved_docs=search_results,
# chat_history=chat_history
# )
#
# processing_time = time.time() - start_time
# response["processing_time"] = processing_time
#
# return response
import os
import json
import logging
from pathlib import Path
import pandas as pd
from typing import List, Dict, Any, Optional, Union
import os
os.environ["PATH"] += os.pathsep + "/opt/homebrew/bin"
from unstructured.partition.pdf import partition_pdf
from unstructured.chunking.title import chunk_by_title
import json
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
class DataIngestion:
def __init__(self):
self.stats = {
"files_processed": 0,
"documents_ingested": 0,
"errors": 0
}
logger.info("DataIngestion initialized")
def ingest_directory(self, directory_path: str, file_extension: Optional[str] = None) -> Dict[str, Any]:
logger.info(f"Processing directory: {directory_path}")
try:
directory = Path(directory_path)
if not directory.exists() or not directory.is_dir():
raise ValueError(f"Directory does not exist: {directory_path}")
if file_extension:
files = list(directory.glob(f"*{file_extension}"))
else:
files = [f for f in directory.iterdir() if f.is_file()]
logger.info(f"Found {len(files)} files to process")
for file_path in files:
try:
self.ingest_file(str(file_path))
self.stats["files_processed"] += 1
except Exception as e:
logger.error(f"Error processing file {file_path}: {e}")
self.stats["errors"] += 1
return self.stats
except Exception as e:
logger.error(f"Error processing directory: {e}")
return self.stats
def ingest_file(self, file_path: str) -> Dict[str, Any]:
print(f"Processing file (file_path): {file_path}")
logger.info(f"Processing file: {file_path}")
file_path = Path(file_path)
if not file_path.exists():
raise FileNotFoundError(f"File not found: {file_path}")
if file_path.suffix.lower() == '.txt':
return self._ingest_text_file(file_path)
elif file_path.suffix.lower() == '.pdf':
return self._ingest_pdf_file(file_path)
else:
logger.warning(f"Unsupported file format: {file_path.suffix}")
return {"success": False, "error": "Unsupported file format"}
def _ingest_text_file(self, file_path: Path) -> Dict[str, Any]:
try:
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read()
metadata = {
"source": file_path.name,
"file_type": "txt"
}
document = {
"content": content,
"metadata": metadata
}
logger.info(f"Successfully ingested text file: {file_path}")
self.stats["documents_ingested"] += 1
return {"success": True, "document": document}
except Exception as e:
logger.error(f"Error ingesting text file: {e}")
return {"success": False, "error": str(e)}
def _ingest_pdf_file(self, file_path: Path) -> Dict[str, Any]:
logger.info(f"Processing PDF with: {file_path}")
try:
elements = partition_pdf(
file_path,
extract_images_in_pdf=False,
extract_tables=True,
infer_table_structure=True,
languages=["de"],
ocr_languages = "deu"
# chunking_strategy="by_title"
)
##
with open("debug_elements.txt", "w", encoding="utf-8") as f:
for i, element in enumerate(elements):
f.write(f"\n=== Element {i} ===\n")
f.write(f"Category: {getattr(element, 'category', 'N/A')}\n")
f.write(f"Text: {str(element)}\n")
f.write(f"Metadata: {getattr(element, 'metadata', {})}\n")
##
content_parts = []
tables = []
images = []
metadata_parts = {}
for element in elements:
if hasattr(element, "category"):
element_text = str(element)
print(f"Element text: {element_text}")
element_category = element.category
print(f"Element category: {element_category}")
if element_category == "Title":
content_parts.append(f"\n## {element_text}\n")
elif element_category == "NarrativeText":
content_parts.append(element_text)
elif element_category == "ListItem":
content_parts.append(f"- {element_text}")
elif element_category == "Table":
formatted_table_text = self._format_table_as_markdown(element_text)
content_parts.append(f"\n{formatted_table_text}\n")
if hasattr(element, "metadata") and element.metadata:
table_metadata = element.metadata.__dict__ if hasattr(element.metadata,
"__dict__") else element.metadata
tables.append({
"text": formatted_table_text,
"raw_text": element_text,
"metadata": table_metadata
})
elif element_category == "Image":
content_parts.append(f"\n[IMAGE: {element_text}]\n")
if hasattr(element, "metadata") and element.metadata:
images.append({
"text": element_text,
"metadata": element.metadata.__dict__ if hasattr(element.metadata,
"__dict__") else element.metadata
})
else:
content_parts.append(element_text)
if hasattr(element, "metadata") and element.metadata:
metadata = element.metadata.__dict__ if hasattr(element.metadata,
"__dict__") else element.metadata
for key, value in metadata.items():
if key not in metadata_parts:
metadata_parts[key] = value
content = "\n".join(content_parts)
metadata = {
"source": file_path.name,
"file_type": "pdf",
"has_tables": len(tables) > 0,
"table_count": len(tables),
"has_images": len(images) > 0,
"image_count": len(images)
}
metadata.update(metadata_parts)
document = {
"content": content,
"metadata": metadata,
"tables": tables,
"images": images,
"elements": [{"category": getattr(e, "category", "Unknown"), "text": str(e)} for e in elements]
}
logger.info(f"Successfully ingested PDF file using unstructured.io: {file_path}")
self.stats["documents_ingested"] += 1
return {"success": True, "document": document}
except Exception as e:
logger.error(f"Error ingesting PDF file with unstructured.io: {e}")
return {"success": False, "error": str(e)}
def _format_table_as_markdown(self, table_text: str) -> str:
try:
lines = table_text.strip().split('\n')
if not lines:
return table_text
if '|' in table_text:
rows = [line.strip() for line in lines]
if not rows[0].startswith('|'):
rows[0] = '| ' + rows[0] + ' |'
if len(rows) > 1 and not rows[1].startswith('|---'):
cols = rows[0].count('|') - 1
separator = '|' + '|'.join(['---' for _ in range(cols)]) + '|'
rows.insert(1, separator)
for i in range(len(rows)):
if i > 1 and not rows[i].startswith('|'):
rows[i] = '| ' + rows[i] + ' |'
return '\n'.join(rows)
else:
rows = []
for line in lines:
if line.strip():
rows.append(line.split())
if not rows:
return table_text
num_cols = len(rows[0]) if rows else 0
if num_cols == 0:
return table_text
for i in range(len(rows)):
while len(rows[i]) < num_cols:
rows[i].append('')
rows[i] = rows[i][:num_cols]
header = '| ' + ' | '.join(rows[0]) + ' |'
separator = '|' + '|'.join(['---' for _ in range(num_cols)]) + '|'
formatted_rows = [header, separator]
for row in rows[1:]:
formatted_rows.append('| ' + ' | '.join(row) + ' |')
return '\n'.join(formatted_rows)
except Exception as e:
logger.warning(f"Error formatting table as markdown: {e}")
return f"Table:\n{table_text}"
import re
import uuid
import logging
from typing import List, Dict, Any, Optional, Tuple, Union
import os
from pathlib import Path
import hashlib
from datetime import datetime
import nltk
from nltk.tokenize import sent_tokenize
import json
# from sparse_embedder import GermanBM25SparseEmbedder
from fastembed import SparseTextEmbedding
try:
nltk.data.find('tokenizers/punkt')
except LookupError:
nltk.download('punkt', quiet=True)
def make_serializable(metadata: Dict[str, Any]) -> Dict[str, Any]:
serializable_metadata = {}
for key, value in metadata.items():
try:
json.dumps(value)
serializable_metadata[key] = value
except TypeError:
if isinstance(value, Header):
serializable_metadata[key] = str(value) # Convert Header to string
else:
serializable_metadata[key] = str(value) # Convert other non-serializable objects to string
return serializable_metadata
class DocumentProcessor:
def __init__(self, config, embedding_model):
self.logger = logging.getLogger(__name__)
self.embedding_model = embedding_model
self.chunk_size = config.rag.chunk_size
self.chunk_overlap = config.rag.chunk_overlap
self.processed_docs_dir = Path(config.rag.processed_docs_dir)
self.processed_docs_dir.mkdir(parents=True, exist_ok=True)
self.chunking_strategy = getattr(config.rag, "chunking_strategy", "sliding_window")
# self.sparse_embedder = GermanBM25SparseEmbedder()
self.sparse_embedder = SparseTextEmbedding("Qdrant/bm25")
def process_document(self, content: str, metadata: Dict[str, Any], elements: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
try:
doc_id_base = hashlib.md5(content.encode()).hexdigest()
doc_id = str(uuid.UUID(doc_id_base[:32]))
enhanced_metadata = metadata.copy()
enhanced_metadata['processing_timestamp'] = datetime.now().isoformat()
if self.chunking_strategy == "semantic":
chunks = self._create_semantic_chunks(content, elements)
elif self.chunking_strategy == "sliding_window":
chunks = self._create_sliding_window_chunks(content)
else:
chunks = self._create_sliding_window_chunks(content)
processed_chunks = []
for i, chunk_info in enumerate(chunks):
if isinstance(chunk_info, tuple):
chunk_text, section, level = chunk_info[0], chunk_info[1], chunk_info[2] if len(chunk_info) > 2 else "standard"
else:
chunk_text, section, level = chunk_info, "general", "standard"
# Generate chunk ID as a UUID with a suffix
chunk_id = str(uuid.UUID(doc_id_base[:24] + f"{i:08}"))
print(f"DEBUGG chunk_text: {chunk_text}")
print(f"DEBUGG section: {section}")
# Generate embedding
dense_embedding = self.embedding_model.embed_documents([chunk_text])[0]
sparse_embeddings = list(
self.sparse_embedder.passage_embed([chunk_text])) # Lấy sparse embedding từ passage_embed
# Đảm bảo rằng sparse_embeddings không rỗng và lấy phần tử đầu tiên
if sparse_embeddings:
sparse_embedding = sparse_embeddings[0]
# In ra các chỉ số và giá trị của sparse_embedding
print(f"DEBUGG sparse_embedding indices (first 5): {sparse_embedding.indices[:5]}")
print(f"DEBUGG sparse_embedding values (first 5): {sparse_embedding.values[:5]}")
else:
sparse_embedding = {}
# Print embeddings to debug
print(f"DEBUGG dense_embedding (first 5 elements): {dense_embedding[:5]}") # Print first 5 elements for brevity
chunk_metadata = enhanced_metadata.copy()
chunk_metadata["chunk_number"] = i
chunk_metadata["total_chunks"] = len(chunks)
chunk_metadata["chunking_strategy"] = self.chunking_strategy
chunk_metadata["word_count"] = len(chunk_text.split())
processed_chunks.append({
"id": chunk_id,
"content": chunk_text,
"embedding": {
"dense": dense_embedding,
"sparse": sparse_embedding
},
"metadata": make_serializable(chunk_metadata)
})
return processed_chunks
except Exception as e:
self.logger.error(f"Error processing document: {e}")
raise
def _create_semantic_chunks(self, content: str, elements: List[Dict[str, Any]]) -> List[Tuple[str, str, str]]:
"""
Create semantic chunks from document based on 'elements' with category-based sectioning.
Falls back to paragraph-based chunking if elements are missing.
Returns:
List of (chunk_text, section_name, chunk_type) tuples
"""
chunks = []
current_section = None
section_text = []
encountered_title = False
for element in elements:
category = element.get("category", "")
element_text = element.get('text', '').strip()
# if not element_text:
# continue
if category == 'Title':
if section_text:
if not encountered_title:
chunks.append((' '.join(section_text), "preamble", "section"))
else:
chunks.append((' '.join(section_text), current_section, "section"))
current_section = element_text
section_text = [element_text]
encountered_title = True
else:
section_text.append(element_text)
if section_text:
if not encountered_title:
chunks.append((' '.join(section_text), "preamble", "section"))
else:
chunks.append((' '.join(section_text), current_section, "section"))
if not chunks and content:
paragraphs = re.split(r'\n\s*\n', content)
for para in paragraphs:
if para.strip():
chunks.append((para.strip(), "paragraph", "standard"))
print("Chunks:", chunks)
# TODO: Split section into paragraphs if it's too large
return chunks
def _create_sliding_window_chunks(self, text: str) -> List[str]:
sentences = sent_tokenize(text)
chunks = []
if len(sentences) <= 3:
return [text]
stride = max(1, (self.chunk_size - self.chunk_overlap) // 20)
for i in range(0, len(sentences), stride):
window_size = min(i + max(3, self.chunk_size // 20), len(sentences))
window_text = " ".join(sentences[i:window_size])
chunks.append(window_text)
return chunks
class Header:
def __init__(self, text):
self.text = text
def __str__(self):
return self.text
\ No newline at end of file
from sentence_transformers import SentenceTransformer
class LocalEmbeddingModel:
def __init__(self, model_name='paraphrase-multilingual-MiniLM-L12-v2'):
self.model = SentenceTransformer(model_name)
def embed_documents(self, texts):
return self.model.encode(texts, convert_to_numpy=True)
def embed_query(self, text):
return self.model.encode([text], convert_to_numpy=True)[0]
import logging
from typing import List, Dict, Any, Optional
import re
class ResponseGenerator:
def __init__(self, config, llm):
def __init__(self, config):
self.logger = logging.getLogger(__name__)
self.llm = llm
self.llm = config.rag.llm
self.max_context_length = config.rag.max_context_length
self.include_sources = getattr(config.rag, "include_sources", True)
default_instructions = """
Bei der Formatierung der Antwort:
1. Verwenden Sie bei Bedarf klare Abschnitte mit Überschriften
2. Stellen Sie tabellarische Daten in korrektem Markdown-Tabellenformat dar:
| Kopf1 | Kopf2 | Kopf3 |
|-------|-------|-------|
| Daten1| Daten2| Daten3|
3. Für geordnete Listen verwenden Sie nummerierte Punkte
4. Für ungeordnete Listen verwenden Sie Aufzählungszeichen
5. Zitieren Sie Ihre Quellen, wenn Sie spezifische Informationen bereitstellen
6. Verwenden Sie präzise, wissenschaftliche Sprache
7. Antworten Sie auf Deutsch
def _build_prompt(
self,
query: str,
context: str,
chat_history: Optional[List[Dict[str, str]]] = None
) -> str:
"""
Build the prompt for the language model.
self.response_format_instructions = getattr(config.rag, "response_format_instructions", default_instructions)
self.include_sources = getattr(config.rag, "include_sources", True)
Args:
query: User query
context: Formatted context from retrieved documents
chat_history: Optional chat history
def generate_response(self, query: str, retrieved_docs: List[Dict[str, Any]],
chat_history: Optional[List[Dict[str, str]]] = None) -> Dict[str, Any]:
Returns:
Complete prompt string
"""
try:
# tai cai nay
if not retrieved_docs:
return self._generate_no_documents_response(query)
table_instructions = """
Einige der abgerufenen Informationen werden in Tabellenform präsentiert. Beim Verwenden von Informationen aus Tabellen:
1. Stelle tabellarische Daten mit korrekter Markdown-Tabellenformatierung und Kopfzeilen dar, zum Beispiel:
| Spalte1 | Spalte2 | Spalte3 |
|---------|---------|---------|
| Wert1 | Wert2 | Wert3 |
2. Formatiere die Tabellenstruktur so um, dass sie leichter zu lesen und zu verstehen ist.
3. Wenn beim Umformatieren der Tabelle eine neue Komponente hinzugefügt wird, erwähne dies ausdrücklich.
4. Interpretiere die tabellarischen Daten in deiner Antwort klar und verständlich.
5. Beziehe dich beim Präsentieren spezifischer Datenpunkte auf die jeweilige Tabelle.
6. Fasse, wenn sinnvoll, Trends oder Muster aus den Tabellen zusammen.
7. Falls nur Referenznummern genannt werden und du dazugehörige Werte wie den Titel einer wissenschaftlichen Arbeit oder Autoren aus dem Kontext holen kannst, ersetze die Referenznummern durch die tatsächlichen Werte.
"""
formatted_context = self._format_context(retrieved_docs)
response_format_instructions = """Anweisungen:
1. Beantworte die Frage ausschließlich auf Basis der im Kontext bereitgestellten Informationen.
2. Falls der Kontext keine relevanten Informationen zur Beantwortung der Frage enthält, sage: "Ich habe nicht genügend Informationen, um diese Frage auf Basis des bereitgestellten Kontexts zu beantworten."
3. Nutze kein Vorwissen, das nicht im Kontext enthalten ist.
5. Sei präzise und korrekt.
6. Gib eine gut strukturierte Antwort mit Überschriften, Unterüberschriften und ggf. tabellarischer Darstellung im Markdown-Format basierend auf dem abgerufenen Wissen. Halte Überschriften und Unterüberschriften möglichst kurz.
7. Führe nur solche Abschnitte auf, die in einer Chatbot-Antwort sinnvoll sind. Nenne z. B. keine Referenzen explizit.
8. Wenn Werte enthalten sind, gib exakt die im Kontext vorhandenen Werte wieder. Erfinde keine Werte.
9. Wiederhole die Frage nicht in der Antwort."""
# Build the prompt
prompt = f"""
Du bist ein wissenschaftlicher Chatbot, der folgende Aufgaben erfüllt:
- Zusammenfassung wissenschaftlicher Artikel
- Suche nach aktuellen und vergangenen Forschungsprojekten
- Navigation zu Artikeln: Bereitstellung von Links zu Veröffentlichungen, Anzeige meistzitierter oder neuester Artikel Antworte auf Deutsch.
prompt = self._build_prompt(query, formatted_context, chat_history)
Hier sind die letzten Nachrichten aus unserem Gespräch:
{chat_history}
print(f"DEBUG prompt: \n{prompt}")
Der Nutzer hat folgende Frage gestellt:
{query}
response = self.llm.invoke(prompt)
Ich habe die folgenden Informationen abgerufen, um bei der Beantwortung der Frage zu helfen:
{context}
sources = self._extract_sources(retrieved_docs) if self.include_sources else []
{table_instructions}
result = {
"response": response,
"sources": sources,
"confidence": self._calculate_confidence(retrieved_docs)
}
{response_format_instructions}
return result
except Exception as e:
self.logger.error(f"Error generating response: {e}")
return {
"response": "Entschuldigung, es ist ein Fehler aufgetreten. Bitte versuchen Sie es erneut.",
"sources": [],
"confidence": 0.0
}
Bitte beantworte die Frage des Nutzers auf Grundlage der bereitgestellten Informationen umfassend, aber prägnant.
Falls die Informationen keine Antwort enthalten, weise bitte auf die Grenzen der verfügbaren Informationen hin.
Gib keinen Quell-Link an, der nicht im Kontext enthalten ist. Erfinde keinen Quell-Link.
def _format_context(self, documents: List[Dict[str, Any]]) -> str:
context_parts = []
total_length = 0
Assistant Response:"""
for i, doc in enumerate(documents):
doc_text = doc.get("content", "").strip()
metadata = doc.get("metadata", {})
return prompt
if not metadata and isinstance(doc, dict):
metadata = doc
def generate_response(
self,
query: str,
retrieved_docs: List[Dict[str, Any]],
# picture_paths: List[str],
chat_history: Optional[List[Dict[str, str]]] = None,
) -> Dict[str, Any]:
"""
Generate a response based on retrieved documents.
source = metadata.get("source", "Unknown Source")
Args:
query: User query
retrieved_docs: List of retrieved document dictionaries
chat_history: Optional chat history
is_table = metadata.get("is_table", False)
content_type = metadata.get("content_type", "")
table_index = metadata.get("table_index", None)
Returns:
Dict containing response text and source information
"""
try:
if is_table or content_type == "table":
clean_table_text = self._ensure_markdown_table_format(doc_text)
table_name = f"Table {table_index}" if table_index is not None else "Table"
# Extract content from documents for context
doc_texts = [doc["content"] for doc in retrieved_docs]
formatted_doc = (
f"[Dokument {i + 1} - TABELLE] Folgende Informationen aus {table_name}:\n\n"
f"{clean_table_text}\n\n"
f"(Quelle: {source})"
)
else:
formatted_doc = f"[Dokument {i + 1}] {doc_text} (Quelle: {source})"
if total_length + len(formatted_doc) > self.max_context_length:
break
context_parts.append(formatted_doc)
total_length += len(formatted_doc)
return "\n\n".join(context_parts)
def _ensure_markdown_table_format(self, table_text: str) -> str:
if "|" in table_text:
lines = table_text.strip().split("\n")
has_separator = False
for i, line in enumerate(lines):
if i > 0 and re.match(r"^\s*\|[\s\-\|]+\|\s*$", line):
has_separator = True
break
if not has_separator and len(lines) > 1:
first_row = lines[0]
col_count = first_row.count("|") - 1
if col_count <= 0:
col_count = first_row.count("\t") + 1 # Try tab count
separator = "|" + "|".join(["---"] * col_count) + "|"
lines.insert(1, separator)
for i in range(len(lines)):
if not lines[i].startswith("|"):
lines[i] = "| " + lines[i]
if not lines[i].endswith("|"):
lines[i] = lines[i] + " |"
return "\n".join(lines)
return table_text
lines = table_text.strip().split("\n")
if not lines:
return table_text
rows = []
for line in lines:
if line.strip():
if "\t" in line:
rows.append(line.split("\t"))
else:
rows.append(re.split(r"\s{2,}", line.strip()))
if not rows:
return table_text
max_cols = max(len(row) for row in rows)
for row in rows:
while len(row) < max_cols:
row.append("")
md_table_lines = []
md_table_lines.append("| " + " | ".join(rows[0]) + " |")
md_table_lines.append("| " + " | ".join(["---"] * max_cols) + " |")
for row in rows[1:]:
md_table_lines.append("| " + " | ".join(row) + " |")
return "\n".join(md_table_lines)
def _build_prompt(self, query: str, context: str,
chat_history: Optional[List[Dict[str, str]]] = None) -> str:
project_instructions = """
Du bist ein wissenschaftlicher Chatbot, der folgende Aufgaben erfüllt:
- Zusammenfassung wissenschaftlicher Artikel
- Suche nach aktuellen und vergangenen Forschungsprojekten
- Navigation zu Artikeln: Bereitstellung von Links zu Veröffentlichungen, Anzeige meistzitierter oder neuester Artikel
Antworte auf Deutsch.
"""
table_instructions = ""
if "TABELLE" in context:
table_instructions = """
Einige Informationen liegen in Tabellenformat vor. Verwende:
1. Markdown-Tabellen mit Kopfzeilen
2. Erhalten Sie die Struktur bei der Darstellung
3. Interpretieren und referenzieren Sie die Tabellen deutlich
"""
chat_str = ""
if chat_history and isinstance(chat_history, list):
chat_lines = []
for msg in chat_history:
if isinstance(msg, dict):
role = "User" if msg.get("type") == "human" else "Assistant"
content = msg.get("content", "")
chat_lines.append(f"{role}: {content}")
else:
chat_lines.append(str(msg))
chat_str = "\n".join(chat_lines)
# Combine retrieved documents into a single context
context = "\n\n===DOCUMENT SECTION===\n\n".join(doc_texts)
print(f"chat_str: {chat_str}")
# Build the prompt
prompt = self._build_prompt(query, context, chat_history)
prompt = f"""
{project_instructions}
# Generate response
response = self.llm.invoke(prompt)
Bisherige Konversation:
{chat_str}
# Extract sources for citation
sources = self._extract_sources(retrieved_docs) if hasattr(self,
'include_sources') and self.include_sources else []
Benutzeranfrage:
{query}
# Calculate confidence
confidence = self._calculate_confidence(retrieved_docs)
Verfügbare Informationen:
{context}
# Add sources to response
if hasattr(self, 'include_sources') and self.include_sources:
response_with_source = response.content + "\n\n##### Source documents:"
for current_source in sources:
source_path = current_source['path']
source_title = current_source['title']
response_with_source += f"\n- [{source_title}]({source_path})"
else:
response_with_source = response.content
{table_instructions}
# # Add picture paths to response
# response_with_source_and_picture_paths = response_with_source + "\n\n##### Reference images:"
# for picture_path in picture_paths:
# response_with_source_and_picture_paths += f"\n- [{picture_path.split('/')[-1]}]({picture_path})"
{self.response_format_instructions}
# Format final response
result = {
# "response": response_with_source_and_picture_paths,
"response": response_with_source,
"sources": sources,
"confidence": confidence
}
Wissenschaftliche Chatbot-Antwort:"""
return result
return prompt
except Exception as e:
self.logger.error(f"Error generating response: {e}")
return {
"response": "I apologize, but I encountered an error while generating a response. Please try rephrasing your question.",
"sources": [],
"confidence": 0.0
}
def _extract_sources(self, documents: List[Dict[str, Any]]) -> List[Dict[str, str]]:
"""
Extract source information from retrieved documents for citation.
Args:
documents: List of retrieved document dictionaries
Returns:
List of source information dictionaries
"""
sources = []
seen_sources = set()
seen_sources = set() # Track unique sources to avoid duplicates
for doc in documents:
metadata = doc.get("metadata", {})
if not metadata and isinstance(doc, dict):
metadata = doc
source = metadata.get("source", "Unknown Source")
# Extract source and source_path
source = doc.get("source")
source_path = doc.get("source_path")
is_table = metadata.get("is_table", False)
content_type = metadata.get("content_type", "")
table_index = metadata.get("table_index", None)
# Skip if no source information is available
if not source:
continue
source_key = f"{source}_{is_table}_{table_index}" if (is_table or content_type == "table") else source
# Create a unique identifier for this source
source_id = f"{source}|{source_path}"
if source_key in seen_sources:
# Skip if we've already included this source
if source_id in seen_sources:
continue
# Add to our sources list
source_info = {
"title": source,
"section": metadata.get("section", ""),
"publication_date": metadata.get("publication_date", "")
"path": source_path,
"score": doc.get("combined_score", doc.get("rerank_score", doc.get("score", 0.0)))
}
if is_table or content_type == "table":
table_descriptor = f"Table {table_index}" if table_index is not None else "Table"
source_info["title"] = f"{source} ({table_descriptor})"
source_info["content_type"] = "table"
sources.append(source_info)
seen_sources.add(source_key)
if len(sources) >= 5:
break
return sources
seen_sources.add(source_id)
def _generate_no_documents_response(self, query: str) -> Dict[str, Any]:
prompt = f"""
Du bist ein wissenschaftlicher Chatbot. Der Benutzer fragt:
\"{query}\"
Leider liegen keine relevanten Dokumente vor. Bitte gib eine allgemeine Antwort und schlage gegebenenfalls Suchstrategien vor.
Wissenschaftliche Chatbot-Antwort:"""
# Sort sources by score from highest to lowest
sources.sort(key=lambda x: x.get("score", 0), reverse=True)
response = self.llm.invoke(prompt)
# Format the final sources list, removing the scores which were just used for sorting
formatted_sources = []
for source in sources:
formatted_source = {
"title": source["title"],
"path": source["path"]
}
formatted_sources.append(formatted_source)
return {
"response": response,
"sources": [],
"confidence": 0.0
}
return formatted_sources
def _calculate_confidence(self, documents: List[Dict[str, Any]]) -> float:
"""
Calculate confidence score based on retrieved documents.
Args:
documents: Retrieved documents
Returns:
Confidence score between 0 and 1
"""
if not documents:
return 0.0
# Use combined score (both reranker and cosine similarity) if available, otherwise use original score
if "combined_score" in documents[0]:
scores = [doc.get("combined_score", 0) for doc in documents[:3]]
elif "rerank_score" in documents[0]:
scores = [doc.get("rerank_score", 0) for doc in documents[:3]]
else:
scores = [doc.get("score", 0) for doc in documents[:3]]
return sum(scores) / len(scores) if scores else 0.0
# Average of top 3 document scores or fewer if less than 3
return sum(scores) / len(scores) if scores else 0.0
\ No newline at end of file
This diff is collapsed.
......@@ -4,6 +4,8 @@ from fastembed import SparseTextEmbedding
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from openai import OpenAI
from agents.rag_agent.local_embedding_model import LocalEmbeddingModel
load_dotenv()
class AgentDecisoinConfig:
......@@ -21,41 +23,51 @@ class RAGConfig:
self.distance_metric = "Cosine"
self.use_local = True
self.vector_local_path = "./data/qdrant_db" # Add this with a default value
self.doc_local_path = "./data/docs_db"
self.parsed_content_dir = "./data/parsed_docs"
self.url = os.getenv("QDRANT_URL")
self.api_key = os.getenv("QDRANT_API_KEY")
self.collection_name = "research_assistant_rag"
self.collection_name = "pro_assistant_rag"
self.chunk_size = 512
self.chunk_overlap = 50
self.processed_docs_dir = "./data/processed"
self.embedding_model = OpenAIEmbeddings(
model=os.getenv("embedding_model_name"),
api_key=os.getenv("openai_api_key")
)
self.sparse_embedder = SparseTextEmbedding("Qdrant/bm25")
self.embedding_for_chunk_model = LocalEmbeddingModel('paraphrase-multilingual-MiniLM-L12-v2')
self.summarizer_model = ChatOpenAI(
model=os.getenv("model_name"),
api_key=os.getenv("openai_api_key"),
temperature=0.5
)
self.chunker_model = ChatOpenAI(
model=os.getenv("model_name"),
api_key=os.getenv("openai_api_key"),
temperature=0.0
)
self.llm = ChatOpenAI(
model=os.getenv("model_name"),
api_key=os.getenv("openai_api_key"),
temperature=0.3
)
self.sparse_embedder = SparseTextEmbedding("Qdrant/bm25")
self.max_queries = 3
self.top_k = 5
self.similarity_threshold = 0.75
self.huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
self.vector_search_type = 'similarity' # or 'mmr'
self.chunking_strategy = "semantic"
self.reranker_model = "cross-encoder/ms-marco-TinyBERT-L-6"
self.reranker_top_k = 3
# self.chunking_strategy = "semantic"
self.max_context_length = 8192
self.include_sources = True
self.response_format_instructions = """Instructions:
1. Answer the query based ONLY on the information provided in the context.
2. If the context doesn't contain relevant information to answer the query, state: "I don't have enough information to answer this question based on the provided context."
3. Do not use prior knowledge not contained in the context.
5. Be concise and accurate.
6. Provide a well-structured response based on retrieved knowledge."""
self.include_sources = True
self.metrics_save_path = "./logs/rag_metrics.json"
self.min_retrieval_confidence = 0.8 # the auto routing from RAG agent to WEB_SEARCH agent is dependent on this value
self.context_limit = 20 # include last 20 messsages (10 Q&A pairs) in history
......
This diff is collapsed.
fastapi
fastapi~=0.115.12
uvicorn==0.34.0
dotenv
dotenv~=0.9.9
qdrant-client==1.13.3
nltk==3.9.1
pandas==2.2.3
......@@ -15,7 +15,7 @@ langchain-core==0.3.44
langchain-openai==0.3.8
langchain-text-splitters==0.3.6
langchain-qdrant
langgraph
langgraph~=0.3.9
langgraph-checkpoint
langgraph-prebuilt
langgraph-sdk
......@@ -24,7 +24,13 @@ pikepdf==9.5.2
pillow==11.1.0
pipreqs==0.5.0
pdf2image==1.17.0
scikit-learn
streamlit
requests
fastembed
\ No newline at end of file
scikit-learn~=1.6.1
streamlit~=1.45.0
requests~=2.32.3
fastembed~=0.6.1
sentence-transformers~=4.1.0
pydantic~=2.11.3
numpy~=1.26.4
docling~=2.31.0
python-dotenv~=1.1.0
openai~=1.76.0
\ No newline at end of file
This diff is collapsed.
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment