app.py

import streamlit as st
import requests
import json

def get_ollama_models():
    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:
        return ['llama2', 'codellama', 'mistral']  # Default models if server not accessible

def generate_response(prompt, model):
    url = "http://localhost:11434/api/generate"
    data = {
        "model": model,
        "prompt": prompt,
        "stream": False
    }
    
    try:
        response = requests.post(url, json=data)
        if response.status_code == 200:
            return response.json()['response']
        return "Error: Unable to get response from model"
    except Exception as e:
        return f"Error: {str(e)}"

def main():
    st.title("🤖 SAWAS Prompt Enhancer")
    
    # Initialize chat history
    if "messages" not in st.session_state:
        st.session_state.messages = []

    # Sidebar for model selection
    with st.sidebar:
        st.header("Settings")
        models = get_ollama_models()
        selected_model = st.selectbox("Select Model", models)
        
        st.markdown("---")
        st.markdown("### About")
        st.markdown("This is a simple chat interface for Ollama models.")
        st.markdown("Make sure Ollama is running locally!")

    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # Chat input
    if prompt := st.chat_input("What would you like to ask?"):
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

        # Generate response
        with st.chat_message("assistant"):
            with st.spinner("Thinking..."):
                response = generate_response(prompt, selected_model)
                st.markdown(response)
                st.session_state.messages.append({"role": "assistant", "content": response})

if __name__ == "__main__":
    main()

    #!/usr/bin/env python
    """
    Author: Veera Babu Manyam
    Organization: SAWAS
    License: Apache License
    """
    
    # Frontend application entry point