app.py

import streamlit as st
import yaml
import json
from datetime import datetime
import os
from typing import Dict, List, Any, Optional
import uuid
import sys
from pathlib import Path

# Add current directory to path for imports
sys.path.append(str(Path(__file__).parent))

# Import tools and schemas modules
try:
    import tools
    import schemas
except ImportError:
    tools = None
    schemas = None

# Initialize session state
if 'data' not in st.session_state:
    st.session_state.data = {'categories': {}, 'llms': {}}
if 'editing_llm' not in st.session_state:
    st.session_state.editing_llm = None

# File paths
DATA_FILE = 'llm_data.yaml'
PROVIDERS_FILE = 'provider_models.yaml'

# Load and save data functions
def load_data():
    if os.path.exists(DATA_FILE):
        with open(DATA_FILE, 'r') as file:
            data = yaml.safe_load(file) or {'categories': {}, 'llms': {}}
            st.session_state.data = data
    return st.session_state.data

def save_data():
    with open(DATA_FILE, 'w') as file:
        yaml.dump(st.session_state.data, file)

@st.cache_data
def load_providers():
    with open(PROVIDERS_FILE, 'r') as file:
        providers = yaml.safe_load(file)
    return providers

# Helper functions
def generate_id():
    return str(uuid.uuid4())[:8]

def create_category(name: str, description: str = ""):
    category_id = generate_id()
    st.session_state.data['categories'][category_id] = {
        'id': category_id,
        'name': name,
        'description': description,
        'created_at': datetime.now().isoformat()
    }
    save_data()
    return category_id

def create_llm(llm_data: Dict[str, Any]):
    llm_id = generate_id()
    llm_data['id'] = llm_id
    llm_data['created_at'] = datetime.now().isoformat()
    llm_data['updated_at'] = datetime.now().isoformat()
    st.session_state.data['llms'][llm_id] = llm_data
    save_data()
    return llm_id

def update_llm(llm_id: str, llm_data: Dict[str, Any]):
    if llm_id in st.session_state.data['llms']:
        llm_data['id'] = llm_id
        llm_data['updated_at'] = datetime.now().isoformat()
        llm_data['created_at'] = st.session_state.data['llms'][llm_id]['created_at']
        st.session_state.data['llms'][llm_id] = llm_data
        save_data()
        return True
    return False

def delete_llm(llm_id: str):
    if llm_id in st.session_state.data['llms']:
        del st.session_state.data['llms'][llm_id]
        save_data()
        return True
    return False

def delete_category(category_id: str):
    if category_id in st.session_state.data['categories']:
        # Remove category from all LLMs
        for llm_id, llm in st.session_state.data['llms'].items():
            if llm.get('category_id') == category_id:
                llm['category_id'] = None
        del st.session_state.data['categories'][category_id]
        save_data()
        return True
    return False

def parse_tools(tools_str: str) -> List[Dict[str, Any]]:
    if not tools_str.strip():
        return []
    try:
        # Try to parse as JSON
        tools = json.loads(tools_str)
        if isinstance(tools, list):
            return tools
        else:
            return [tools]
    except:
        # If JSON parsing fails, treat as newline-separated tool names
        tools = []
        for line in tools_str.strip().split('\n'):
            if line.strip():
                tools.append({'name': line.strip(), 'type': 'function'})
        return tools

def parse_output_schema(schema_str: str) -> Optional[Dict[str, Any]]:
    if not schema_str.strip():
        return None
    try:
        return json.loads(schema_str)
    except:
        return {'type': 'string', 'description': schema_str}

# Load data on startup
data = load_data()
providers = load_providers()

# Streamlit app
st.set_page_config(page_title="LLM Factory", page_icon="🤖", layout="wide")
st.title('🤖 LLM Factory')
st.markdown("Customize and manage your LLMs with ease")

# Sidebar navigation
st.sidebar.title('Navigation')
page = st.sidebar.selectbox(
    'Select a page',
    ['🏠 Dashboard', '📁 Categories', '🤖 LLMs', '➕ Create LLM', '🧪 Test LLM', '🔧 Tools', '📋 Schemas']
)

