460 lines
18 KiB
Python
460 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
import os
|
|
import re
|
|
import json
|
|
import ssl
|
|
import argparse
|
|
import requests
|
|
from bs4 import BeautifulSoup
|
|
from urllib.parse import quote
|
|
from langchain_community.embeddings import HuggingFaceEmbeddings
|
|
from langchain_community.vectorstores import Chroma
|
|
from langchain_core.documents import Document
|
|
import traceback
|
|
|
|
# Disable SSL warnings and proxy settings
|
|
ssl._create_default_https_context = ssl._create_unverified_context
|
|
requests.packages.urllib3.disable_warnings()
|
|
|
|
def clear_proxy_settings():
|
|
"""Remove proxy environment variables that might cause connection issues."""
|
|
for var in ["HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy"]:
|
|
if var in os.environ:
|
|
print(f"Removing proxy env var: {var}")
|
|
del os.environ[var]
|
|
|
|
# Run at module load time
|
|
clear_proxy_settings()
|
|
|
|
# Configuration
|
|
DOCUMENT_PATHS = [
|
|
r'doc1.txt',
|
|
r'doc2.txt',
|
|
r'doc3.txt',
|
|
r'doc4.txt',
|
|
r'doc5.txt',
|
|
r'doc6.txt'
|
|
]
|
|
EMBEDDING_MODEL = 'sentence-transformers/paraphrase-multilingual-mpnet-base-v2'
|
|
LLM_MODEL = 'gemma3'
|
|
CHUNK_SIZE = 1000
|
|
OVERLAP = 200
|
|
CHROMA_PERSIST_DIR = 'chroma_db'
|
|
|
|
# Confidence thresholds
|
|
THRESHOLDS = {
|
|
'direct_answer': 0.7,
|
|
'rag_confidence': 0.6,
|
|
'web_search': 0.5
|
|
}
|
|
|
|
def query_llm(prompt, model='gemma3'):
|
|
"""Query the LLM model directly using Ollama API."""
|
|
try:
|
|
ollama_endpoint = "http://localhost:11434/api/generate"
|
|
payload = {
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"stream": False
|
|
}
|
|
response = requests.post(ollama_endpoint, json=payload)
|
|
|
|
if response.status_code == 200:
|
|
result = response.json()
|
|
return result.get('response', '')
|
|
else:
|
|
print(f"Ollama API error: {response.status_code}")
|
|
return f"Error calling Ollama API: {response.status_code}"
|
|
except Exception as e:
|
|
print(f"Error querying LLM: {e}")
|
|
return f"Error: {str(e)}"
|
|
|
|
class BM25Retriever:
|
|
"""BM25 retriever implementation for text similarity search"""
|
|
|
|
@classmethod
|
|
def from_documents(cls, documents):
|
|
"""Create a BM25 retriever from documents"""
|
|
retriever = cls()
|
|
retriever.documents = documents
|
|
retriever.k = 4
|
|
return retriever
|
|
|
|
def get_relevant_documents(self, query):
|
|
"""Get relevant documents using BM25 algorithm"""
|
|
# Simple BM25-like implementation
|
|
scores = []
|
|
query_terms = set(re.findall(r'\b\w+\b', query.lower()))
|
|
|
|
for doc in self.documents:
|
|
doc_terms = set(re.findall(r'\b\w+\b', doc.page_content.lower()))
|
|
# Calculate term overlap as a simple approximation of BM25
|
|
overlap = len(query_terms.intersection(doc_terms))
|
|
scores.append((doc, overlap))
|
|
|
|
# Sort by score and return top k
|
|
sorted_docs = [doc for doc, score in sorted(scores, key=lambda x: x[1], reverse=True)]
|
|
return sorted_docs[:self.k]
|
|
|
|
class HybridRetriever:
|
|
"""Hybrid retriever combining BM25 and vector search with configurable weights"""
|
|
|
|
def __init__(self, vector_retriever, bm25_retriever, vector_weight=0.3):
|
|
"""Initialize with separate retrievers and weights"""
|
|
self._vector_retriever = vector_retriever
|
|
self._bm25_retriever = bm25_retriever
|
|
self._vector_weight = vector_weight
|
|
self._bm25_weight = 1.0 - vector_weight
|
|
|
|
def get_relevant_documents(self, query):
|
|
"""Get relevant documents using weighted combination of retrievers"""
|
|
try:
|
|
# Get results from both retrievers
|
|
vector_docs = self._vector_retriever.get_relevant_documents(query)
|
|
bm25_docs = self._bm25_retriever.get_relevant_documents(query)
|
|
|
|
# Create dictionary to track unique documents and their scores
|
|
doc_dict = {}
|
|
|
|
# Add vector docs with their weights
|
|
for i, doc in enumerate(vector_docs):
|
|
# Score based on position (inverse rank)
|
|
score = (len(vector_docs) - i) * self._vector_weight
|
|
doc_id = doc.page_content[:50] # Use first 50 chars as a simple ID
|
|
if doc_id in doc_dict:
|
|
doc_dict[doc_id]["score"] += score
|
|
else:
|
|
doc_dict[doc_id] = {"doc": doc, "score": score}
|
|
|
|
# Add BM25 docs with their weights
|
|
for i, doc in enumerate(bm25_docs):
|
|
# Score based on position (inverse rank)
|
|
score = (len(bm25_docs) - i) * self._bm25_weight
|
|
doc_id = doc.page_content[:50] # Use first 50 chars as a simple ID
|
|
if doc_id in doc_dict:
|
|
doc_dict[doc_id]["score"] += score
|
|
else:
|
|
doc_dict[doc_id] = {"doc": doc, "score": score}
|
|
|
|
# Sort by combined score (highest first)
|
|
sorted_docs = sorted(doc_dict.values(), key=lambda x: x["score"], reverse=True)
|
|
|
|
# Return just the document objects
|
|
return [item["doc"] for item in sorted_docs]
|
|
except Exception as e:
|
|
print(f"Error in hybrid retrieval: {e}")
|
|
return []
|
|
|
|
class AgenticQASystem:
|
|
"""QA system implementing the specified architecture"""
|
|
|
|
def __init__(self):
|
|
"""Initialize the QA system with retrievers"""
|
|
# Load embeddings
|
|
self.embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
|
|
# Load documents and retrievers
|
|
self.documents = self.load_documents()
|
|
self.retriever = self.initialize_retriever()
|
|
|
|
def load_documents(self):
|
|
"""Load documents from configured paths with sliding window chunking"""
|
|
print("Loading documents...")
|
|
docs = []
|
|
for path in DOCUMENT_PATHS:
|
|
try:
|
|
with open(path, 'r', encoding='utf-8') as f:
|
|
text = re.sub(r'\s+', ' ', f.read()).strip()
|
|
# Sliding window chunking
|
|
chunks = [text[i:i+CHUNK_SIZE] for i in range(0, len(text), CHUNK_SIZE - OVERLAP)]
|
|
for chunk in chunks:
|
|
docs.append(Document(
|
|
page_content=chunk,
|
|
metadata={"source": os.path.basename(path)}
|
|
))
|
|
except Exception as e:
|
|
print(f"Error loading document {path}: {e}")
|
|
print(f"Loaded {len(docs)} document chunks")
|
|
return docs
|
|
|
|
def initialize_retriever(self):
|
|
"""Initialize the hybrid retriever with BM25 and direct Chroma queries"""
|
|
if not self.documents:
|
|
print("No documents loaded, retriever initialization failed")
|
|
return None
|
|
|
|
try:
|
|
# Create BM25 retriever
|
|
bm25_retriever = BM25Retriever.from_documents(self.documents)
|
|
bm25_retriever.k = 4 # Top k results to retrieve
|
|
|
|
# Initialize vector store with KNN search
|
|
import shutil
|
|
if os.path.exists(CHROMA_PERSIST_DIR):
|
|
print(f"Removing existing Chroma DB to prevent dimension mismatch")
|
|
shutil.rmtree(CHROMA_PERSIST_DIR)
|
|
|
|
# Create vector store directly from Chroma
|
|
print("Creating vector store...")
|
|
vector_store = Chroma.from_documents(
|
|
documents=self.documents,
|
|
embedding=self.embeddings,
|
|
persist_directory=CHROMA_PERSIST_DIR
|
|
)
|
|
|
|
vector_retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 4})
|
|
print(f"Vector retriever created: {type(vector_retriever)}")
|
|
|
|
# Create hybrid retriever - BM25 (70%) and Vector (30%)
|
|
print("Creating hybrid retriever")
|
|
hybrid_retriever = HybridRetriever(vector_retriever, bm25_retriever, vector_weight=0.3)
|
|
print("Hybrid retriever initialized successfully")
|
|
return hybrid_retriever
|
|
|
|
except Exception as e:
|
|
print(f"Error initializing retriever: {e}")
|
|
traceback.print_exc()
|
|
return None
|
|
|
|
def estimate_confidence(self, text, query, context=None):
|
|
"""Estimate confidence of response"""
|
|
# Start with baseline confidence
|
|
confidence = 0.5
|
|
|
|
# Check for uncertainty markers
|
|
uncertainty_phrases = [
|
|
"نمیدانم", "مطمئن نیستم", "ممکن است", "شاید", "احتمالاً",
|
|
"فکر میکنم", "به نظر میرسد"
|
|
]
|
|
|
|
if any(phrase in text.lower() for phrase in uncertainty_phrases):
|
|
confidence -= 0.2
|
|
|
|
# Check for question relevance
|
|
query_words = set(re.findall(r'\b\w+\b', query.lower()))
|
|
text_words = set(re.findall(r'\b\w+\b', text.lower()))
|
|
|
|
# Calculate overlap between query and response
|
|
if query_words:
|
|
overlap_ratio = len(query_words.intersection(text_words)) / len(query_words)
|
|
if overlap_ratio > 0.5:
|
|
confidence += 0.2
|
|
elif overlap_ratio < 0.2:
|
|
confidence -= 0.2
|
|
|
|
# If context provided, check context relevance
|
|
if context:
|
|
context_words = set(re.findall(r'\b\w+\b', context.lower()))
|
|
if context_words:
|
|
context_overlap = len(context_words.intersection(text_words)) / len(context_words)
|
|
if context_overlap > 0.3:
|
|
confidence += 0.2
|
|
else:
|
|
confidence -= 0.1
|
|
|
|
# Ensure confidence is within [0,1]
|
|
return max(0.0, min(1.0, confidence))
|
|
|
|
def check_direct_knowledge(self, query):
|
|
"""Check if the LLM can answer directly from its knowledge"""
|
|
print("Checking LLM's direct knowledge...")
|
|
prompt = f"""به این سوال با استفاده از دانش خود پاسخ دهید. فقط به زبان فارسی پاسخ دهید.
|
|
|
|
سوال: {query}
|
|
|
|
پاسخ فارسی:"""
|
|
|
|
response = query_llm(prompt, model=LLM_MODEL)
|
|
confidence = self.estimate_confidence(response, query)
|
|
print(f"LLM direct knowledge confidence: {confidence:.2f}")
|
|
|
|
return response, confidence
|
|
|
|
def rag_query(self, query):
|
|
"""Use RAG to retrieve and generate answer"""
|
|
if not self.retriever:
|
|
print("Retriever not initialized, skipping RAG")
|
|
return None, 0.0
|
|
|
|
print("Retrieving documents for RAG...")
|
|
# Retrieve relevant documents
|
|
docs = self.retriever.get_relevant_documents(query)
|
|
|
|
if not docs:
|
|
print("No relevant documents found")
|
|
return None, 0.0
|
|
|
|
print(f"Retrieved {len(docs)} relevant documents")
|
|
|
|
# Prepare context
|
|
context = "\n\n".join([doc.page_content for doc in docs])
|
|
sources = [doc.metadata.get("source", "Unknown") for doc in docs]
|
|
|
|
# Query LLM with context
|
|
prompt = f"""با توجه به اطلاعات زیر، به سوال پاسخ دهید. فقط به زبان فارسی پاسخ دهید.
|
|
|
|
اطلاعات:
|
|
{context}
|
|
|
|
سوال: {query}
|
|
|
|
پاسخ فارسی:"""
|
|
|
|
response = query_llm(prompt, model=LLM_MODEL)
|
|
confidence = self.estimate_confidence(response, query, context)
|
|
print(f"RAG confidence: {confidence:.2f}")
|
|
|
|
return {
|
|
"response": response,
|
|
"confidence": confidence,
|
|
"sources": list(set(sources))
|
|
}, confidence
|
|
|
|
def web_search(self, query):
|
|
"""Search the web for an answer"""
|
|
print("Searching web for answer...")
|
|
# Search DuckDuckGo
|
|
search_url = f"https://html.duckduckgo.com/html/?q={quote(query)}"
|
|
response = requests.get(search_url, verify=False, timeout=10)
|
|
|
|
if response.status_code != 200:
|
|
print(f"Error searching web: HTTP {response.status_code}")
|
|
return None, 0.0
|
|
|
|
# Parse results
|
|
soup = BeautifulSoup(response.text, 'html.parser')
|
|
results = []
|
|
|
|
for element in soup.select('.result__url, .result__a')[:4]:
|
|
href = element.get('href') if 'href' in element.attrs else None
|
|
|
|
if href and not href.startswith('/') and (href.startswith('http://') or href.startswith('https://')):
|
|
results.append(href)
|
|
elif not href and element.find('a') and 'href' in element.find('a').attrs:
|
|
href = element.find('a')['href']
|
|
if href and not href.startswith('/'):
|
|
results.append(href)
|
|
|
|
if not results:
|
|
print("No web results found")
|
|
return None, 0.0
|
|
|
|
# Crawl top results
|
|
web_content = []
|
|
for url in results[:3]:
|
|
try:
|
|
headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"}
|
|
page = requests.get(url, headers=headers, timeout=10, verify=False)
|
|
page.raise_for_status()
|
|
|
|
soup = BeautifulSoup(page.text, 'html.parser')
|
|
|
|
# Remove non-content elements
|
|
for tag in ['script', 'style', 'nav', 'footer', 'header']:
|
|
for element in soup.find_all(tag):
|
|
element.decompose()
|
|
|
|
# Get paragraphs
|
|
paragraphs = [p.get_text(strip=True) for p in soup.find_all('p')
|
|
if len(p.get_text(strip=True)) > 20]
|
|
|
|
if paragraphs:
|
|
web_content.append(f"[Source: {url}] " + " ".join(paragraphs[:5]))
|
|
except Exception as e:
|
|
print(f"Error crawling {url}: {e}")
|
|
|
|
if not web_content:
|
|
print("No useful content found from web results")
|
|
return None, 0.0
|
|
|
|
# Query LLM with web content
|
|
context = "\n\n".join(web_content)
|
|
prompt = f"""با توجه به اطلاعات زیر که از وب بدست آمده، به سوال پاسخ دهید. فقط به زبان فارسی پاسخ دهید.
|
|
|
|
اطلاعات:
|
|
{context}
|
|
|
|
سوال: {query}
|
|
|
|
پاسخ فارسی:"""
|
|
|
|
response = query_llm(prompt, model=LLM_MODEL)
|
|
confidence = self.estimate_confidence(response, query, context)
|
|
print(f"Web search confidence: {confidence:.2f}")
|
|
|
|
return {
|
|
"response": response,
|
|
"confidence": confidence,
|
|
"sources": results[:3]
|
|
}, confidence
|
|
|
|
def get_answer(self, query):
|
|
"""Main method to get an answer following the specified architecture"""
|
|
print(f"Processing query: {query}")
|
|
|
|
# STEP 1: Try direct LLM knowledge
|
|
direct_response, direct_confidence = self.check_direct_knowledge(query)
|
|
|
|
if direct_confidence >= THRESHOLDS['direct_answer']:
|
|
print("Using direct LLM knowledge (high confidence)")
|
|
return f"{direct_response}\n\n[Source: LLM Knowledge, Confidence: {direct_confidence:.2f}]"
|
|
|
|
# STEP 2: Try RAG with local documents
|
|
rag_result, rag_confidence = self.rag_query(query)
|
|
|
|
if rag_result and rag_confidence >= THRESHOLDS['rag_confidence']:
|
|
print("Using RAG response (sufficient confidence)")
|
|
sources_text = ", ".join(rag_result["sources"][:3])
|
|
return f"{rag_result['response']}\n\n[Source: Local Documents, Confidence: {rag_confidence:.2f}, Sources: {sources_text}]"
|
|
|
|
# STEP 3: Try web search
|
|
web_result, web_confidence = self.web_search(query)
|
|
|
|
if web_result and web_confidence >= THRESHOLDS['web_search']:
|
|
print("Using web search response (sufficient confidence)")
|
|
sources_text = ", ".join(web_result["sources"])
|
|
return f"{web_result['response']}\n\n[Source: Web Search, Confidence: {web_confidence:.2f}, Sources: {sources_text}]"
|
|
|
|
# STEP 4: Fall back to direct response with warning
|
|
print("No high-confidence source found, using direct response with warning")
|
|
return f"{direct_response}\n\n[Warning: Low confidence ({direct_confidence:.2f}). Please verify information.]"
|
|
|
|
# Simple API functions
|
|
def get_answer(query):
|
|
"""Get an answer for a query"""
|
|
system = AgenticQASystem()
|
|
return system.get_answer(query)
|
|
|
|
# Main entry point
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="QA System")
|
|
|
|
mode_group = parser.add_mutually_exclusive_group(required=True)
|
|
mode_group.add_argument("--query", "-q", help="Query to answer")
|
|
mode_group.add_argument("--interactive", "-i", action="store_true", help="Run in interactive chat mode")
|
|
mode_group.add_argument("--test", "-t", action="store_true", help="Run tests")
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.interactive:
|
|
# Simple interactive mode without memory
|
|
qa_system = AgenticQASystem()
|
|
print("=== QA System ===")
|
|
print("Type 'exit' or 'quit' to end")
|
|
|
|
while True:
|
|
user_input = input("\nYou: ")
|
|
if not user_input.strip():
|
|
continue
|
|
|
|
if user_input.lower() in ['exit', 'quit', 'خروج']:
|
|
break
|
|
|
|
response = qa_system.get_answer(user_input)
|
|
print(f"\nBot: {response}")
|
|
elif args.query:
|
|
qa_system = AgenticQASystem()
|
|
print(qa_system.get_answer(args.query))
|
|
elif args.test:
|
|
print("Running tests...") |