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_files: 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 documents if citations are enabled if enable_citations: # Handle plain text input if doc_type in ["plain_text", "combined"] and text_content.strip(): latest_message["content"].append({ "type": "document", "source": { "type": "text", "media_type": "text/plain", "data": text_content.strip() }, "title": "User Text Document", "citations": {"enabled": True} }) # Handle PDF input if doc_type in ["pdf", "combined"] and pdf_files: # Handle pdf_files as a list if isinstance(pdf_files, str): pdf_files = [pdf_files] # Convert single path to list # Add each PDF as a separate document for i, pdf_file in enumerate(pdf_files): try: 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": f"User PDF Document {i+1}", "citations": {"enabled": True} }) except Exception as e: print(f"Error processing PDF {i+1}: {str(e)}") continue # If no documents were added and citations are enabled, use default document if not latest_message["content"]: latest_message["content"].append({ "type": "document", "source": { "type": "text", "media_type": "text/plain", "data": DEFAULT_DOC }, "title": "Sample 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 in ["plain_text", "combined"]) pdf_visible = enable_citations and (doc_type in ["pdf", "combined"]) 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(theme="ocean", fill_height=True) as demo: gr.Markdown("# Chat with Anthropic Claude's 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", "combined"], value="plain_text", label="Document Type", info="Choose the type of document(s) to reference" ) text_input = gr.Textbox( label="Document Content", placeholder=f"Enter your document text here.\nDefault text will be picked if citations are enabled and you don't provide the documents. Default document is --{DEFAULT_DOC}", lines=10, info="Enter the text you want to reference" ) pdf_input = gr.File( label="Upload PDF", file_count="multiple", 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)