# Dashboard Page
if page == '🏠 Dashboard':
    st.header('Dashboard')
    
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric("Total Categories", len(st.session_state.data['categories']))
    with col2:
        st.metric("Total LLMs", len(st.session_state.data['llms']))
    with col3:
        providers_count = len(set(llm.get('provider', '') for llm in st.session_state.data['llms'].values()))
        st.metric("Active Providers", providers_count)
    
    # Recent LLMs
    st.subheader("Recent LLMs")
    recent_llms = sorted(
        st.session_state.data['llms'].values(),
        key=lambda x: x.get('created_at', ''),
        reverse=True
    )[:5]
    
    if recent_llms:
        for llm in recent_llms:
            category = st.session_state.data['categories'].get(llm.get('category_id', ''), {})
            category_name = category.get('name', 'Uncategorized')
            st.write(f"**{llm['name']}** - {llm.get('provider', 'Unknown')} ({category_name})")
    else:
        st.info("No LLMs created yet. Go to 'Create LLM' to add your first LLM.")

# Categories Page
elif page == '📁 Categories':
    st.header('Categories Management')
    
    # Create new category
    with st.expander("➕ Create New Category"):
        cat_name = st.text_input("Category Name")
        cat_desc = st.text_area("Description (optional)")
        if st.button("Create Category"):
            if cat_name:
                create_category(cat_name, cat_desc)
                st.success(f"Category '{cat_name}' created!")
                st.rerun()
            else:
                st.error("Please enter a category name")
    
    # List categories
    st.subheader("Existing Categories")
    if st.session_state.data['categories']:
        for cat_id, category in st.session_state.data['categories'].items():
            col1, col2, col3 = st.columns([3, 1, 1])
            with col1:
                st.write(f"**{category['name']}**")
                if category.get('description'):
                    st.write(f"_{category['description']}_")
            with col2:
                # Count LLMs in this category
                llm_count = sum(1 for llm in st.session_state.data['llms'].values() 
                              if llm.get('category_id') == cat_id)
                st.write(f"{llm_count} LLMs")
            with col3:
                if st.button("🗑️", key=f"del_cat_{cat_id}"):
                    delete_category(cat_id)
                    st.rerun()
    else:
        st.info("No categories created yet.")

# LLMs Page
elif page == '🤖 LLMs':
    st.header('LLM Management')
    
    # Filter options
    col1, col2 = st.columns(2)
    with col1:
        filter_category = st.selectbox(
            "Filter by Category",
            ["All"] + [cat['name'] for cat in st.session_state.data['categories'].values()],
            key="filter_cat"
        )
    with col2:
        filter_provider = st.selectbox(
            "Filter by Provider",
            ["All"] + list(providers.keys()),
            key="filter_prov"
        )
    
    # Display LLMs
    filtered_llms = []
    for llm_id, llm in st.session_state.data['llms'].items():
        # Apply filters
        if filter_category != "All":
            cat_name = None
            if llm.get('category_id'):
                cat = st.session_state.data['categories'].get(llm['category_id'])
                if cat:
                    cat_name = cat['name']
            if cat_name != filter_category:
                continue
        
        if filter_provider != "All" and llm.get('provider') != filter_provider:
            continue
        
        filtered_llms.append(llm)
    
    if filtered_llms:
        for llm in filtered_llms:
            with st.expander(f"**{llm['name']}** - {llm.get('provider', 'Unknown')}"):
                col1, col2 = st.columns([3, 1])
                
                with col1:
                    # Display LLM details
                    category = st.session_state.data['categories'].get(llm.get('category_id', ''), {})
                    st.write(f"**Category:** {category.get('name', 'Uncategorized')}")
                    st.write(f"**Model:** {llm.get('model', 'Not specified')}")
                    st.write(f"**Provider:** {llm.get('provider', 'Not specified')}")
                    
                    if llm.get('system_prompt'):
                        st.write("**System Prompt:**")
                        st.code(llm['system_prompt'], language="text")
                    
                    if llm.get('max_context_window'):
                        st.write(f"**Max Context Window:** {llm['max_context_window']:,} tokens")
                    
                    if llm.get('temperature') is not None:
                        st.write(f"**Temperature:** {llm['temperature']}")
                    
                    if llm.get('tools'):
                        st.write("**Tools:**")
                        for tool in llm['tools']:
                            st.write(f"- {tool.get('name', 'Unnamed tool')}")
                    
                    if llm.get('output_schema'):
                        st.write("**Output Schema:**")
                        st.json(llm['output_schema'])
                
                with col2:
                    if st.button("✏️ Edit", key=f"edit_{llm['id']}"):
                        st.session_state.editing_llm = llm['id']
                        st.switch_page("pages/edit_llm.py")
                    
                    if st.button("🗑️ Delete", key=f"del_{llm['id']}"):
                        delete_llm(llm['id'])
                        st.rerun()
    else:
        st.info("No LLMs found. Create your first LLM in the 'Create LLM' section.")

