Spaces:

edgemate
/

edgellm

Runtime error

File size: 10,079 Bytes

import { useState, useEffect } from 'react'
import { useNavigate } from 'react-router-dom'
import { Card, CardHeader, CardTitle, CardContent } from '@/components/ui/card'
import { Button } from '@/components/ui/button'
import { Badge } from '@/components/ui/badge'
import { 
  BookOpen, 
  Brain, 
  Zap, 
  Download, 
  Trash2, 
  Loader2,
  Info,
  CheckCircle,
  Cloud,
  HardDrive
} from 'lucide-react'

interface ModelInfo {
  model_name: string
  name: string
  supports_thinking: boolean
  ram_required_gb: string
  size_gb: string
  is_loaded: boolean
  type: 'local' | 'api'
}

interface ModelsResponse {
  models: ModelInfo[]
  current_model: string
}

export function Models() {
  const navigate = useNavigate()
  const [models, setModels] = useState<ModelInfo[]>([])
  const [loading, setLoading] = useState(true)
  const [modelLoading, setModelLoading] = useState<string | null>(null)

  useEffect(() => {
    fetchModels()
  }, [])

  const fetchModels = async () => {
    try {
      const baseUrl = `${window.location.protocol}//${window.location.host}`
      const res = await fetch(`${baseUrl}/models`)
      
      if (!res.ok) {
        throw new Error(`Failed to fetch models: ${res.status}`)
      }
      
      const data: ModelsResponse = await res.json()
      setModels(data.models)
    } catch (error) {
      console.error('Error fetching models:', error)
    } finally {
      setLoading(false)
    }
  }

  const loadModel = async (modelName: string) => {
    setModelLoading(modelName)
    try {
      const baseUrl = `${window.location.protocol}//${window.location.host}`
      const res = await fetch(`${baseUrl}/load-model`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
        },
        body: JSON.stringify({ model_name: modelName }),
      })

      if (!res.ok) {
        throw new Error(`Failed to load model: ${res.status}`)
      }

      // Refresh models list
      fetchModels()
    } catch (error) {
      console.error('Error loading model:', error)
    } finally {
      setModelLoading(null)
    }
  }

  const unloadModel = async (modelName: string) => {
    setModelLoading(modelName)
    try {
      const baseUrl = `${window.location.protocol}//${window.location.host}`
      const res = await fetch(`${baseUrl}/unload-model`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
        },
        body: JSON.stringify({ model_name: modelName }),
      })

      if (!res.ok) {
        throw new Error(`Failed to unload model: ${res.status}`)
      }

      // Refresh models list
      fetchModels()
    } catch (error) {
      console.error('Error unloading model:', error)
    } finally {
      setModelLoading(null)
    }
  }

  if (loading) {
    return (
      <div className="min-h-screen bg-background flex items-center justify-center">
        <Loader2 className="h-8 w-8 animate-spin" />
      </div>
    )
  }

  return (
    <div className="min-h-screen bg-background">
      {/* Header */}
      <div className="border-b bg-background/95 backdrop-blur supports-[backdrop-filter]:bg-background/60">
        <div className="max-w-6xl mx-auto p-6">
          <div className="flex items-center justify-between">
            <div className="flex items-center gap-3">
              <div className="w-8 h-8 bg-blue-600 rounded-lg flex items-center justify-center">
                <BookOpen className="h-5 w-5 text-white" />
              </div>
              <div>
                <h1 className="text-2xl font-bold">Model Catalog</h1>
                <p className="text-sm text-muted-foreground">
                  Browse and manage AI models for your conversations
                </p>
              </div>
            </div>
            <Button 
              onClick={() => navigate('/playground')} 
              className="flex items-center gap-2"
            >
              <Zap className="h-4 w-4" />
              Go to Playground
            </Button>
          </div>
        </div>
      </div>

      <div className="flex-1 p-6">
        <div className="max-w-6xl mx-auto space-y-6">
          
          {/* Info Card */}
          <Card className="bg-blue-50 border-blue-200">
            <CardContent className="pt-6">
              <div className="flex items-start gap-3">
                <Info className="h-5 w-5 text-blue-600 mt-0.5" />
                <div>
                  <h3 className="font-medium text-blue-900">Model Management</h3>
                  <p className="text-sm text-blue-700 mt-1">
                    Load models to use them in the playground. Models are cached locally for faster access.
                    Each model requires significant storage space and initial download time.
                  </p>
                </div>
              </div>
            </CardContent>
          </Card>

          {/* API Models Section */}
          <div>
            <h2 className="text-xl font-semibold mb-4 flex items-center gap-2">
              <Cloud className="h-5 w-5" />
              API Models
              <Badge variant="outline" className="text-xs">Cloud-Powered</Badge>
            </h2>
            <div className="grid grid-cols-1 md:grid-cols-2 gap-6 mb-8">
              {models.filter(m => m.type === 'api').map((model) => (
                <ModelCard 
                  key={model.model_name}
                  model={model}
                  modelLoading={modelLoading}
                  onLoad={loadModel}
                  onUnload={unloadModel}
                />
              ))}
            </div>
          </div>

          {/* Local Models Section */}
          <div>
            <h2 className="text-xl font-semibold mb-4 flex items-center gap-2">
              <HardDrive className="h-5 w-5" />
              Local Models
              <Badge variant="outline" className="text-xs">Self-Hosted</Badge>
            </h2>
            <div className="grid grid-cols-1 md:grid-cols-2 gap-6">
              {models.filter(m => m.type === 'local').map((model) => (
                <ModelCard 
                  key={model.model_name}
                  model={model}
                  modelLoading={modelLoading}
                  onLoad={loadModel}
                  onUnload={unloadModel}
                />
              ))}
            </div>
          </div>


        </div>
      </div>
    </div>
  )
}

// ModelCard component for reusability
interface ModelCardProps {
  model: ModelInfo
  modelLoading: string | null
  onLoad: (modelName: string) => void
  onUnload: (modelName: string) => void
}

function ModelCard({ model, modelLoading, onLoad, onUnload }: ModelCardProps) {
  const isApiModel = model.type === 'api'
  const isLoading = modelLoading === model.model_name
  const isLoaded = model.is_loaded

  return (
    <Card className="hover:shadow-md transition-shadow">
      <CardHeader className="pb-3">
        <div className="flex items-start justify-between">
          <div className="flex-1">
            <CardTitle className="text-base flex items-center gap-2">
              {isApiModel ? (
                <Cloud className="h-4 w-4 text-blue-600" />
              ) : (
                <HardDrive className="h-4 w-4 text-green-600" />
              )}
              {model.name}
            </CardTitle>
            <p className="text-xs text-muted-foreground mt-1">
              {model.model_name}
            </p>
          </div>
          {isLoaded && <CheckCircle className="h-5 w-5 text-green-600" />}
        </div>
      </CardHeader>
      
      <CardContent className="space-y-3">
        {/* Model Info */}
        <div className="space-y-2 text-sm">
          {!isApiModel && (
            <>
              <div className="flex justify-between">
                <span className="text-muted-foreground">Size:</span>
                <Badge variant="outline" className="text-xs">
                  {model.size_gb}
                </Badge>
              </div>
              <div className="flex justify-between">
                <span className="text-muted-foreground">RAM Required:</span>
                <Badge variant="outline" className="text-xs">
                  {model.ram_required_gb}
                </Badge>
              </div>
            </>
          )}
          
          <div className="flex justify-between">
            <span className="text-muted-foreground">Type:</span>
            <Badge variant={isApiModel ? "default" : "secondary"} className="text-xs">
              {isApiModel ? 'API' : 'Local'}
            </Badge>
          </div>
          
          {model.supports_thinking && (
            <div className="flex justify-between">
              <span className="text-muted-foreground">Features:</span>
              <Badge variant="outline" className="text-xs">
                <Brain className="h-3 w-3 mr-1" />
                Thinking
              </Badge>
            </div>
          )}
        </div>

        {/* Action Button */}
        <div className="pt-2">
          {isLoaded ? (
            <Button
              size="sm"
              variant="outline"
              onClick={() => onUnload(model.model_name)}
              disabled={isLoading}
              className="w-full"
            >
              {isLoading ? (
                <Loader2 className="h-4 w-4 mr-2 animate-spin" />
              ) : (
                <Trash2 className="h-4 w-4 mr-2" />
              )}
              {isLoading ? 'Unloading...' : 'Unload Model'}
            </Button>
          ) : (
            <Button
              size="sm"
              onClick={() => onLoad(model.model_name)}
              disabled={isLoading}
              className="w-full"
            >
              {isLoading ? (
                <Loader2 className="h-4 w-4 mr-2 animate-spin" />
              ) : (
                <Download className="h-4 w-4 mr-2" />
              )}
              {isLoading ? 'Loading...' : 'Load Model'}
            </Button>
          )}
        </div>
      </CardContent>
    </Card>
  )
}