webscout 8.2.4__py3-none-any.whl → 8.2.5__py3-none-any.whl

webscout/Local/server.py

@@ -1,721 +0,0 @@
-"""
-API server with OpenAI compatibility
-"""
-
-import json
-import time
-import os
-import datetime
-from typing import Dict, Any, List, Optional, AsyncGenerator, Union
-
-import uvicorn
-from fastapi import FastAPI, HTTPException, BackgroundTasks
-from fastapi.responses import StreamingResponse
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel, Field
-
-from .config import config
-from .model_manager import ModelManager
-from .llm import LLMInterface
-
-app = FastAPI(title="webscout.local API", description="OpenAI-compatible API for webscout.local")
-
-# Add CORS middleware to allow cross-origin requests
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  # Allow all origins
-    allow_credentials=True,
-    allow_methods=["*"],  # Allow all methods
-    allow_headers=["*"],  # Allow all headers
-)
-
-# Models
-model_manager: ModelManager = ModelManager()
-loaded_models: Dict[str, LLMInterface] = {}  # Cache for loaded models
-
-class Image(BaseModel):
-    """
-    Represents an image for multimodal models.
-    """
-    url: Optional[str] = None
-    data: Optional[str] = None
-
-class ToolFunction(BaseModel):
-    """
-    Represents a function for function calling.
-    """
-    name: str
-    description: Optional[str] = None
-    parameters: Dict[str, Any] = Field(default_factory=dict)
-
-class Tool(BaseModel):
-    """
-    Represents a tool for function calling.
-    """
-    type: str = "function"
-    function: ToolFunction
-
-class ToolCall(BaseModel):
-    """
-    Represents a tool call from the model.
-    """
-    function: Dict[str, Any]
-
-class ChatMessage(BaseModel):
-    """
-    Represents a single chat message for the chat completion endpoint.
-    """
-    role: str
-    content: Union[str, List[Dict[str, Any]], None] = ""
-    images: Optional[List[str]] = None
-    tool_calls: Optional[List[ToolCall]] = None
-
-class ChatCompletionRequest(BaseModel):
-    """
-    Request model for chat completions.
-    """
-    model: str
-    messages: List[ChatMessage]
-    temperature: float = 0.7
-    top_p: float = 0.95
-    max_tokens: int = 256
-    stream: bool = False
-    stop: Optional[List[str]] = None
-    tools: Optional[List[Tool]] = None
-    format: Optional[Union[str, Dict[str, Any]]] = None
-    options: Optional[Dict[str, Any]] = None
-    keep_alive: Optional[str] = "5m"
-
-class CompletionRequest(BaseModel):
-    """
-    Request model for text completions.
-    """
-    model: str
-    prompt: str
-    suffix: Optional[str] = None
-    temperature: float = 0.7
-    top_p: float = 0.95
-    max_tokens: int = 256
-    stream: bool = False
-    stop: Optional[List[str]] = None
-    images: Optional[List[str]] = None
-    format: Optional[Union[str, Dict[str, Any]]] = None
-    options: Optional[Dict[str, Any]] = None
-    system: Optional[str] = None
-    template: Optional[str] = None
-    context: Optional[List[int]] = None
-    raw: bool = False
-    keep_alive: Optional[str] = "5m"
-
-class ModelInfo(BaseModel):
-    """
-    Model information for listing available models.
-    """
-    id: str
-    object: str = "model"
-    created: int
-    owned_by: str = "webscout.local"
-
-class ModelList(BaseModel):
-    """
-    List of available models.
-    """
-    object: str = "list"
-    data: List[ModelInfo]
-
-class EmbeddingRequest(BaseModel):
-    """
-    Request model for embeddings.
-    """
-    model: str
-    input: Union[str, List[str]]
-    truncate: bool = True
-    options: Optional[Dict[str, Any]] = None
-    keep_alive: Optional[str] = "5m"
-
-class CopyRequest(BaseModel):
-    """
-    Request model for copying a model.
-    """
-    source: str
-    destination: str
-
-class DeleteRequest(BaseModel):
-    """
-    Request model for deleting a model.
-    """
-    model: str
-
-class ShowRequest(BaseModel):
-    """
-    Request model for showing model information.
-    """
-    model: str
-    verbose: bool = False
-
-class PullRequest(BaseModel):
-    """
-    Request model for pulling a model.
-    """
-    model: str
-    insecure: bool = False
-    stream: bool = True
-
-class ModelDetails(BaseModel):
-    """
-    Detailed model information.
-    """
-    format: str = "gguf"
-    family: str = "llama"
-    families: Optional[List[str]] = None
-    parameter_size: Optional[str] = None
-    quantization_level: Optional[str] = None
-
-class ModelResponse(BaseModel):
-    """
-    Response model for model information.
-    """
-    name: str
-    modified_at: str
-    size: int
-    digest: Optional[str] = None
-    details: Optional[ModelDetails] = None
-
-class ModelsResponse(BaseModel):
-    """
-    Response model for listing models.
-    """
-    models: List[ModelResponse]
-
-class RunningModel(BaseModel):
-    """
-    Information about a running model.
-    """
-    name: str
-    model: str
-    size: int
-    digest: Optional[str] = None
-    details: Optional[ModelDetails] = None
-    expires_at: Optional[str] = None
-    size_vram: Optional[int] = None
-
-class RunningModelsResponse(BaseModel):
-    """
-    Response model for listing running models.
-    """
-    models: List[RunningModel]
-
-class VersionResponse(BaseModel):
-    """
-    Response model for version information.
-    """
-    version: str
-
-def parse_keep_alive(keep_alive: Optional[str]) -> int:
-    """
-    Parse the keep_alive parameter to seconds.
-    Args:
-        keep_alive (Optional[str]): Keep alive duration string.
-    Returns:
-        int: Keep alive duration in seconds.
-    """
-    if not keep_alive:
-        return 300  # Default 5 minutes
-
-    if keep_alive.endswith("ms"):
-        return int(keep_alive[:-2]) // 1000
-    elif keep_alive.endswith("s"):
-        return int(keep_alive[:-1])
-    elif keep_alive.endswith("m"):
-        return int(keep_alive[:-1]) * 60
-    elif keep_alive.endswith("h"):
-        return int(keep_alive[:-1]) * 3600
-    elif keep_alive == "0":
-        return 0
-    else:
-        try:
-            return int(keep_alive)
-        except ValueError:
-            return 300  # Default 5 minutes
-
-def get_model(model_name: str, options: Optional[Dict[str, Any]] = None) -> LLMInterface:
-    """
-    Get or load a model by name, using a cache for efficiency.
-    Args:
-        model_name (str): Name of the model to load.
-        options (Optional[Dict[str, Any]]): Additional options for loading the model.
-    Returns:
-        LLMInterface: Loaded model interface.
-    Raises:
-        HTTPException: If the model cannot be loaded.
-    """
-    if model_name not in loaded_models:
-        try:
-            loaded_models[model_name] = LLMInterface(model_name)
-
-            # Extract options if provided
-            n_gpu_layers = options.get("n_gpu_layers", None) if options else None
-            n_ctx = options.get("num_ctx", None) if options else None
-            verbose = options.get("verbose", False) if options else False
-
-            loaded_models[model_name].load_model(
-                n_gpu_layers=n_gpu_layers,
-                n_ctx=n_ctx,
-                verbose=verbose
-            )
-        except Exception as e:
-            raise HTTPException(status_code=404, detail=f"Model {model_name} not found: {str(e)}")
-    return loaded_models[model_name]
-
-def unload_model(model_name: str) -> None:
-    """
-    Unload a model from memory.
-    Args:
-        model_name (str): Name of the model to unload.
-    """
-    if model_name in loaded_models:
-        # Free the model resources
-        loaded_models[model_name].llm = None
-        # Remove from loaded models cache
-        del loaded_models[model_name]
-
-@app.get("/v1/models", response_model=ModelList)
-async def list_models() -> ModelList:
-    """
-    List available models.
-    Returns:
-        ModelList: List of available models.
-    """
-    models = model_manager.list_models()
-    model_list: List[ModelInfo] = []
-    for model in models:
-        model_list.append(
-            ModelInfo(
-                id=model["name"],
-                created=int(time.time()),
-            )
-        )
-    return ModelList(object="list", data=model_list)
-
-@app.post("/v1/chat/completions")
-async def create_chat_completion(request: ChatCompletionRequest, background_tasks: BackgroundTasks) -> Any:
-    """
-    Create a chat completion.
-    Args:
-        request (ChatCompletionRequest): Chat completion request.
-        background_tasks (BackgroundTasks): Background tasks for model unloading.
-    Returns:
-        StreamingResponse or dict: Streaming or regular response.
-    """
-    # Parse keep_alive parameter
-    keep_alive_seconds = parse_keep_alive(request.keep_alive)
-
-    # If messages is empty, just load the model and return
-    if not request.messages:
-        model = get_model(request.model, request.options)
-
-        # Schedule unloading if keep_alive is 0
-        if keep_alive_seconds == 0:
-            background_tasks.add_task(unload_model, request.model)
-            done_reason = "unload"
-        else:
-            done_reason = "load"
-
-        return {
-            "model": request.model,
-            "created_at": datetime.datetime.now().isoformat(),
-            "message": {
-                "role": "assistant",
-                "content": ""
-            },
-            "done_reason": done_reason,
-            "done": True
-        }
-
-    # Process messages with images if present
-    processed_messages = []
-    for msg in request.messages:
-        message_dict = {"role": msg.role, "content": msg.content}
-
-        # Handle images for multimodal models
-        if msg.images:
-            message_dict["images"] = msg.images
-
-        # Handle tool calls
-        if msg.tool_calls:
-            message_dict["tool_calls"] = [tc.model_dump() for tc in msg.tool_calls]
-
-        processed_messages.append(message_dict)
-
-    # Get or load the model
-    model = get_model(request.model, request.options)
-
-    # Handle streaming response
-    if request.stream:
-        async def generate() -> AsyncGenerator[str, None]:
-            stream = model.create_chat_completion(
-                messages=processed_messages,
-                max_tokens=request.max_tokens,
-                temperature=request.temperature,
-                top_p=request.top_p,
-                stream=True,
-                stop=request.stop,
-                tools=request.tools,
-                format=request.format,
-            )
-            for chunk in stream:
-                yield f"data: {json.dumps(chunk)}\n\n"
-            yield "data: [DONE]\n\n"
-
-            # Schedule unloading if keep_alive is 0
-            if keep_alive_seconds == 0:
-                background_tasks.add_task(unload_model, request.model)
-
-        return StreamingResponse(generate(), media_type="text/event-stream")
-    else:
-        # Non-streaming response
-        response = model.create_chat_completion(
-            messages=processed_messages,
-            max_tokens=request.max_tokens,
-            temperature=request.temperature,
-            top_p=request.top_p,
-            stream=False,
-            stop=request.stop,
-            tools=request.tools,
-            format=request.format,
-        )
-
-        # Schedule unloading if keep_alive is 0
-        if keep_alive_seconds == 0:
-            background_tasks.add_task(unload_model, request.model)
-
-        return response
-
-@app.post("/v1/completions")
-async def create_completion(request: CompletionRequest, background_tasks: BackgroundTasks) -> Any:
-    """
-    Create a text completion.
-    Args:
-        request (CompletionRequest): Completion request.
-        background_tasks (BackgroundTasks): Background tasks for model unloading.
-    Returns:
-        StreamingResponse or dict: Streaming or regular response.
-    """
-    # Parse keep_alive parameter
-    keep_alive_seconds = parse_keep_alive(request.keep_alive)
-
-    # If prompt is empty, just load the model and return
-    if not request.prompt:
-        model = get_model(request.model, request.options)
-
-        # Schedule unloading if keep_alive is 0
-        if keep_alive_seconds == 0:
-            background_tasks.add_task(unload_model, request.model)
-            done_reason = "unload"
-        else:
-            done_reason = "load"
-
-        return {
-            "model": request.model,
-            "created_at": datetime.datetime.now().isoformat(),
-            "response": "",
-            "done": True,
-            "done_reason": done_reason
-        }
-
-    # Get or load the model
-    model = get_model(request.model, request.options)
-
-    # Handle streaming response
-    if request.stream:
-        async def generate() -> AsyncGenerator[str, None]:
-            stream = model.create_completion(
-                prompt=request.prompt,
-                max_tokens=request.max_tokens,
-                temperature=request.temperature,
-                top_p=request.top_p,
-                stream=True,
-                stop=request.stop,
-                suffix=request.suffix,
-                images=request.images,
-                system=request.system,
-                template=request.template,
-                context=request.context,
-                raw=request.raw,
-                format=request.format,
-            )
-            for chunk in stream:
-                yield f"data: {json.dumps(chunk)}\n\n"
-            yield "data: [DONE]\n\n"
-
-            # Schedule unloading if keep_alive is 0
-            if keep_alive_seconds == 0:
-                background_tasks.add_task(unload_model, request.model)
-
-        return StreamingResponse(generate(), media_type="text/event-stream")
-    else:
-        # Non-streaming response
-        response = model.create_completion(
-            prompt=request.prompt,
-            max_tokens=request.max_tokens,
-            temperature=request.temperature,
-            top_p=request.top_p,
-            stream=False,
-            stop=request.stop,
-            suffix=request.suffix,
-            images=request.images,
-            system=request.system,
-            template=request.template,
-            context=request.context,
-            raw=request.raw,
-            format=request.format,
-        )
-
-        # Schedule unloading if keep_alive is 0
-        if keep_alive_seconds == 0:
-            background_tasks.add_task(unload_model, request.model)
-
-        return response
-
-@app.post("/api/embed")
-async def create_embeddings(request: EmbeddingRequest, background_tasks: BackgroundTasks) -> Any:
-    """
-    Generate embeddings from a model.
-    Args:
-        request (EmbeddingRequest): Embedding request.
-        background_tasks (BackgroundTasks): Background tasks for model unloading.
-    Returns:
-        dict: Embedding response.
-    """
-    # Parse keep_alive parameter
-    keep_alive_seconds = parse_keep_alive(request.keep_alive)
-
-    # Get or load the model
-    model = get_model(request.model, request.options)
-
-    try:
-        # Generate embeddings
-        embeddings = model.create_embeddings(
-            input=request.input,
-            truncate=request.truncate
-        )
-
-        # Schedule unloading if keep_alive is 0
-        if keep_alive_seconds == 0:
-            background_tasks.add_task(unload_model, request.model)
-
-        return embeddings
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Failed to generate embeddings: {str(e)}")
-
-@app.post("/api/tags")
-async def list_all_models() -> ModelsResponse:
-    """
-    List all available models.
-    Returns:
-        ModelsResponse: List of available models.
-    """
-    models = model_manager.list_models()
-    response_models = []
-
-    for model in models:
-        # Extract model details
-        details = ModelDetails(
-            format="gguf",
-            family="llama",  # Default, could be improved with model metadata
-            parameter_size=model.get("parameter_size", "Unknown"),
-            quantization_level=model.get("quantization_level", "Unknown")
-        )
-
-        # Create model response
-        response_models.append(
-            ModelResponse(
-                name=model["name"],
-                modified_at=model.get("downloaded_at", datetime.datetime.now().isoformat()),
-                size=os.path.getsize(model["path"]) if "path" in model else 0,
-                digest=None,  # Could be improved with actual digest calculation
-                details=details
-            )
-        )
-
-    return ModelsResponse(models=response_models)
-
-@app.get("/api/ps")
-async def list_running_models() -> RunningModelsResponse:
-    """
-    List all running models.
-    Returns:
-        RunningModelsResponse: List of running models.
-    """
-    running_models = []
-
-    for name, _ in loaded_models.items():
-        # Get model info
-        model_info = model_manager.get_model_info(name)
-
-        if model_info:
-            # Extract model details
-            details = ModelDetails(
-                format="gguf",
-                family="llama",  # Default, could be improved with model metadata
-                parameter_size=model_info.get("parameter_size", "Unknown"),
-                quantization_level=model_info.get("quantization_level", "Unknown")
-            )
-
-            # Create running model response
-            running_models.append(
-                RunningModel(
-                    name=name,
-                    model=name,
-                    size=os.path.getsize(model_info["path"]) if "path" in model_info else 0,
-                    digest=None,  # Could be improved with actual digest calculation
-                    details=details,
-                    expires_at=None,  # Could be improved with actual expiration time
-                    size_vram=None  # Could be improved with actual VRAM usage
-                )
-            )
-
-    return RunningModelsResponse(models=running_models)
-
-@app.post("/api/show")
-async def show_model(request: ShowRequest) -> Dict[str, Any]:
-    """
-    Show model information.
-    Args:
-        request (ShowRequest): Show request.
-    Returns:
-        Dict[str, Any]: Model information.
-    """
-    model_info = model_manager.get_model_info(request.model)
-
-    if not model_info:
-        raise HTTPException(status_code=404, detail=f"Model {request.model} not found")
-
-    # Extract model details
-    details = {
-        "parent_model": "",
-        "format": "gguf",
-        "family": "llama",
-        "families": ["llama"],
-        "parameter_size": model_info.get("parameter_size", "Unknown"),
-        "quantization_level": model_info.get("quantization_level", "Unknown")
-    }
-
-    # Get model capabilities
-    capabilities = ["completion"]
-
-    # Add vision capability if model supports it
-    if "vision" in model_info.get("capabilities", []):
-        capabilities.append("vision")
-
-    # Create response
-    response = {
-        "modelfile": model_info.get("modelfile", ""),
-        "parameters": model_info.get("parameters", ""),
-        "template": model_info.get("template", ""),
-        "details": details,
-        "capabilities": capabilities
-    }
-
-    # Add model_info if verbose is requested
-    if request.verbose:
-        response["model_info"] = model_info.get("model_info", {})
-
-    return response
-
-@app.post("/api/copy")
-async def copy_model(request: CopyRequest) -> Dict[str, str]:
-    """
-    Copy a model.
-    Args:
-        request (CopyRequest): Copy request.
-    Returns:
-        Dict[str, str]: Success message.
-    """
-    try:
-        # Check if source model exists
-        if not model_manager.get_model_info(request.source):
-            raise HTTPException(status_code=404, detail=f"Source model {request.source} not found")
-
-        # Copy the model
-        model_manager.copy_model(request.source, request.destination)
-
-        return {"status": "success"}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Failed to copy model: {str(e)}")
-
-@app.delete("/api/delete")
-async def delete_model(request: DeleteRequest) -> Dict[str, str]:
-    """
-    Delete a model.
-    Args:
-        request (DeleteRequest): Delete request.
-    Returns:
-        Dict[str, str]: Success message.
-    """
-    # Check if model exists
-    if not model_manager.get_model_info(request.model):
-        raise HTTPException(status_code=404, detail=f"Model {request.model} not found")
-
-    # Delete the model
-    if model_manager.remove_model(request.model):
-        return {"status": "success"}
-    else:
-        raise HTTPException(status_code=500, detail=f"Failed to delete model {request.model}")
-
-@app.post("/api/pull")
-async def pull_model(request: PullRequest) -> Union[Dict[str, str], StreamingResponse]:
-    """
-    Pull a model from Hugging Face.
-    Args:
-        request (PullRequest): Pull request.
-    Returns:
-        Union[Dict[str, str], StreamingResponse]: Success message or streaming response.
-    """
-    if request.stream:
-        async def generate_stream():
-            # Initial response
-            yield json.dumps({"status": "pulling manifest"}) + "\n"
-
-            try:
-                # Pull the model
-                _, _ = model_manager.download_model(request.model)
-
-                # Success responses
-                yield json.dumps({"status": "downloading model"}) + "\n"
-                yield json.dumps({"status": "verifying sha256 digest"}) + "\n"
-                yield json.dumps({"status": "writing manifest"}) + "\n"
-                yield json.dumps({"status": "success"}) + "\n"
-            except Exception as e:
-                yield json.dumps({"status": "error", "error": str(e)}) + "\n"
-
-        return StreamingResponse(generate_stream(), media_type="application/json")
-    else:
-        try:
-            # Pull the model
-            model_manager.download_model(request.model)
-
-            return {"status": "success"}
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Failed to pull model: {str(e)}")
-
-@app.get("/api/version")
-async def get_version() -> VersionResponse:
-    """
-    Get the version of the API.
-    Returns:
-        VersionResponse: Version information.
-    """
-    from webscout.Local import __version__
-    return VersionResponse(version=__version__)
-
-def start_server(host: Optional[str] = None, port: Optional[int] = None) -> None:
-    """
-    Start the API server.
-    Args:
-        host (Optional[str]): Host to bind the server to.
-        port (Optional[int]): Port to bind the server to.
-    """
-    host = host or config.get("api_host", "127.0.0.1")
-    port = port or config.get("api_port", 8000)
-    uvicorn.run(app, host=host, port=port)