Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import os | |
| from groq import Groq | |
| from datetime import datetime | |
| import re | |
| import json | |
| import hashlib | |
| # Page configuration | |
| st.set_page_config( | |
| page_title="AI Artifact Chat", | |
| page_icon="💬", | |
| layout="wide", | |
| initial_sidebar_state="expanded" | |
| ) | |
| # Initialize session state | |
| if "messages" not in st.session_state: | |
| st.session_state.messages = [] | |
| if "artifacts" not in st.session_state: | |
| st.session_state.artifacts = [] | |
| if "conversation_start" not in st.session_state: | |
| st.session_state.conversation_start = datetime.now() | |
| if "system_prompt" not in st.session_state: | |
| st.session_state.system_prompt = """You are a helpful assistant that generates content in various formats including HTML, React components, and Three.js visualizations. When showing code examples, always use markdown code blocks with appropriate language tags (e.g. ```html, ```javascript, ```jsx).""" | |
| def clean_and_format_code(response): | |
| """Clean and format the code from the response""" | |
| # Remove thinking tags | |
| thinking_patterns = [ | |
| r'<think>.*?</think>', | |
| r'<thinking>.*?</thinking>', | |
| r'<antThinking>.*?</antThinking>', | |
| r'<thoughts>.*?</thoughts>', | |
| r'<antThought>.*?</antThought>' | |
| ] | |
| for pattern in thinking_patterns: | |
| response = re.sub(pattern, '', response, flags=re.DOTALL) | |
| # Remove markdown-style emphasis | |
| response = re.sub(r'\*\*(.*?)\*\*', r'\1', response) # Remove **bold** | |
| response = re.sub(r'\*(.*?)\*', r'\1', response) # Remove *italic* | |
| # If the response contains both script and HTML elements, wrap it in HTML structure | |
| if '<script' in response and ('<!DOCTYPE' not in response and '<html' not in response): | |
| response = f""" | |
| <!DOCTYPE html> | |
| <html> | |
| <head> | |
| <meta charset="UTF-8"> | |
| <meta name="viewport" content="width=device-width, initial-scale=1.0"> | |
| <title>3D Visualization</title> | |
| </head> | |
| <body> | |
| {response} | |
| </body> | |
| </html> | |
| """ | |
| return response.strip() | |
| def extract_code_and_type(response): | |
| """Extract code blocks and determine artifact type""" | |
| # First, clean and format the code | |
| response = clean_and_format_code(response) | |
| # Look for code blocks first | |
| code_blocks = list(re.finditer(r'```(\w+)?\n(.*?)```', response, flags=re.DOTALL)) | |
| if code_blocks: | |
| html_code = None | |
| js_code = None | |
| react_code = None | |
| for block in code_blocks: | |
| lang = block.group(1).lower() if block.group(1) else '' | |
| code = block.group(2).strip() | |
| if lang in ['html', 'htm']: | |
| html_code = code | |
| elif lang in ['javascript', 'js']: | |
| js_code = code | |
| elif lang in ['jsx', 'react', 'tsx']: | |
| react_code = code | |
| return code, "application/vnd.ant.react" | |
| # Combine HTML and JavaScript if both exist | |
| if html_code and js_code: | |
| if '</body>' in html_code: | |
| combined = html_code.replace('</body>', f'<script>{js_code}</script></body>') | |
| else: | |
| combined = f"{html_code}<script>{js_code}</script>" | |
| return combined, "text/html" | |
| elif html_code: | |
| return html_code, "text/html" | |
| elif js_code: | |
| return js_code, "application/vnd.ant.react" | |
| # If no code blocks found but content contains HTML/script tags | |
| if '<script' in response or '<html' in response or '<div' in response: | |
| return response, "text/html" | |
| elif 'function' in response or 'class' in response or 'import React' in response: | |
| return response, "application/vnd.ant.react" | |
| return response, "text/markdown" | |
| # Sidebar | |
| with st.sidebar: | |
| st.title("⚙️ Settings") | |
| # System Prompt | |
| st.markdown("### System Prompt") | |
| new_system_prompt = st.text_area( | |
| "Customize the system prompt:", | |
| value=st.session_state.system_prompt, | |
| height=150, | |
| help="This prompt will guide the AI's behavior during the current session" | |
| ) | |
| if new_system_prompt != st.session_state.system_prompt: | |
| st.session_state.system_prompt = new_system_prompt | |
| # Model Selection | |
| st.markdown("### Model Settings") | |
| model_options = [ | |
| "llama3-70b-8192", | |
| "gemma2-9b-it", | |
| "llama-3.3-70b-versatile", | |
| "deepseek-r1-distill-llama-70b" | |
| ] | |
| selected_model = st.selectbox( | |
| "Select Model:", | |
| model_options, | |
| help="Choose the AI model for generation" | |
| ) | |
| # Session Info | |
| st.markdown("### Session Info") | |
| st.markdown(f"**Started:** {st.session_state.conversation_start.strftime('%Y-%m-%d %H:%M')}") | |
| st.markdown(f"**Messages:** {len(st.session_state.messages)}") | |
| st.markdown(f"**Artifacts:** {len(st.session_state.artifacts)}") | |
| if st.button("Clear Conversation", type="secondary"): | |
| st.session_state.messages = [] | |
| st.session_state.artifacts = [] | |
| st.rerun() | |
| # Main chat interface | |
| st.title("💬 AI Artifact Chat") | |
| st.markdown("Generate content and view artifacts in the Artifacts Viewer page") | |
| # Message display | |
| for message in st.session_state.messages: | |
| with st.chat_message(message["role"]): | |
| st.markdown(message["content"]) | |
| if "artifact_id" in message: | |
| st.info(f"🎨 Artifact generated! View it in the Artifacts Viewer page (ID: {message['artifact_id']})") | |
| # Input area | |
| prompt = st.chat_input("Enter your message...") | |
| if prompt: | |
| # Add user message | |
| st.session_state.messages.append({"role": "user", "content": prompt}) | |
| # Show typing indicator | |
| with st.chat_message("assistant"): | |
| with st.spinner("Thinking..."): | |
| # Generate response using Groq | |
| client = Groq(api_key=os.getenv('GROQ_API_KEY')) | |
| chat_completion = client.chat.completions.create( | |
| messages=[ | |
| { | |
| "role": "system", | |
| "content": st.session_state.system_prompt | |
| }, | |
| *[{"role": m["role"], "content": m["content"]} | |
| for m in st.session_state.messages[-10:]] # Include last 10 messages for context | |
| ], | |
| model=selected_model, | |
| max_tokens=3000 | |
| ) | |
| response = chat_completion.choices[0].message.content | |
| # Remove thinking tags and their contents | |
| thinking_patterns = [ | |
| r'<think>.*?</think>', | |
| r'<thinking>.*?</thinking>', | |
| r'<antThinking>.*?</antThinking>', | |
| r'<thoughts>.*?</thoughts>', | |
| r'<antThought>.*?</antThought>' | |
| ] | |
| for pattern in thinking_patterns: | |
| response = re.sub(pattern, '', response, flags=re.DOTALL) | |
| # Clean up extra newlines | |
| response = re.sub(r'\n\s*\n\s*\n', '\n\n', response) | |
| # Try to extract code and determine type | |
| artifact_content, artifact_type = extract_code_and_type(response) | |
| if artifact_content: | |
| # Generate artifact ID | |
| artifact_id = hashlib.md5(artifact_content.encode()).hexdigest()[:8] | |
| # Store artifact | |
| st.session_state.artifacts.append({ | |
| "id": artifact_id, | |
| "type": artifact_type, | |
| "content": artifact_content, | |
| "created_at": datetime.now().isoformat() | |
| }) | |
| # Add assistant message with artifact reference | |
| st.session_state.messages.append({ | |
| "role": "assistant", | |
| "content": response.replace(artifact_content, '').strip(), | |
| "artifact_id": artifact_id | |
| }) | |
| else: | |
| # Add regular assistant message | |
| st.session_state.messages.append({ | |
| "role": "assistant", | |
| "content": response | |
| }) | |
| st.rerun() |