solved the api masking issue - used ChatSettings class
This commit is contained in:
Madhu
parent
dd2eda6df8
commit
c9ea9655bd
1 changed files with 86 additions and 170 deletions
|
|
@ -79,182 +79,98 @@ def perform_search(query: str) -> List[dict]:
|
||||||
logger.error(f"Search error: {str(e)}")
|
logger.error(f"Search error: {str(e)}")
|
||||||
return []
|
return []
|
||||||
|
|
||||||
|
@cl.on_settings_update
|
||||||
|
async def handle_settings_update(settings: dict):
|
||||||
|
"""Handle settings updates when user submits the form."""
|
||||||
|
try:
|
||||||
|
# Validate API keys
|
||||||
|
def validate_key(key: str, key_type: str, valid_prefixes: tuple) -> bool:
|
||||||
|
if not key:
|
||||||
|
raise ValueError(f"{key_type} API key is required")
|
||||||
|
if valid_prefixes and not any(key.startswith(prefix) for prefix in valid_prefixes):
|
||||||
|
raise ValueError(f"Invalid {key_type} API key format")
|
||||||
|
return True
|
||||||
|
|
||||||
|
# Validate DB URL
|
||||||
|
def validate_db_url(url: str) -> bool:
|
||||||
|
valid_prefixes = ('postgresql://', 'mysql://', 'sqlite:///')
|
||||||
|
if not url:
|
||||||
|
raise ValueError("Database URL is required")
|
||||||
|
if not any(url.startswith(prefix) for prefix in valid_prefixes):
|
||||||
|
raise ValueError("Invalid database URL format")
|
||||||
|
return True
|
||||||
|
|
||||||
|
# Validate all inputs
|
||||||
|
validate_key(settings["OpenAIApiKey"], "OpenAI", ("sk-", "sk-proj-"))
|
||||||
|
validate_key(settings["AnthropicApiKey"], "Anthropic", ("sk-ant-",))
|
||||||
|
validate_key(settings["CohereApiKey"], "Cohere", tuple())
|
||||||
|
validate_db_url(settings["DBUrl"])
|
||||||
|
|
||||||
|
# Store validated values in user_env
|
||||||
|
user_env = {
|
||||||
|
"OPENAI_API_KEY": settings["OpenAIApiKey"],
|
||||||
|
"ANTHROPIC_API_KEY": settings["AnthropicApiKey"],
|
||||||
|
"COHERE_API_KEY": settings["CohereApiKey"],
|
||||||
|
"DB_URL": settings["DBUrl"]
|
||||||
|
}
|
||||||
|
|
||||||
|
# Store in user session
|
||||||
|
cl.user_session.set("env", user_env)
|
||||||
|
|
||||||
|
# Initialize RAGLite config
|
||||||
|
global my_config
|
||||||
|
my_config = initialize_config(user_env)
|
||||||
|
|
||||||
|
await cl.Message(content="✅ Successfully configured with your API keys!").send()
|
||||||
|
|
||||||
|
# Automatically prompt for PDF upload
|
||||||
|
await cl.AskFileMessage(
|
||||||
|
content="Please upload one or more PDF documents to begin!",
|
||||||
|
accept=["application/pdf"],
|
||||||
|
max_size_mb=20,
|
||||||
|
max_files=5
|
||||||
|
).send()
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
error_msg = f"❌ Error with provided settings: {str(e)}"
|
||||||
|
logger.error(error_msg)
|
||||||
|
await cl.Message(content=error_msg).send()
|
||||||
|
|
||||||
@cl.on_chat_start
|
@cl.on_chat_start
|
||||||
async def start() -> None:
|
async def start() -> None:
|
||||||
try:
|
try:
|
||||||
logger.info("Chat session started")
|
logger.info("Chat session started")
|
||||||
cl.user_session.set("chat_history", [])
|
cl.user_session.set("chat_history", [])
|
||||||
|
|
||||||
# Helper function to validate and get API key - so that if the user enters the wrong key, they can try again
|
# Just show the settings form
|
||||||
async def get_valid_api_key(key_type: str, validation_prefix: tuple) -> str:
|
await cl.ChatSettings(
|
||||||
while True:
|
[
|
||||||
key_response = await cl.AskUserMessage(
|
TextInput(
|
||||||
content=f"Please enter your {key_type} API key:",
|
id="OpenAIApiKey",
|
||||||
timeout=180
|
label="OpenAI API Key",
|
||||||
).send()
|
initial="",
|
||||||
|
placeholder="Enter your OpenAI API Key (starts with 'sk-')"
|
||||||
if not key_response or 'output' not in key_response:
|
),
|
||||||
await cl.Message(content=f"❌ {key_type} API key is required").send()
|
TextInput(
|
||||||
continue
|
id="AnthropicApiKey",
|
||||||
|
label="Anthropic API Key",
|
||||||
key_value = key_response['output'].strip()
|
initial="",
|
||||||
if not any(key_value.startswith(prefix) for prefix in validation_prefix):
|
placeholder="Enter your Anthropic API Key (starts with 'sk-ant-')"
|
||||||
await cl.Message(
|
),
|
||||||
content=f"❌ Invalid {key_type} API key format. Please try again."
|
TextInput(
|
||||||
).send()
|
id="CohereApiKey",
|
||||||
continue
|
label="Cohere API Key",
|
||||||
|
initial="",
|
||||||
return key_value
|
placeholder="Enter your Cohere API Key"
|
||||||
|
),
|
||||||
# this is the helper function to validate and get DB URL, so that if the user enters the wrong URL, they can try again
|
TextInput(
|
||||||
async def get_valid_db_url() -> str:
|
id="DBUrl",
|
||||||
valid_db_prefixes = ('postgresql://', 'mysql://', 'sqlite:///')
|
label="Database URL",
|
||||||
while True:
|
initial="",
|
||||||
db_url_response = await cl.AskUserMessage(
|
placeholder="Enter your Database URL (e.g., postgresql://user:pass@host:port/db)"
|
||||||
content="Please enter your Database URL:\nSupported formats:\n- PostgreSQL: postgresql://user:pass@host:port/db\n- MySQL: mysql://user:pass@host:port/db\n- SQLite: sqlite:///path/to/db.sqlite",
|
),
|
||||||
timeout=180
|
]
|
||||||
).send()
|
).send()
|
||||||
|
|
||||||
if not db_url_response or 'output' not in db_url_response:
|
|
||||||
await cl.Message(content="❌ Database URL is required").send()
|
|
||||||
continue
|
|
||||||
|
|
||||||
db_url_value = db_url_response['output'].strip()
|
|
||||||
if not any(db_url_value.startswith(prefix) for prefix in valid_db_prefixes):
|
|
||||||
await cl.Message(
|
|
||||||
content="❌ Invalid database URL format. Must start with one of:\n- postgresql://\n- mysql://\n- sqlite://\nPlease try again."
|
|
||||||
).send()
|
|
||||||
continue
|
|
||||||
|
|
||||||
return db_url_value
|
|
||||||
|
|
||||||
# Get and validate API keys with retry
|
|
||||||
openai_key = await get_valid_api_key("OpenAI", ("sk-", "sk-proj-"))
|
|
||||||
anthropic_key = await get_valid_api_key("Anthropic", ("sk-ant-",))
|
|
||||||
cohere_key = await get_valid_api_key("Cohere", ("",)) # Cohere keys don't have a specific prefix
|
|
||||||
|
|
||||||
# Get and validate DB URL with retry
|
|
||||||
db_url = await get_valid_db_url()
|
|
||||||
|
|
||||||
# Store validated values in user_env
|
|
||||||
user_env = {
|
|
||||||
"OPENAI_API_KEY": openai_key,
|
|
||||||
"ANTHROPIC_API_KEY": anthropic_key,
|
|
||||||
"COHERE_API_KEY": cohere_key,
|
|
||||||
"DB_URL": db_url
|
|
||||||
}
|
|
||||||
|
|
||||||
# Store in user session
|
|
||||||
cl.user_session.set("env", user_env)
|
|
||||||
logger.info("API keys stored in user session")
|
|
||||||
|
|
||||||
# Initialize RAGLite config with retry
|
|
||||||
while True:
|
|
||||||
try:
|
|
||||||
global my_config
|
|
||||||
my_config = initialize_config(user_env)
|
|
||||||
await cl.Message(content="✅ Successfully configured with your API keys!").send()
|
|
||||||
break
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Configuration error: {str(e)}")
|
|
||||||
await cl.Message(content=f"❌ Error configuring with provided keys: {str(e)}\nPlease check your credentials and try again.").send()
|
|
||||||
# Retry getting all credentials
|
|
||||||
openai_key = await get_valid_api_key("OpenAI", ("sk-", "sk-proj-"))
|
|
||||||
anthropic_key = await get_valid_api_key("Anthropic", ("sk-ant-",))
|
|
||||||
cohere_key = await get_valid_api_key("Cohere", ("",))
|
|
||||||
db_url = await get_valid_db_url()
|
|
||||||
user_env.update({
|
|
||||||
"OPENAI_API_KEY": openai_key,
|
|
||||||
"ANTHROPIC_API_KEY": anthropic_key,
|
|
||||||
"COHERE_API_KEY": cohere_key,
|
|
||||||
"DB_URL": db_url
|
|
||||||
})
|
|
||||||
cl.user_session.set("env", user_env)
|
|
||||||
|
|
||||||
async def get_valid_documents() -> List[cl.File]:
|
|
||||||
while True:
|
|
||||||
files = await cl.AskFileMessage(
|
|
||||||
content="Please upload one or more PDF documents to begin!",
|
|
||||||
accept=["application/pdf"],
|
|
||||||
max_size_mb=20,
|
|
||||||
max_files=5
|
|
||||||
).send()
|
|
||||||
|
|
||||||
if not files:
|
|
||||||
await cl.Message(content="❌ No files were uploaded. Please try again.").send()
|
|
||||||
continue
|
|
||||||
|
|
||||||
return files
|
|
||||||
|
|
||||||
# Process documents with retry for each file
|
|
||||||
async def process_documents(files: List[cl.File]) -> bool:
|
|
||||||
"""Process uploaded documents with retry functionality for failed files."""
|
|
||||||
processed_files = set()
|
|
||||||
files_to_process = files.copy()
|
|
||||||
|
|
||||||
while files_to_process:
|
|
||||||
current_file = files_to_process.pop(0)
|
|
||||||
|
|
||||||
if current_file.name in processed_files:
|
|
||||||
continue
|
|
||||||
|
|
||||||
logger.info(f"Processing file: {current_file.name}")
|
|
||||||
step = cl.Step(name=f"Processing {current_file.name}...")
|
|
||||||
async with step: # Use step as context manager
|
|
||||||
try:
|
|
||||||
process_document(current_file.path)
|
|
||||||
processed_files.add(current_file.name)
|
|
||||||
await cl.Message(
|
|
||||||
content=f"✅ The Document '{current_file.name}' is processed successfully."
|
|
||||||
).send()
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Failed to process '{current_file.name}': {str(e)}")
|
|
||||||
error_message = f"❌ Failed to process '{current_file.name}': {str(e)}"
|
|
||||||
|
|
||||||
# Ask user if they want to retry this file
|
|
||||||
retry = await cl.AskUserMessage(
|
|
||||||
content=f"{error_message}\nWould you like to try uploading this file again? (yes/no)",
|
|
||||||
timeout=180
|
|
||||||
).send()
|
|
||||||
|
|
||||||
if retry and retry['output'].lower().strip() == 'yes':
|
|
||||||
new_file = await cl.AskFileMessage(
|
|
||||||
content=f"Please upload '{current_file.name}' again:",
|
|
||||||
accept=["application/pdf"],
|
|
||||||
max_size_mb=20,
|
|
||||||
max_files=1
|
|
||||||
).send()
|
|
||||||
|
|
||||||
if new_file:
|
|
||||||
files_to_process.append(new_file[0])
|
|
||||||
|
|
||||||
# Ensure step is completed
|
|
||||||
await step.end()
|
|
||||||
|
|
||||||
if not processed_files:
|
|
||||||
await cl.Message(content="❌ No documents were processed successfully. Please try uploading new documents.").send()
|
|
||||||
return False
|
|
||||||
|
|
||||||
# Send final success message and return to chat
|
|
||||||
final_msg = cl.Message(content="✅ Document processing completed. You can now ask questions!")
|
|
||||||
await final_msg.send()
|
|
||||||
return True
|
|
||||||
|
|
||||||
# Main document processing loop
|
|
||||||
while True:
|
|
||||||
files = await get_valid_documents()
|
|
||||||
if await process_documents(files):
|
|
||||||
break
|
|
||||||
|
|
||||||
# Ask if user wants to try uploading different documents
|
|
||||||
retry = await cl.AskUserMessage(
|
|
||||||
content="Would you like to try uploading different documents? (yes/no)",
|
|
||||||
timeout=180
|
|
||||||
).send()
|
|
||||||
|
|
||||||
if not retry or retry['output'].lower().strip() != 'yes':
|
|
||||||
await cl.Message(content="❌ Stopping due to document processing failures.").send()
|
|
||||||
return
|
|
||||||
|
|
||||||
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)}")
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue