webscout 8.2.1__py3-none-any.whl → 8.2.3__py3-none-any.whl

webscout/Local/server.py

@@ -4,12 +4,15 @@ API server with OpenAI compatibility
 
 import json
 import time
-from typing import Dict, Any, List, Optional, AsyncGenerator
+import os
+import datetime
+from typing import Dict, Any, List, Optional, AsyncGenerator, Union
 
 import uvicorn
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, BackgroundTasks
 from fastapi.responses import StreamingResponse
-from pydantic import BaseModel
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
 
 from .config import config
 from .model_manager import ModelManager
@@ -17,16 +20,55 @@ 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: str
+    content: Union[str, List[Dict[str, Any]], None] = ""
+    images: Optional[List[str]] = None
+    tool_calls: Optional[List[ToolCall]] = None
 
 class ChatCompletionRequest(BaseModel):
     """
@@ -39,6 +81,10 @@ class ChatCompletionRequest(BaseModel):
     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):
     """
@@ -46,11 +92,20 @@ class CompletionRequest(BaseModel):
     """
     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):
     """
@@ -68,11 +123,127 @@ class ModelList(BaseModel):
     object: str = "list"
     data: List[ModelInfo]
 
-def get_model(model_name: str) -> LLMInterface:
+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:
@@ -81,11 +252,33 @@ def get_model(model_name: str) -> LLMInterface:
     if model_name not in loaded_models:
         try:
             loaded_models[model_name] = LLMInterface(model_name)
-            loaded_models[model_name].load_model()
+
+            # 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:
     """
@@ -105,51 +298,135 @@ async def list_models() -> ModelList:
     return ModelList(object="list", data=model_list)
 
 @app.post("/v1/chat/completions")
-async def create_chat_completion(request: ChatCompletionRequest) -> Any:
+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.
     """
-    model = get_model(request.model)
-    messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
+    # 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=messages,
+                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=messages,
+            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) -> Any:
+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.
     """
-    model = get_model(request.model)
+    # 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(
@@ -159,12 +436,25 @@ async def create_completion(request: CompletionRequest) -> Any:
                 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,
@@ -172,9 +462,253 @@ async def create_completion(request: CompletionRequest) -> Any:
             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.