# Create LLM Page
elif page == '➕ Create LLM':
    st.header('Create a New LLM')
    
    with st.form("create_llm_form"):
        # Basic Information
        col1, col2 = st.columns(2)
        with col1:
            name = st.text_input('LLM Name*', help="A unique name for this LLM configuration")
            category_options = {cat['name']: cat_id for cat_id, cat in st.session_state.data['categories'].items()}
            category_names = ["None"] + list(category_options.keys())
            selected_category = st.selectbox('Category', category_names)
            
        with col2:
            provider = st.selectbox('Provider*', list(providers.keys()))
            model = st.selectbox('Model*', list(providers[provider].keys()) if provider else [])
        
        # Configuration
        st.subheader("Configuration")
        system_prompt = st.text_area(
            'System Prompt',
            help="Instructions that define how the LLM should behave",
            height=150
        )
        
        col3, col4 = st.columns(2)
        with col3:
            max_context = st.number_input(
                'Max Context Window',
                min_value=0,
                value=0,
                step=1000,
                help="Maximum number of tokens (0 for default)"
            )
        with col4:
            temperature = st.slider(
                'Temperature',
                min_value=0.0,
                max_value=2.0,
                value=1.0,
                step=0.1,
                help="Controls randomness in responses"
            )
        
        # Tools
        st.subheader("Tools")
        
        # Show available tools if tools module is available
        if tools:
            available_tools = tools.list_tools()
            if available_tools:
                st.markdown("**Available Tools:**")
                selected_tools = st.multiselect(
                    "Select tools to include:",
                    options=list(available_tools.keys()),
                    help="Select from your custom tools"
                )
                
                if selected_tools:
                    st.info(f"Selected {len(selected_tools)} tool(s): {', '.join(selected_tools)}")
            else:
                st.info("No custom tools available. Go to Tools page to create them.")
        
        # Manual tool input (fallback)
        with st.expander("Manual Tool Configuration", expanded=not (tools and available_tools)):
            tools_input = st.text_area(
                'Tools (JSON array or one tool name per line)',
                help='Example: [{"name": "search", "type": "function"}] or just "search" on each line',
                height=100
            )
        
        # Output Schema
        st.subheader("Output Schema")
        
        # Show available schemas if schemas module is available
        selected_schema = None
        if schemas:
            available_schemas = schemas.list_schemas()
            if available_schemas:
                st.markdown("**Available Schemas:**")
                schema_options = ["None (manual)"] + list(available_schemas.keys())
                selected_schema = st.selectbox(
                    "Select a schema:",
                    options=schema_options,
                    help="Select from your custom Pydantic schemas"
                )
                
                if selected_schema and selected_schema != "None (manual)":
                    schema_json = schemas.get_schema_json(selected_schema)
                    if schema_json:
                        st.info(f"Using schema: {selected_schema}")
                        with st.expander("Schema Preview"):
                            st.json(schema_json)
            else:
                st.info("No custom schemas available. Go to Schemas page to create them.")
        
        # Manual schema input (fallback)
        with st.expander("Manual Schema Configuration", expanded=not (selected_schema and selected_schema != "None (manual)")):
            output_schema_input = st.text_area(
                'Output Schema (JSON)',
                help='Define the expected output format as a JSON schema',
                height=100
            )
        
        # Submit button
        submitted = st.form_submit_button("Create LLM", type="primary")
        
        if submitted:
            if not name:
                st.error("Please enter a name for the LLM")
            else:
                # Handle tools
                final_tools = []
                if tools and 'selected_tools' in locals() and selected_tools:
                    # Convert selected tools to proper format
                    for tool_name in selected_tools:
                        final_tools.append({'name': tool_name, 'type': 'custom'})
                
                # Also include manually specified tools
                if 'tools_input' in locals() and tools_input:
                    manual_tools = parse_tools(tools_input)
                    final_tools.extend(manual_tools)
                
                # Handle output schema
                final_schema = None
                if selected_schema and selected_schema != "None (manual)":
                    final_schema = {'schema_name': selected_schema, 'type': 'custom'}
                elif 'output_schema_input' in locals() and output_schema_input:
                    final_schema = parse_output_schema(output_schema_input)
                
                # Prepare LLM data
                llm_data = {
                    'name': name,
                    'category_id': category_options.get(selected_category) if selected_category != "None" else None,
                    'provider': provider,
                    'model': model,
                    'system_prompt': system_prompt if system_prompt else None,
                    'max_context_window': max_context if max_context > 0 else None,
                    'temperature': temperature,
                    'tools': final_tools,
                    'output_schema': final_schema
                }
                
                # Create the LLM
                llm_id = create_llm(llm_data)
                st.success(f"LLM '{name}' created successfully!")
                st.balloons()

