Create app.py

app.py

@@ -0,0 +1,363 @@
+import streamlit as st
+import requests
+import json
+import datetime
+import tempfile
+import subprocess
+import black
+from streamlit_ace import st_ace
+from streamlit_extras.colored_header import colored_header
+from streamlit_extras.add_vertical_space import add_vertical_space
+import re
+from typing import Optional, Dict, List
+import ast
+
+def get_ollama_models():
+    """Fetch available models from Ollama"""
+    try:
+        response = requests.get("http://localhost:11434/api/tags")
+        if response.status_code == 200:
+            models = [model["name"] for model in response.json()["models"]]
+            return models
+        return []
+    except Exception as e:
+        st.error(f"Error fetching models: {str(e)}")
+        return []
+
+def clear_chat():
+    """Clear the chat history"""
+    if 'messages' in st.session_state:
+        st.session_state.messages = []
+    if 'current_session' in st.session_state:
+        st.session_state.current_session = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+
+def handle_file_upload():
+    """Handle file upload functionality"""
+    uploaded_file = st.file_uploader(
+        "Upload a file",
+        type=["txt", "pdf", "py", "json", "csv"],
+        help="Upload a file to discuss with the AI"
+    )
+    
+    if uploaded_file is not None:
+        file_contents = uploaded_file.read()
+        if uploaded_file.type == "application/pdf":
+            return f"Uploaded PDF: {uploaded_file.name}"
+        else:
+            try:
+                return file_contents.decode()
+            except UnicodeDecodeError:
+                return "Binary file uploaded"
+    return None
+
+def generate_response(prompt, model, temperature, max_tokens, system_prompt, stream=False):
+    """Generate response from Ollama model with parameters"""
+    url = "http://localhost:11434/api/generate"
+    data = {
+        "model": model,
+        "prompt": prompt,
+        "system": system_prompt,
+        "temperature": temperature,
+        "max_tokens": max_tokens,
+        "stream": stream
+    }
+    
+    try:
+        if stream:
+            response_text = ""
+            with requests.post(url, json=data, stream=True) as response:
+                for line in response.iter_lines():
+                    if line:
+                        json_response = json.loads(line)
+                        response_text += json_response.get("response", "")
+                        yield response_text
+        else:
+            response = requests.post(url, json=data)
+            if response.status_code == 200:
+                return response.json()["response"]
+            else:
+                st.error(f"Error {response.status_code}: {response.text}")
+                return "Error: Unable to get response from the model."
+    except Exception as e:
+        st.error(f"Error generating response: {str(e)}")
+        return f"Error: {str(e)}"
+
+class CodeAnalyzer:
+    @staticmethod
+    def extract_code_blocks(text: str) -> List[str]:
+        """Extract code blocks from markdown text"""
+        code_blocks = re.findall(r'```(?:python)?\n(.*?)\n```', text, re.DOTALL)
+        return code_blocks
+    
+    @staticmethod
+    def is_code_complete(code: str) -> bool:
+        """Check if the code block is syntactically complete"""
+        try:
+            ast.parse(code)
+            return True
+        except SyntaxError:
+            return False
+    
+    @staticmethod
+    def get_context(code: str) -> Dict:
+        """Analyze code to extract context (variables, functions, classes)"""
+        context = {
+            'variables': [],
+            'functions': [],
+            'classes': []
+        }
+        
+        try:
+            tree = ast.parse(code)
+            for node in ast.walk(tree):
+                if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Store):
+                    context['variables'].append(node.id)
+                elif isinstance(node, ast.FunctionDef):
+                    context['functions'].append(node.name)
+                elif isinstance(node, ast.ClassDef):
+                    context['classes'].append(node.name)
+        except:
+            pass
+        
+        return context
+
+class CodeCompletion:
+    def __init__(self, model: str):
+        self.model = model
+    
+    def get_completion_suggestions(self, code: str, context: Dict) -> str:
+        """Generate code completion suggestions based on context"""
+        prompt = f"""Given the following code context:
+Code:
+{code}
+
+Context:
+Variables: {', '.join(context['variables'])}
+Functions: {', '.join(context['functions'])}
+Classes: {', '.join(context['classes'])}
+
+Please complete or continue this code in a natural way."""
+        
+        response = generate_response(prompt, self.model, 0.3, 500, 
+                                  "You are a Python coding assistant. Provide only code completion, no explanations.")
+        return response
+
+def handle_code_continuation(incomplete_code: str, model: str) -> str:
+    """Handle continuation of incomplete code"""
+    prompt = f"""Complete the following Python code:
+{incomplete_code}
+
+Provide only the completion part that would make this code syntactically complete and logical."""
+    
+    response = generate_response(prompt, model, 0.3, 500,
+                               "You are a Python coding assistant. Complete the code naturally.")
+    return response
+
+def format_code(code: str) -> str:
+    """Format Python code using black"""
+    try:
+        return black.format_str(code, mode=black.FileMode())
+    except:
+        return code
+
+def init_session_state():
+    """Initialize session state variables"""
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+    if "sessions" not in st.session_state:
+        st.session_state.sessions = {}
+    if "current_session" not in st.session_state:
+        st.session_state.current_session = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    if "system_prompt" not in st.session_state:
+        st.session_state.system_prompt = "You are a helpful AI assistant."
+    if "saved_code_snippets" not in st.session_state:
+        st.session_state.saved_code_snippets = []
+    if "code_context" not in st.session_state:
+        st.session_state.code_context = {}
+    if "current_code_block" not in st.session_state:
+        st.session_state.current_code_block = None
+    if "code_history" not in st.session_state:
+        st.session_state.code_history = []
+    if "last_code_state" not in st.session_state:
+        st.session_state.last_code_state = None
+
+def setup_page_config():
+    """Setup page configuration and styling"""
+    st.set_page_config(
+        page_title="Enhanced Ollama Chat",
+        page_icon="🤖",
+        layout="wide",
+        initial_sidebar_state="expanded"
+    )
+    
+    st.markdown("""
+    <style>
+        /* Main container styling */
+        .main {
+            max-width: 1200px;
+            margin: 0 auto;
+            padding: 2rem;
+        }
+        
+        /* Message container styling */
+        .stChatMessage {
+            background-color: #ffffff;
+            border-radius: 8px;
+            padding: 1rem;
+            margin: 0.5rem 0;
+            box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
+        }
+        
+        /* Chat input styling */
+        .stChatInputContainer {
+            border-radius: 8px;
+            border: 1px solid #e0e0e0;
+            padding: 0.5rem;
+            background-color: #ffffff;
+        }
+        
+        /* Code editor styling */
+        .code-editor {
+            border-radius: 8px;
+            margin: 1rem 0;
+            border: 1px solid #e0e0e0;
+        }
+        
+        /* Code snippet container */
+        .code-snippet {
+            background-color: #f8fafc;
+            padding: 1rem;
+            border-radius: 8px;
+            margin: 0.5rem 0;
+        }
+        
+        /* Code completion suggestions */
+        .completion-suggestion {
+            background-color: #f1f5f9;
+            padding: 0.5rem;
+            border-left: 3px solid #0284c7;
+            margin: 0.25rem 0;
+        }
+    </style>
+    """, unsafe_allow_html=True)
+
+def code_editor_section():
+    """Render the code editor section"""
+    st.subheader("📝 Code Editor")
+    
+    code_content = st_ace(
+        value=st.session_state.current_code_block or "",
+        language="python",
+        theme="monokai",
+        key="code_editor",
+        height=300,
+        show_gutter=True,
+        wrap=True,
+        auto_update=True
+    )
+    
+    col1, col2 = st.columns(2)
+    
+    with col1:
+        if st.button("Format Code"):
+            st.session_state.current_code_block = format_code(code_content)
+    
+    with col2:
+        if st.button("Get Completion Suggestions"):
+            if code_content:
+                code_analyzer = CodeAnalyzer()
+                context = code_analyzer.get_context(code_content)
+                completion = CodeCompletion(st.session_state.selected_model)
+                suggestions = completion.get_completion_suggestions(code_content, context)
+                st.code(suggestions, language="python")
+
+def main():
+    """Main application logic"""
+    setup_page_config()
+    init_session_state()
+    
+    # Sidebar configuration
+    with st.sidebar:
+        colored_header(label="Model Settings", description="Configure your chat parameters", color_name="blue-70")
+        
+        available_models = get_ollama_models()
+        if not available_models:
+            st.error("⚠️ No Ollama models found. Please make sure Ollama is running and models are installed.")
+            st.stop()
+        
+        selected_model = st.selectbox("Choose a model", available_models, index=0 if available_models else None)
+        st.session_state.selected_model = selected_model
+        
+        with st.expander("Advanced Settings", expanded=False):
+            temperature = st.slider("Temperature", 0.0, 2.0, 0.7, 0.1)
+            max_tokens = st.number_input("Max Tokens", 50, 4096, 2048)
+            system_prompt = st.text_area("System Prompt", st.session_state.system_prompt)
+            stream_output = st.checkbox("Stream Output", value=True)
+        
+        if st.button("Clear Chat"):
+            clear_chat()
+    
+    st.title("🤖 Enhanced Ollama Chat")
+    st.caption(f"Currently using: {selected_model}")
+    
+    # Main interface tabs
+    tab1, tab2 = st.tabs(["Chat", "Code Editor"])
+    
+    with tab1:
+        # File upload section
+        uploaded_content = handle_file_upload()
+        if uploaded_content:
+            st.session_state.messages.append({
+                "role": "user",
+                "content": f"I've uploaded the following content:\n\n{uploaded_content}"
+            })
+        
+        # Display chat messages
+        for message in st.session_state.messages:
+            with st.chat_message(message["role"]):
+                st.markdown(message["content"])
+                
+                # Check for code blocks in the message
+                code_blocks = CodeAnalyzer.extract_code_blocks(message["content"])
+                if code_blocks and message["role"] == "assistant":
+                    for code in code_blocks:
+                        if not CodeAnalyzer.is_code_complete(code):
+                            st.info("This code block appears to be incomplete. Would you like to complete it?")
+                            if st.button("Complete Code", key=f"complete_{len(code)}"):
+                                completion = handle_code_continuation(code, selected_model)
+                                st.code(completion, language="python")
+        
+        # Chat input
+        if prompt := st.chat_input("Message (use @ to attach a file, / for commands)"):
+            with st.chat_message("user"):
+                st.markdown(prompt)
+            st.session_state.messages.append({"role": "user", "content": prompt})
+            
+            with st.chat_message("assistant"):
+                if stream_output:
+                    message_placeholder = st.empty()
+                    for response in generate_response(prompt, selected_model, temperature, max_tokens, system_prompt, stream=True):
+                        message_placeholder.markdown(response)
+                    final_response = response
+                else:
+                    with st.spinner("Thinking..."):
+                        final_response = generate_response(prompt, selected_model, temperature, max_tokens, system_prompt, stream=False)
+                        st.markdown(final_response)
+                
+                # Store code blocks in context
+                code_blocks = CodeAnalyzer.extract_code_blocks(final_response)
+                if code_blocks:
+                    st.session_state.last_code_state = code_blocks[-1]
+                
+            st.session_state.messages.append({"role": "assistant", "content": final_response})
+    
+    with tab2:
+        code_editor_section()
+    
+    # Footer
+    add_vertical_space(2)
+    st.markdown("---")
+    st.markdown("Made with ❤️ using Streamlit and Ollama")
+
+if __name__ == "__main__":
+    main()