lightspeed-stack 0.1.0__py3-none-any.whl

models/responses.py

@@ -0,0 +1,244 @@
+"""Models for service responses."""
+
+from typing import Any, Optional
+
+from pydantic import BaseModel
+
+
+class ModelsResponse(BaseModel):
+    """Model representing a response to models request."""
+
+    models: list[dict[str, Any]]
+
+
+# TODO(lucasagomes): a lot of fields to add to QueryResponse. For now
+# we are keeping it simple. The missing fields are:
+# - referenced_documents: The optional URLs and titles for the documents used
+#   to generate the response.
+# - truncated: Set to True if conversation history was truncated to be within context window.
+# - input_tokens: Number of tokens sent to LLM
+# - output_tokens: Number of tokens received from LLM
+# - available_quotas: Quota available as measured by all configured quota limiters
+# - tool_calls: List of tool requests.
+# - tool_results: List of tool results.
+# See LLMResponse in ols-service for more details.
+class QueryResponse(BaseModel):
+    """Model representing LLM response to a query.
+
+    Attributes:
+        conversation_id: The optional conversation ID (UUID).
+        response: The response.
+    """
+
+    conversation_id: Optional[str] = None
+    response: str
+
+    # provides examples for /docs endpoint
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "conversation_id": "123e4567-e89b-12d3-a456-426614174000",
+                    "response": "Operator Lifecycle Manager (OLM) helps users install...",
+                }
+            ]
+        }
+    }
+
+
+class InfoResponse(BaseModel):
+    """Model representing a response to a info request.
+
+    Attributes:
+        name: Service name.
+        version: Service version.
+
+    Example:
+        ```python
+        info_response = InfoResponse(
+            name="Lightspeed Stack",
+            version="1.0.0",
+        )
+        ```
+    """
+
+    name: str
+    version: str
+
+    # provides examples for /docs endpoint
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "name": "Lightspeed Stack",
+                    "version": "1.0.0",
+                }
+            ]
+        }
+    }
+
+
+class ProviderHealthStatus(BaseModel):
+    """Model representing the health status of a provider.
+
+    Attributes:
+        provider_id: The ID of the provider.
+        status: The health status ('ok', 'unhealthy', 'not_implemented').
+        message: Optional message about the health status.
+    """
+
+    provider_id: str
+    status: str
+    message: Optional[str] = None
+
+
+class ReadinessResponse(BaseModel):
+    """Model representing response to a readiness request.
+
+    Attributes:
+        ready: If service is ready.
+        reason: The reason for the readiness.
+        providers: List of unhealthy providers in case of readiness failure.
+
+    Example:
+        ```python
+        readiness_response = ReadinessResponse(
+            ready=False,
+            reason="Service is not ready",
+            providers=[
+                ProviderHealthStatus(
+                    provider_id="ollama",
+                    status="Error",
+                    message="Server is unavailable"
+                )
+            ]
+        )
+        ```
+    """
+
+    ready: bool
+    reason: str
+    providers: list[ProviderHealthStatus]
+
+    # provides examples for /docs endpoint
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "ready": True,
+                    "reason": "Service is ready",
+                    "providers": [],
+                }
+            ]
+        }
+    }
+
+
+class LivenessResponse(BaseModel):
+    """Model representing a response to a liveness request.
+
+    Attributes:
+        alive: If app is alive.
+
+    Example:
+        ```python
+        liveness_response = LivenessResponse(alive=True)
+        ```
+    """
+
+    alive: bool
+
+    # provides examples for /docs endpoint
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "alive": True,
+                }
+            ]
+        }
+    }
+
+
+class NotAvailableResponse(BaseModel):
+    """Model representing error response for readiness endpoint."""
+
+    detail: dict[str, str]
+
+    # provides examples for /docs endpoint
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "detail": {
+                        "response": "Service is not ready",
+                        "cause": "Index is not ready",
+                    }
+                },
+                {
+                    "detail": {
+                        "response": "Service is not ready",
+                        "cause": "LLM is not ready",
+                    },
+                },
+            ]
+        }
+    }
+
+
+class FeedbackResponse(BaseModel):
+    """Model representing a response to a feedback request.
+
+    Attributes:
+        response: The response of the feedback request.
+
+    Example:
+        ```python
+        feedback_response = FeedbackResponse(response="feedback received")
+        ```
+    """
+
+    response: str
+
+    # provides examples for /docs endpoint
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "response": "feedback received",
+                }
+            ]
+        }
+    }
+
+
+class StatusResponse(BaseModel):
+    """Model representing a response to a status request.
+
+    Attributes:
+        functionality: The functionality of the service.
+        status: The status of the service.
+
+    Example:
+        ```python
+        status_response = StatusResponse(
+            functionality="feedback",
+            status={"enabled": True},
+        )
+        ```
+    """
+
+    functionality: str
+    status: dict
+
+    # provides examples for /docs endpoint
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "functionality": "feedback",
+                    "status": {"enabled": True},
+                }
+            ]
+        }
+    }

