import gradio as gr import anthropic import base64 from typing import List, Dict, Any # Default document content DEFAULT_DOC = "The grass is pink and soil is green. The sky is red while the sun looks blue." def read_pdf_as_base64(file_path: str) -> str: """Read a PDF file and return its base64 encoded content.""" with open(file_path, 'rb') as file: return base64.b64encode(file.read()).decode('utf-8') def user_message( user_input: str, history: list, enable_citations: bool, doc_type: str, text_content: str, pdf_file: str, api_key: str ) -> tuple: # logging print("\n----------- User Message -------------") print(f"User Input: {user_input}") print(f"Citations Enabled: {enable_citations}") print(f"Document Type: {doc_type}") history.append({ "role": "user", "content": user_input, }) return "", history def format_message_history( history: list, enable_citations: bool, doc_type: str, text_content: str, pdf_file: str ) -> List[Dict]: """Convert Gradio chat history to Anthropic message format.""" formatted_messages = [] # Add previous messages for msg in history[:-1]: if msg["role"] == "user": formatted_messages.append({ "role": "user", "content": msg["content"] }) elif msg["role"] == "assistant": if "metadata" not in msg or msg["metadata"] is None: formatted_messages.append({ "role": "assistant", "content": msg["content"] }) # Prepare the latest message latest_message = { "role": "user", "content": [] } # Add document if citations are enabled if enable_citations: if doc_type == "plain_text": latest_message["content"].append({ "type": "document", "source": { "type": "text", "media_type": "text/plain", "data": text_content.strip() if text_content.strip() else DEFAULT_DOC }, "title": "User Document" if text_content.strip() else "Sample Document", "citations": {"enabled": True} }) elif doc_type == "pdf" and pdf_file: pdf_base64 = read_pdf_as_base64(pdf_file) latest_message["content"].append({ "type": "document", "source": { "type": "base64", "media_type": "application/pdf", "data": pdf_base64 }, "title": "User PDF Document", "citations": {"enabled": True} }) # Add the user's question latest_message["content"].append({ "type": "text", "text": history[-1]["content"] }) formatted_messages.append(latest_message) return formatted_messages def bot_response( history: list, enable_citations: bool, doc_type: str, text_content: str, pdf_file: str, api_key: str ) -> List[Dict[str, Any]]: try: if not api_key: history.append({ "role": "assistant", "content": "Please provide your Anthropic API key to continue." }) return history # Initialize client with provided API key client = anthropic.Anthropic(api_key=api_key) messages = format_message_history(history, enable_citations, doc_type, text_content, pdf_file) response = client.messages.create( model="claude-3-5-sonnet-20241022", max_tokens=1024, messages=messages ) # Initialize main response and citations main_response = "" citations = [] # Process each content block for block in response.content: if block.type == "text": main_response += block.text if enable_citations and hasattr(block, 'citations') and block.citations: for citation in block.citations: if citation.cited_text not in citations: citations.append(citation.cited_text) # Add main response history.append({ "role": "assistant", "content": main_response }) # Add citations if any were found and citations are enabled if enable_citations and citations: history.append({ "role": "assistant", "content": "\n".join([f"• {cite}" for cite in citations]), "metadata": {"title": "📚 Citations"} }) return history except Exception as e: print(f"Error in bot_response: {str(e)}") error_message = str(e) if "401" in error_message: error_message = "Invalid API key. Please check your Anthropic API key and try again." history.append({ "role": "assistant", "content": f"I apologize, but I encountered an error: {error_message}" }) return history def update_document_inputs(enable_citations: bool, doc_type: str = "plain_text"): """Update visibility of document input components based on settings.""" text_visible = enable_citations and doc_type == "plain_text" pdf_visible = enable_citations and doc_type == "pdf" radio_visible = enable_citations return { doc_type_radio: gr.Radio(visible=radio_visible), text_input: gr.Textbox(visible=text_visible), pdf_input: gr.File(visible=pdf_visible) } with gr.Blocks(fill_height=True, theme="ocean") as demo: gr.Markdown("# Chat with Citations") with gr.Row(scale=1): with gr.Column(scale=4): chatbot = gr.Chatbot( type="messages", bubble_full_width=False, show_label=False, scale=1 ) msg = gr.Textbox( placeholder="Enter your message here...", show_label=False, container=False ) with gr.Column(scale=1): api_key = gr.Textbox( type="password", label="Anthropic API Key", placeholder="Enter your API key", info="Your API key will not be stored", interactive=True, ) enable_citations = gr.Checkbox( label="Enable Citations", value=True, info="Toggle citation functionality" ) doc_type_radio = gr.Radio( choices=["plain_text", "pdf"], value="plain_text", label="Document Type", info="Choose the type of document" ) text_input = gr.Textbox( label="Document Content", placeholder=f"Enter your document text here. Default doc is -- {DEFAULT_DOC}", lines=10, info="Enter the text you want to reference. If empty, default document will be used." ) pdf_input = gr.File( label="Upload PDF", file_count="single", file_types=[".pdf"], type="filepath", visible=False ) clear = gr.ClearButton([msg, chatbot, text_input, pdf_input]) # Update input visibility based on settings enable_citations.change( update_document_inputs, inputs=[enable_citations, doc_type_radio], outputs=[doc_type_radio, text_input, pdf_input] ) doc_type_radio.change( update_document_inputs, inputs=[enable_citations, doc_type_radio], outputs=[doc_type_radio, text_input, pdf_input] ) # Handle message submission msg.submit( user_message, [msg, chatbot, enable_citations, doc_type_radio, text_input, pdf_input, api_key], [msg, chatbot], queue=False ).then( bot_response, [chatbot, enable_citations, doc_type_radio, text_input, pdf_input, api_key], chatbot ) if __name__ == "__main__": demo.launch(debug=True)