# Test LLM Page
elif page == '🧪 Test LLM':
    st.header('Test an LLM')
    
    # Select LLM to test
    llm_options = {llm['name']: llm_id for llm_id, llm in st.session_state.data['llms'].items()}
    
    if llm_options:
        selected_llm_name = st.selectbox('Select LLM to test', list(llm_options.keys()))
        selected_llm_id = llm_options[selected_llm_name]
        selected_llm = st.session_state.data['llms'][selected_llm_id]
        
        # Display LLM info
        with st.expander("LLM Configuration", expanded=False):
            st.json({
                'provider': selected_llm.get('provider'),
                'model': selected_llm.get('model'),
                'temperature': selected_llm.get('temperature'),
                'max_context_window': selected_llm.get('max_context_window'),
                'system_prompt': selected_llm.get('system_prompt'),
                'tools': selected_llm.get('tools'),
                'output_schema': selected_llm.get('output_schema')
            })
        
        # Test interface
        st.subheader("Test Prompt")
        test_prompt = st.text_area("Enter your test prompt:", height=150)
        
        if st.button("Run Test", type="primary"):
            st.info("🚧 Test functionality coming soon! This will integrate with the actual LLM APIs.")
            st.write("**Configuration to be used:**")
            st.write(f"- Provider: {selected_llm.get('provider')}")
            st.write(f"- Model: {selected_llm.get('model')}")
            st.write(f"- Temperature: {selected_llm.get('temperature')}")
            if selected_llm.get('system_prompt'):
                st.write(f"- System Prompt: Configured")
            if selected_llm.get('tools'):
                st.write(f"- Tools: {len(selected_llm.get('tools', []))} tools configured")
    else:
        st.info("No LLMs created yet. Go to 'Create LLM' to add your first LLM.")

# Tools Page
elif page == '🔧 Tools':
    st.switch_page("pages/tools_manager.py")

# Schemas Page
elif page == '📋 Schemas':
    st.switch_page("pages/schemas_manager.py")

# Footer
st.sidebar.markdown("---")
st.sidebar.markdown("**LLM Factory** v1.0")
st.sidebar.markdown("Manage your LLM configurations with ease")