runners/__init__.py

@@ -0,0 +1 @@
+"""Runners."""

runners/uvicorn.py

@@ -0,0 +1,31 @@
+"""Uvicorn runner."""
+
+import logging
+
+import uvicorn
+
+from models.config import ServiceConfiguration
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+def start_uvicorn(configuration: ServiceConfiguration) -> None:
+    """Start Uvicorn-based REST API service."""
+    logger.info("Starting Uvicorn")
+
+    log_level = logging.INFO
+
+    # please note:
+    # TLS fields can be None, which means we will pass those values as None to uvicorn.run
+    uvicorn.run(
+        "app.main:app",
+        host=configuration.host,
+        port=configuration.port,
+        workers=configuration.workers,
+        log_level=log_level,
+        ssl_keyfile=configuration.tls_config.tls_key_path,
+        ssl_certfile=configuration.tls_config.tls_certificate_path,
+        ssl_keyfile_password=str(configuration.tls_config.tls_key_password or ""),
+        use_colors=True,
+        access_log=True,
+    )

utils/.ruff_cache/.gitignore

@@ -0,0 +1,2 @@
+# Automatically created by ruff.
+*

utils/.ruff_cache/CACHEDIR.TAG

@@ -0,0 +1 @@
+Signature: 8a477f597d28d172789f06886806bc55

utils/__init__.py

@@ -0,0 +1 @@
+"""Utils."""

utils/checks.py

@@ -0,0 +1,27 @@
+"""Checks that are performed to configuration options."""
+
+import os
+from typing import Optional
+
+from pydantic import FilePath
+
+
+class InvalidConfigurationError(Exception):
+    """Lightspeed configuration is invalid."""
+
+
+def get_attribute_from_file(data: dict, file_name_key: str) -> Optional[str]:
+    """Retrieve value of an attribute from a file."""
+    file_path = data.get(file_name_key)
+    if file_path is not None:
+        with open(file_path, encoding="utf-8") as f:
+            return f.read().rstrip()
+    return None
+
+
+def file_check(path: FilePath, desc: str) -> None:
+    """Check that path is a readable regular file."""
+    if not os.path.isfile(path):
+        raise InvalidConfigurationError(f"{desc} '{path}' is not a file")
+    if not os.access(path, os.R_OK):
+        raise InvalidConfigurationError(f"{desc} '{path}' is not readable")

utils/common.py

@@ -0,0 +1,111 @@
+"""Common utilities for the project."""
+
+from typing import Any, List, cast
+from logging import Logger
+
+from llama_stack_client import LlamaStackClient
+
+from llama_stack.distribution.library_client import (
+    LlamaStackAsLibraryClient,
+    AsyncLlamaStackAsLibraryClient,
+)
+
+from client import get_llama_stack_client
+from models.config import Configuration, ModelContextProtocolServer
+
+
+# TODO(lucasagomes): implement this function to retrieve user ID from auth
+def retrieve_user_id(auth: Any) -> str:  # pylint: disable=unused-argument
+    """Retrieve the user ID from the authentication handler.
+
+    Args:
+        auth: The Authentication handler (FastAPI Depends) that will
+            handle authentication Logic.
+
+    Returns:
+        str: The user ID.
+    """
+    return "user_id_placeholder"
+
+
+async def register_mcp_servers_async(
+    logger: Logger, configuration: Configuration
+) -> None:
+    """Register Model Context Protocol (MCP) servers with the LlamaStack client (async)."""
+    # Skip MCP registration if no MCP servers are configured
+    if not configuration.mcp_servers:
+        logger.debug("No MCP servers configured, skipping registration")
+        return
+
+    if configuration.llama_stack.use_as_library_client:
+        # Library client - use async interface
+        # config.py validation ensures library_client_config_path is not None
+        # when use_as_library_client is True
+        config_path = cast(str, configuration.llama_stack.library_client_config_path)
+        client = LlamaStackAsLibraryClient(config_path)
+        await client.async_client.initialize()
+
+        await _register_mcp_toolgroups_async(
+            client.async_client, configuration.mcp_servers, logger
+        )
+    else:
+        # Service client - use sync interface
+        client = get_llama_stack_client(configuration.llama_stack)
+
+        _register_mcp_toolgroups_sync(client, configuration.mcp_servers, logger)
+
+
+async def _register_mcp_toolgroups_async(
+    client: AsyncLlamaStackAsLibraryClient,
+    mcp_servers: List[ModelContextProtocolServer],
+    logger: Logger,
+) -> None:
+    """Async logic for registering MCP toolgroups."""
+    # Get registered tools
+    registered_toolgroups = await client.toolgroups.list()
+    registered_toolgroups_ids = [
+        tool_group.provider_resource_id for tool_group in registered_toolgroups
+    ]
+    logger.debug("Registered toolgroups: %s", registered_toolgroups_ids)
+
+    # Register toolgroups for MCP servers if not already registered
+    for mcp in mcp_servers:
+        if mcp.name not in registered_toolgroups_ids:
+            logger.debug("Registering MCP server: %s, %s", mcp.name, mcp.url)
+
+            registration_params = {
+                "toolgroup_id": mcp.name,
+                "provider_id": mcp.provider_id,
+                "mcp_endpoint": {"uri": mcp.url},
+            }
+
+            await client.toolgroups.register(**registration_params)
+            logger.debug("MCP server %s registered successfully", mcp.name)
+
+
+def _register_mcp_toolgroups_sync(
+    client: LlamaStackClient,
+    mcp_servers: List[ModelContextProtocolServer],
+    logger: Logger,
+) -> None:
+    """Sync logic for registering MCP toolgroups."""
+    # Get registered tool groups
+    registered_toolgroups = client.toolgroups.list()
+    registered_toolgroups_ids = [
+        tool_group.provider_resource_id for tool_group in registered_toolgroups
+    ]
+    logger.debug("Registered toolgroups: %s", registered_toolgroups_ids)
+
+    # Register toolgroups for MCP servers if not already registered
+    for mcp in mcp_servers:
+        if mcp.name not in registered_toolgroups_ids:
+            logger.debug("Registering MCP server: %s, %s", mcp.name, mcp.url)
+
+            registration_params = {
+                "toolgroup_id": mcp.name,
+                "provider_id": mcp.provider_id,
+                "mcp_endpoint": {"uri": mcp.url},
+            }
+
+            client.toolgroups.register(**registration_params)
+            logger.debug("MCP server %s registered successfully", mcp.name)

