document processing and response gen smooth
This commit is contained in:
parent
c9ea9655bd
commit
5193d040f1
1 changed files with 139 additions and 96 deletions
|
|
@ -47,13 +47,13 @@ def process_document(file_path: str) -> None:
|
||||||
try:
|
try:
|
||||||
import time
|
import time
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
|
|
||||||
# Insert document into PostgreSQL database
|
# Insert document into PostgreSQL database
|
||||||
insert_document(Path(file_path), config=my_config)
|
insert_document(Path(file_path), config=my_config)
|
||||||
|
|
||||||
processing_time = time.time() - start_time
|
processing_time = time.time() - start_time
|
||||||
logger.info(f"Document processed and embedded in {processing_time:.2f} seconds")
|
logger.info(f"Document processed and embedded in {processing_time:.2f} seconds")
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error processing document: {str(e)}")
|
logger.error(f"Error processing document: {str(e)}")
|
||||||
raise
|
raise
|
||||||
|
|
@ -65,15 +65,15 @@ def perform_search(query: str) -> List[dict]:
|
||||||
# First try hybrid search in the database
|
# First try hybrid search in the database
|
||||||
chunk_ids, scores = hybrid_search(query, num_results=10, config=my_config)
|
chunk_ids, scores = hybrid_search(query, num_results=10, config=my_config)
|
||||||
logger.debug(f"Found {len(chunk_ids)} chunks with scores: {scores}")
|
logger.debug(f"Found {len(chunk_ids)} chunks with scores: {scores}")
|
||||||
|
|
||||||
if not chunk_ids:
|
if not chunk_ids:
|
||||||
logger.info("No relevant chunks found in database")
|
logger.info("No relevant chunks found in database")
|
||||||
return []
|
return []
|
||||||
|
|
||||||
# Retrieve and rerank chunks
|
# Retrieve and rerank chunks
|
||||||
chunks = retrieve_chunks(chunk_ids, config=my_config)
|
chunks = retrieve_chunks(chunk_ids, config=my_config)
|
||||||
reranked_chunks = rerank_chunks(query, chunks, config=my_config)
|
reranked_chunks = rerank_chunks(query, chunks, config=my_config)
|
||||||
|
|
||||||
return reranked_chunks
|
return reranked_chunks
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Search error: {str(e)}")
|
logger.error(f"Search error: {str(e)}")
|
||||||
|
|
@ -83,15 +83,13 @@ def perform_search(query: str) -> List[dict]:
|
||||||
async def handle_settings_update(settings: dict):
|
async def handle_settings_update(settings: dict):
|
||||||
"""Handle settings updates when user submits the form."""
|
"""Handle settings updates when user submits the form."""
|
||||||
try:
|
try:
|
||||||
# Validate API keys
|
|
||||||
def validate_key(key: str, key_type: str, valid_prefixes: tuple) -> bool:
|
def validate_key(key: str, key_type: str, valid_prefixes: tuple) -> bool:
|
||||||
if not key:
|
if not key:
|
||||||
raise ValueError(f"{key_type} API key is required")
|
raise ValueError(f"{key_type} API key is required")
|
||||||
if valid_prefixes and not any(key.startswith(prefix) for prefix in valid_prefixes):
|
if valid_prefixes and not any(key.startswith(prefix) for prefix in valid_prefixes):
|
||||||
raise ValueError(f"Invalid {key_type} API key format")
|
raise ValueError(f"Invalid {key_type} API key format")
|
||||||
return True
|
return True
|
||||||
|
# Validate DB URL
|
||||||
# Validate DB URL
|
|
||||||
def validate_db_url(url: str) -> bool:
|
def validate_db_url(url: str) -> bool:
|
||||||
valid_prefixes = ('postgresql://', 'mysql://', 'sqlite:///')
|
valid_prefixes = ('postgresql://', 'mysql://', 'sqlite:///')
|
||||||
if not url:
|
if not url:
|
||||||
|
|
@ -100,12 +98,10 @@ async def handle_settings_update(settings: dict):
|
||||||
raise ValueError("Invalid database URL format")
|
raise ValueError("Invalid database URL format")
|
||||||
return True
|
return True
|
||||||
|
|
||||||
# Validate all inputs
|
|
||||||
validate_key(settings["OpenAIApiKey"], "OpenAI", ("sk-", "sk-proj-"))
|
validate_key(settings["OpenAIApiKey"], "OpenAI", ("sk-", "sk-proj-"))
|
||||||
validate_key(settings["AnthropicApiKey"], "Anthropic", ("sk-ant-",))
|
validate_key(settings["AnthropicApiKey"], "Anthropic", ("sk-ant-",))
|
||||||
validate_key(settings["CohereApiKey"], "Cohere", tuple())
|
validate_key(settings["CohereApiKey"], "Cohere", tuple())
|
||||||
validate_db_url(settings["DBUrl"])
|
validate_db_url(settings["DBUrl"])
|
||||||
|
|
||||||
# Store validated values in user_env
|
# Store validated values in user_env
|
||||||
user_env = {
|
user_env = {
|
||||||
"OPENAI_API_KEY": settings["OpenAIApiKey"],
|
"OPENAI_API_KEY": settings["OpenAIApiKey"],
|
||||||
|
|
@ -114,23 +110,65 @@ async def handle_settings_update(settings: dict):
|
||||||
"DB_URL": settings["DBUrl"]
|
"DB_URL": settings["DBUrl"]
|
||||||
}
|
}
|
||||||
|
|
||||||
# Store in user session
|
|
||||||
cl.user_session.set("env", user_env)
|
|
||||||
|
|
||||||
# Initialize RAGLite config
|
|
||||||
global my_config
|
global my_config
|
||||||
my_config = initialize_config(user_env)
|
my_config = initialize_config(user_env)
|
||||||
|
cl.user_session.set("env", user_env)
|
||||||
|
|
||||||
await cl.Message(content="✅ Successfully configured with your API keys!").send()
|
await cl.Message(content="✅ Successfully configured with your API keys!").send()
|
||||||
|
|
||||||
# Automatically prompt for PDF upload
|
# Ask for file upload with proper configuration
|
||||||
await cl.AskFileMessage(
|
files = await cl.AskFileMessage(
|
||||||
content="Please upload one or more PDF documents to begin!",
|
content="Please upload one or more PDF documents to begin!",
|
||||||
accept=["application/pdf"],
|
accept=["application/pdf"],
|
||||||
max_size_mb=20,
|
max_size_mb=20,
|
||||||
|
timeout=300,
|
||||||
max_files=5
|
max_files=5
|
||||||
).send()
|
).send()
|
||||||
|
|
||||||
|
if files:
|
||||||
|
success = False
|
||||||
|
|
||||||
|
# Process uploaded files
|
||||||
|
for file in files:
|
||||||
|
logger.info(f"Starting to process file: {file.name}")
|
||||||
|
|
||||||
|
# Create new message for each file
|
||||||
|
await cl.Message(f"Processing {file.name}...").send()
|
||||||
|
|
||||||
|
try:
|
||||||
|
logger.info(f"Embedding document: {file.path}")
|
||||||
|
process_document(file_path=file.path)
|
||||||
|
|
||||||
|
success = True
|
||||||
|
await cl.Message(f"✅ Successfully processed: {file.name}").send()
|
||||||
|
logger.info(f"Successfully processed and embedded: {file.name}")
|
||||||
|
|
||||||
|
except Exception as proc_error:
|
||||||
|
error_msg = f"Failed to process {file.name}: {str(proc_error)}"
|
||||||
|
logger.error(error_msg)
|
||||||
|
await cl.Message(f"❌ {error_msg}").send()
|
||||||
|
continue
|
||||||
|
|
||||||
|
if success:
|
||||||
|
# Send completion message
|
||||||
|
await cl.Message(
|
||||||
|
content="✅ Documents are ready! You can now ask questions about them."
|
||||||
|
).send()
|
||||||
|
|
||||||
|
# Store session state
|
||||||
|
cl.user_session.set("documents_loaded", True)
|
||||||
|
|
||||||
|
# Reset the chat interface
|
||||||
|
await cl.Message(content="Ask your first question:").send()
|
||||||
|
|
||||||
|
# Clear any existing message elements
|
||||||
|
cl.user_session.set("message_elements", [])
|
||||||
|
|
||||||
|
else:
|
||||||
|
await cl.Message(
|
||||||
|
content="❌ No documents were successfully processed. Please try uploading again."
|
||||||
|
).send()
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
error_msg = f"❌ Error with provided settings: {str(e)}"
|
error_msg = f"❌ Error with provided settings: {str(e)}"
|
||||||
logger.error(error_msg)
|
logger.error(error_msg)
|
||||||
|
|
@ -138,39 +176,38 @@ async def handle_settings_update(settings: dict):
|
||||||
|
|
||||||
@cl.on_chat_start
|
@cl.on_chat_start
|
||||||
async def start() -> None:
|
async def start() -> None:
|
||||||
|
"""Initialize chat and request API keys."""
|
||||||
try:
|
try:
|
||||||
logger.info("Chat session started")
|
logger.info("Chat session started")
|
||||||
cl.user_session.set("chat_history", [])
|
cl.user_session.set("chat_history", [])
|
||||||
|
|
||||||
# Just show the settings form
|
# Show settings form first
|
||||||
await cl.ChatSettings(
|
await cl.ChatSettings([
|
||||||
[
|
TextInput(
|
||||||
TextInput(
|
id="OpenAIApiKey",
|
||||||
id="OpenAIApiKey",
|
label="OpenAI API Key",
|
||||||
label="OpenAI API Key",
|
initial="",
|
||||||
initial="",
|
placeholder="Enter your OpenAI API Key (starts with 'sk-')"
|
||||||
placeholder="Enter your OpenAI API Key (starts with 'sk-')"
|
),
|
||||||
),
|
TextInput(
|
||||||
TextInput(
|
id="AnthropicApiKey",
|
||||||
id="AnthropicApiKey",
|
label="Anthropic API Key",
|
||||||
label="Anthropic API Key",
|
initial="",
|
||||||
initial="",
|
placeholder="Enter your Anthropic API Key (starts with 'sk-ant-')"
|
||||||
placeholder="Enter your Anthropic API Key (starts with 'sk-ant-')"
|
),
|
||||||
),
|
TextInput(
|
||||||
TextInput(
|
id="CohereApiKey",
|
||||||
id="CohereApiKey",
|
label="Cohere API Key",
|
||||||
label="Cohere API Key",
|
initial="",
|
||||||
initial="",
|
placeholder="Enter your Cohere API Key"
|
||||||
placeholder="Enter your Cohere API Key"
|
),
|
||||||
),
|
TextInput(
|
||||||
TextInput(
|
id="DBUrl",
|
||||||
id="DBUrl",
|
label="Database URL",
|
||||||
label="Database URL",
|
initial="",
|
||||||
initial="",
|
placeholder="Enter your Database URL (e.g., postgresql://user:pass@host:port/db)"
|
||||||
placeholder="Enter your Database URL (e.g., postgresql://user:pass@host:port/db)"
|
),
|
||||||
),
|
]).send()
|
||||||
]
|
|
||||||
).send()
|
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error in chat start: {str(e)}")
|
logger.error(f"Error in chat start: {str(e)}")
|
||||||
|
|
@ -178,67 +215,73 @@ async def start() -> None:
|
||||||
|
|
||||||
@cl.on_message
|
@cl.on_message
|
||||||
async def message_handler(message: cl.Message) -> None:
|
async def message_handler(message: cl.Message) -> None:
|
||||||
|
"""Handle user queries using RAG."""
|
||||||
try:
|
try:
|
||||||
msg = cl.Message(content="Thinking...")
|
# Check if documents are loaded
|
||||||
|
if not cl.user_session.get("documents_loaded"):
|
||||||
|
await cl.Message(content="❌ Please upload and process documents first!").send()
|
||||||
|
return
|
||||||
|
|
||||||
|
if not my_config:
|
||||||
|
await cl.Message(content="❌ Please configure your API keys first!").send()
|
||||||
|
return
|
||||||
|
|
||||||
|
# Create message for streaming
|
||||||
|
msg = cl.Message(content="")
|
||||||
await msg.send()
|
await msg.send()
|
||||||
|
|
||||||
query = message.content.strip()
|
query = message.content.strip()
|
||||||
chat_history = cl.user_session.get("chat_history", [])
|
logger.info(f"Processing query: {query}")
|
||||||
|
|
||||||
# Search for relevant chunks using global config
|
# Search for relevant chunks
|
||||||
reranked_chunks = perform_search(query)
|
reranked_chunks = perform_search(query)
|
||||||
|
|
||||||
if reranked_chunks:
|
if not reranked_chunks:
|
||||||
logger.info("Using RAG for response generation")
|
|
||||||
try:
|
|
||||||
# Convert chat history to proper format for RAG
|
|
||||||
formatted_messages = []
|
|
||||||
for user_msg, assistant_msg in chat_history:
|
|
||||||
formatted_messages.append({"role": "user", "content": user_msg})
|
|
||||||
formatted_messages.append({"role": "assistant", "content": assistant_msg})
|
|
||||||
|
|
||||||
response_stream = rag(
|
|
||||||
prompt=query,
|
|
||||||
system_prompt=RAG_SYSTEM_PROMPT,
|
|
||||||
search=hybrid_search,
|
|
||||||
messages=formatted_messages,
|
|
||||||
max_contexts=5,
|
|
||||||
config=my_config
|
|
||||||
)
|
|
||||||
|
|
||||||
full_response = ""
|
|
||||||
for chunk in response_stream:
|
|
||||||
full_response += chunk
|
|
||||||
await msg.stream_token(chunk)
|
|
||||||
|
|
||||||
await msg.send()
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"RAG error: {str(e)}")
|
|
||||||
# If RAG fails, fall back to general Claude
|
|
||||||
await handle_fallback(query, msg)
|
|
||||||
return
|
|
||||||
|
|
||||||
else:
|
|
||||||
logger.info("No relevant chunks found, falling back to general Claude response")
|
logger.info("No relevant chunks found, falling back to general Claude response")
|
||||||
await handle_fallback(query, msg)
|
await handle_fallback(query, msg)
|
||||||
return
|
return
|
||||||
|
|
||||||
# Update chat history
|
# Use RAG for response generation
|
||||||
chat_history.append((query, full_response))
|
try:
|
||||||
cl.user_session.set("chat_history", chat_history)
|
chat_history = cl.user_session.get("chat_history", [])
|
||||||
|
formatted_messages = []
|
||||||
|
for user_msg, assistant_msg in chat_history:
|
||||||
|
formatted_messages.append({"role": "user", "content": user_msg})
|
||||||
|
formatted_messages.append({"role": "assistant", "content": assistant_msg})
|
||||||
|
|
||||||
|
response_stream = rag(
|
||||||
|
prompt=query,
|
||||||
|
system_prompt=RAG_SYSTEM_PROMPT,
|
||||||
|
search=hybrid_search,
|
||||||
|
messages=formatted_messages,
|
||||||
|
max_contexts=5,
|
||||||
|
config=my_config
|
||||||
|
)
|
||||||
|
|
||||||
|
full_response = ""
|
||||||
|
for chunk in response_stream:
|
||||||
|
full_response += chunk
|
||||||
|
await msg.stream_token(chunk)
|
||||||
|
|
||||||
|
# Update chat history
|
||||||
|
chat_history.append((query, full_response))
|
||||||
|
cl.user_session.set("chat_history", chat_history)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"RAG error: {str(e)}")
|
||||||
|
await handle_fallback(query, msg)
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
error_msg = f"Error processing your question: {str(e)}"
|
error_msg = f"Error processing your question: {str(e)}"
|
||||||
logger.error(error_msg)
|
logger.error(error_msg)
|
||||||
await msg.send(content=error_msg) # Use send instead of update
|
await cl.Message(content=error_msg).send()
|
||||||
|
|
||||||
async def handle_fallback(query: str, msg: cl.Message) -> None:
|
async def handle_fallback(query: str, msg: cl.Message) -> None:
|
||||||
"""Handle fallback to Claude when RAG is not available or fails."""
|
"""Handle fallback to Claude when RAG is not available or fails."""
|
||||||
try:
|
try:
|
||||||
user_env = cl.user_session.get("env")
|
user_env = cl.user_session.get("env")
|
||||||
client = anthropic.Anthropic(api_key=user_env["ANTHROPIC_API_KEY"])
|
client = anthropic.Anthropic(api_key=user_env["ANTHROPIC_API_KEY"])
|
||||||
|
|
||||||
response = client.messages.create(
|
response = client.messages.create(
|
||||||
model="claude-3-5-sonnet-20241022",
|
model="claude-3-5-sonnet-20241022",
|
||||||
max_tokens=1024,
|
max_tokens=1024,
|
||||||
|
|
@ -246,19 +289,19 @@ async def handle_fallback(query: str, msg: cl.Message) -> None:
|
||||||
{"role": "user", "content": query}
|
{"role": "user", "content": query}
|
||||||
]
|
]
|
||||||
)
|
)
|
||||||
|
|
||||||
full_response = response.content[0].text
|
full_response = response.content[0].text
|
||||||
await msg.send(content=full_response)
|
await msg.send(content=full_response)
|
||||||
|
|
||||||
# Update chat history
|
# Update chat history
|
||||||
chat_history = cl.user_session.get("chat_history", [])
|
chat_history = cl.user_session.get("chat_history", [])
|
||||||
chat_history.append((query, full_response))
|
chat_history.append((query, full_response))
|
||||||
cl.user_session.set("chat_history", chat_history)
|
cl.user_session.set("chat_history", chat_history)
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
error_msg = f"Fallback error: {str(e)}"
|
error_msg = f"Fallback error: {str(e)}"
|
||||||
logger.error(error_msg)
|
logger.error(error_msg)
|
||||||
await msg.send(content=error_msg)
|
await msg.send(content=error_msg)
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
cl.run()
|
cl.run()
|
||||||
Loading…
Reference in a new issue