utils/endpoints.py

@@ -0,0 +1,34 @@
+"""Utility functions for endpoint handlers."""
+
+from fastapi import HTTPException, status
+
+import constants
+from models.requests import QueryRequest
+from configuration import AppConfig
+
+
+def check_configuration_loaded(configuration: AppConfig) -> None:
+    """Check that configuration is loaded and raise exception when it is not."""
+    if configuration is None:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail={"response": "Configuration is not loaded"},
+        )
+
+
+def get_system_prompt(query_request: QueryRequest, configuration: AppConfig) -> str:
+    """Get the system prompt: the provided one, configured one, or default one."""
+    # system prompt defined in query request has precendence
+    if query_request.system_prompt:
+        return query_request.system_prompt
+
+    # customized system prompt should be used when query request
+    # does not contain one
+    if (
+        configuration.customization is not None
+        and configuration.customization.system_prompt is not None
+    ):
+        return configuration.customization.system_prompt
+
+    # default system prompt has the lowest precedence
+    return constants.DEFAULT_SYSTEM_PROMPT

utils/mcp_headers.py

@@ -0,0 +1,48 @@
+"""MCP headers handling."""
+
+import json
+import logging
+from fastapi import Request
+
+logger = logging.getLogger("app.endpoints.dependencies")
+
+
+async def mcp_headers_dependency(_request: Request) -> dict[str, dict[str, str]]:
+    """Get the mcp headers dependency to passed to mcp servers.
+
+    mcp headers is a json dictionary or mcp url paths and their respective headers
+
+    Args:
+        request (Request): The FastAPI request object.
+
+    Returns:
+        The mcp headers dictionary, or empty dictionary if not found or on json decoding error
+    """
+    return extract_mcp_headers(_request)
+
+
+def extract_mcp_headers(request: Request) -> dict[str, dict[str, str]]:
+    """Extract mcp headers from MCP-HEADERS header.
+
+    Args:
+        request: The FastAPI request object
+
+    Returns:
+        The mcp headers dictionary, or empty dictionary if not found or on json decoding error
+    """
+    mcp_headers_string = request.headers.get("MCP-HEADERS", "")
+    mcp_headers = {}
+    if mcp_headers_string:
+        try:
+            mcp_headers = json.loads(mcp_headers_string)
+        except json.decoder.JSONDecodeError as e:
+            logger.error("MCP headers decode error: %s", e)
+
+        if not isinstance(mcp_headers, dict):
+            logger.error(
+                "MCP headers wrong type supplied (mcp headers must be a dictionary), "
+                "but type %s was supplied",
+                type(mcp_headers),
+            )
+            mcp_headers = {}
+    return mcp_headers

utils/suid.py

@@ -0,0 +1,28 @@
+"""Session ID utility functions."""
+
+import uuid
+
+
+def get_suid() -> str:
+    """Generate a unique session ID (SUID) using UUID4.
+
+    Returns:
+        A unique session ID.
+    """
+    return str(uuid.uuid4())
+
+
+def check_suid(suid: str) -> bool:
+    """Check if given string is a proper session ID.
+
+    Args:
+        suid: The string to check.
+
+    Returns True if the string is a valid UUID, False otherwise.
+    """
+    try:
+        # accepts strings and bytes only
+        uuid.UUID(suid)
+        return True
+    except (ValueError, TypeError):
+        return False