lemonade-sdk 8.1.10__py3-none-any.whl → 8.1.11__py3-none-any.whl

lemonade/tools/server/serve.py

@@ -10,11 +10,13 @@ import traceback
 from typing import Optional, Union
 import json
 from pathlib import Path
+import os
 
-from fastapi import FastAPI, HTTPException, status, Request
+from fastapi import FastAPI, HTTPException, status, Request, WebSocket
 from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
+from starlette.websockets import WebSocketDisconnect, WebSocketState
 import uvicorn
 from uvicorn.config import Config
 from uvicorn.server import Server as UvicornServer
@@ -48,6 +50,7 @@ from openai.types.responses import (
 import lemonade.api as lemonade_api
 from lemonade.tools.server.wrapped_server import WrappedServer
 from lemonade.tools.server.llamacpp import LlamaServer
+from lemonade.tools.server.flm import FlmServer
 from lemonade.tools.server.tool_calls import extract_tool_calls, get_tool_call_pattern
 from lemonade.tools.server.webapp import get_webapp_html
 from lemonade.tools.server.utils.port import lifespan
@@ -75,12 +78,62 @@ from lemonade_server.settings import save_setting
 # Tests should use the max_new_tokens argument to set a lower value
 DEFAULT_MAX_NEW_TOKENS = 1500
 
-# Only import tray on Windows
-if platform.system() == "Windows":
+if platform.system() in ["Windows", "Darwin"]:
     # pylint: disable=ungrouped-imports
     from lemonade.tools.server.tray import LemonadeTray, OutputDuplicator
 
 
+class WebsocketTextFilter(logging.Filter):
+    def filter(self, record: logging.LogRecord) -> bool:
+        # Only allow logs that don't include "> TEXT"
+        return "> TEXT" not in record.getMessage()
+
+
+async def log_streamer(websocket: WebSocket, path: str, interval: float = 1.0):
+    logger = logging.getLogger()
+    await websocket.accept()
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            f.seek(0, os.SEEK_END)  # start at end
+            while True:
+                # Try reading a line
+                line = f.readline()
+                if not line:
+                    await asyncio.sleep(interval)
+                    continue
+
+                # Send defensively: if disconnected, bail out
+                if websocket.application_state != WebSocketState.CONNECTED:
+                    # Server-side state says we're not connected anymore
+                    break
+
+                try:
+                    await websocket.send_text(line)
+                except WebSocketDisconnect:
+                    # Client closed — normal path out
+                    break
+                except RuntimeError as re:
+                    # Starlette will raise this if a close has already been sent
+                    logger.debug("RuntimeError during send: %s", re)
+                    break
+
+    except WebSocketDisconnect:
+        # Client closed the socket; do not try to send or close again
+        pass
+    except Exception as e:  # pylint: disable=broad-except
+        # Log server-side; do not attempt to send error over a possibly closed socket
+        logger.exception("Error in log_streamer: %s", e)
+    finally:
+        # Only close if Starlette still thinks we're connected.
+        # This prevents "Cannot call send once a close message has been sent."
+        try:
+            if websocket.application_state == WebSocketState.CONNECTED:
+                await websocket.close()
+        except Exception:  # pylint: disable=broad-except
+            # If close itself races, swallow — we're shutting down anyway.
+            pass
+
+
 class ServerModel(Model):
     """
     An extension of OpenAI's Model class that adds
@@ -204,6 +257,12 @@ class Server:
             allow_headers=["*"],  # Allows all headers
         )
 
+        # Set up debug middleware if debug logging is enabled
+        # This must be done during app initialization, not at runtime
+        self.debug_logging_enabled = log_level == "debug"
+        if self.debug_logging_enabled:
+            self.setup_middleware_timer()
+
         # Set up custom routes
         self.setup_routes(["/api/v0", "/api/v1"])
 
@@ -265,6 +324,7 @@ class Server:
             self.app.post(f"{prefix}/completions")(self.completions)
             self.app.post(f"{prefix}/responses")(self.responses)
             self.app.post(f"{prefix}/log-level")(self.set_log_level)
+            self.app.websocket(f"{prefix}/logs/ws")(self.logs_ws)
 
             # OpenAI-compatible routes
             self.app.post(f"{prefix}/chat/completions")(self.chat_completions)
@@ -394,11 +454,13 @@ class Server:
             )
             file_handler.setLevel(logging_level)
             file_handler.setFormatter(uvicorn_formatter)
+            file_handler.addFilter(WebsocketTextFilter())
 
             # Set up console handler
             console_handler = logging.StreamHandler()
             console_handler.setLevel(logging_level)
             console_handler.setFormatter(uvicorn_formatter)
+            console_handler.addFilter(WebsocketTextFilter())
 
             # Configure root logger with both handlers
             logging.basicConfig(
@@ -421,10 +483,6 @@ class Server:
             ).run()
             sys.exit(0)
 
-        if self.debug_logging_enabled:
-            # Print the elapsed time for each request
-            self.setup_middleware_timer()
-
         # Let the app know what port it's running on, so
         # that the lifespan can access it
         self.app.port = self.port
@@ -521,7 +579,9 @@ class Server:
 
         return lc
 
-    async def completions(self, completion_request: CompletionRequest):
+    async def completions(
+        self, completion_request: CompletionRequest, request: Request
+    ):
         """
         Stream completion responses using HTTP chunked transfer encoding.
         """
@@ -534,7 +594,7 @@ class Server:
         # Load the model if it's different from the currently loaded one
         await self.load_llm(lc)
 
-        if self.llm_loaded.recipe == "llamacpp":
+        if self.llm_loaded.recipe == "llamacpp" or self.llm_loaded.recipe == "flm":
             return self.wrapped_server.completion(completion_request)
 
         # Check if the model supports reasoning
@@ -573,29 +633,43 @@ class Server:
                 # This is necessary because the variable is modified
                 # in the inner function
                 nonlocal reasoning_first_token
+                try:
+                    async for token in self._generate_tokens(**generation_args):
+                        # Handle client disconnect: stop generation and exit
+                        if await request.is_disconnected():
+                            self.stop_event.set()
+                            break
 
-                async for token in self._generate_tokens(**generation_args):
-                    choice = CompletionChoice(
-                        text=("<think>" + token if reasoning_first_token else token),
-                        index=0,
-                        finish_reason="stop",
-                        logprobs=None,
-                    )
+                        choice = CompletionChoice(
+                            text=(
+                                "<think>" + token if reasoning_first_token else token
+                            ),
+                            index=0,
+                            finish_reason="stop",
+                            logprobs=None,
+                        )
 
-                    completion = Completion(
-                        id="0",
-                        choices=[choice],
-                        model=self.llm_loaded.checkpoint,
-                        object="text_completion",
-                        created=int(time.time()),
-                    )
+                        completion = Completion(
+                            id="0",
+                            choices=[choice],
+                            model=self.llm_loaded.checkpoint,
+                            object="text_completion",
+                            created=int(time.time()),
+                        )
 
-                    # Format as SSE
-                    reasoning_first_token = False
-                    yield f"data: {completion.model_dump_json()}\n\n".encode("utf-8")
+                        # Format as SSE
+                        reasoning_first_token = False
+                        yield f"data: {completion.model_dump_json()}\n\n".encode(
+                            "utf-8"
+                        )
 
-                # Send the [DONE] marker
-                yield b"data: [DONE]\n\n"
+                    # Send the [DONE] marker only if still connected
+                    if not await request.is_disconnected():
+                        yield b"data: [DONE]\n\n"
+                except asyncio.CancelledError:
+                    # Propagate cancellation to the generator loop
+                    self.stop_event.set()
+                    return
 
             return StreamingResponse(
                 generate(),
@@ -653,7 +727,9 @@ class Server:
                 created=int(time.time()),
             )
 
-    async def chat_completions(self, chat_completion_request: ChatCompletionRequest):
+    async def chat_completions(
+        self, chat_completion_request: ChatCompletionRequest, request: Request
+    ):
         """
         Stream chat completion responses using HTTP chunked transfer encoding.
         """
@@ -669,7 +745,7 @@ class Server:
         # Load the model if it's different from the currently loaded one
         await self.load_llm(lc)
 
-        if self.llm_loaded.recipe == "llamacpp":
+        if self.llm_loaded.recipe == "llamacpp" or self.llm_loaded.recipe == "flm":
             return self.wrapped_server.chat_completion(chat_completion_request)
 
         # Convert chat messages to text using the model's chat template
@@ -731,69 +807,80 @@ class Server:
 
                 # Keep track of the full response for tool call extraction
                 full_response = ""
+                try:
+                    async for token in self._generate_tokens(**generation_args):
+                        # Handle client disconnect: stop generation and exit
+                        if await request.is_disconnected():
+                            self.stop_event.set()
+                            break
 
-                async for token in self._generate_tokens(**generation_args):
-                    # Continuously look for tool calls embedded into the generated text
-                    openai_tool_calls = None
-                    if chat_completion_request.tools:
+                        # Continuously look for tool calls embedded into the generated text
+                        openai_tool_calls = None
+                        if chat_completion_request.tools:
 
-                        # Append the token to the full response
-                        full_response += token
+                            # Append the token to the full response
+                            full_response += token
 
-                        tool_calls, _ = extract_tool_calls(
-                            full_response,
-                            tool_call_pattern,
-                        )
+                            tool_calls, _ = extract_tool_calls(
+                                full_response,
+                                tool_call_pattern,
+                            )
 
-                        # If there are tool calls, reset the full response for the next tool call
-                        if tool_calls:
-                            openai_tool_calls = []
-                            full_response = ""
-                        for tool_call in tool_calls:
-                            openai_tool_calls.append(
-                                ChoiceDeltaToolCall(
-                                    index=0,
-                                    id="-",
-                                    function=ChoiceDeltaToolCallFunction(
-                                        arguments=json.dumps(tool_call["arguments"]),
-                                        name=tool_call["name"],
-                                    ),
-                                    type="function",
+                            # If there are tool calls, reset the full response for the next call
+                            if tool_calls:
+                                openai_tool_calls = []
+                                full_response = ""
+                            for tool_call in tool_calls:
+                                openai_tool_calls.append(
+                                    ChoiceDeltaToolCall(
+                                        index=0,
+                                        id="-",
+                                        function=ChoiceDeltaToolCallFunction(
+                                            arguments=json.dumps(
+                                                tool_call["arguments"]
+                                            ),
+                                            name=tool_call["name"],
+                                        ),
+                                        type="function",
+                                    )
                                 )
-                            )
 
-                    # Create a ChatCompletionChunk
-                    chunk = ChatCompletionChunk.model_construct(
-                        id="0",
-                        object="chat.completion.chunk",
-                        created=int(time.time()),
-                        model=self.llm_loaded.checkpoint,
-                        choices=[
-                            Choice.model_construct(
-                                index=0,
-                                delta=ChoiceDelta(
-                                    content=(
-                                        "<think>" + token
-                                        if reasoning_first_token
-                                        else token
+                        # Create a ChatCompletionChunk
+                        chunk = ChatCompletionChunk.model_construct(
+                            id="0",
+                            object="chat.completion.chunk",
+                            created=int(time.time()),
+                            model=self.llm_loaded.checkpoint,
+                            choices=[
+                                Choice.model_construct(
+                                    index=0,
+                                    delta=ChoiceDelta(
+                                        content=(
+                                            "<think>" + token
+                                            if reasoning_first_token
+                                            else token
+                                        ),
+                                        function_call=None,
+                                        role="assistant",
+                                        tool_calls=openai_tool_calls,
+                                        refusal=None,
                                     ),
-                                    function_call=None,
-                                    role="assistant",
-                                    tool_calls=openai_tool_calls,
-                                    refusal=None,
-                                ),
-                                finish_reason=None,
-                                logprobs=None,
-                            )
-                        ],
-                    )
+                                    finish_reason=None,
+                                    logprobs=None,
+                                )
+                            ],
+                        )
 
-                    # Format as SSE
-                    reasoning_first_token = False
-                    yield f"data: {chunk.model_dump_json()}\n\n".encode("utf-8")
+                        # Format as SSE
+                        reasoning_first_token = False
+                        yield f"data: {chunk.model_dump_json()}\n\n".encode("utf-8")
 
-                # Send the [DONE] marker
-                yield b"data: [DONE]\n\n"
+                    # Send the [DONE] marker only if still connected
+                    if not await request.is_disconnected():
+                        yield b"data: [DONE]\n\n"
+                except asyncio.CancelledError:
+                    self.stop_event.set()
+                    return
 
             return StreamingResponse(
                 generate(),
@@ -952,7 +1039,7 @@ class Server:
             formatted_messages.append(f"{role_marker}\n{content} <|end|>")
         return "\n".join(formatted_messages) + "\n<|assistant|>"
 
-    async def responses(self, responses_request: ResponsesRequest):
+    async def responses(self, responses_request: ResponsesRequest, request: Request):
         """
         Stream responses using HTTP chunked transfer encoding.
         """
@@ -965,6 +1052,12 @@ class Server:
         # Load the model if it's different from the currently loaded one
         await self.load_llm(lc)
 
+        if self.llm_loaded.recipe == "llamacpp":
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=f"Responses API not supported for recipe: {self.llm_loaded.recipe}",
+            )
+
         # Convert chat messages to text using the model's chat template
         if isinstance(responses_request.input, str):
             text = responses_request.input
@@ -1018,56 +1111,71 @@ class Server:
 
                 full_response = "<think>" if reasoning_first_token else ""
 
-                async for token in self._generate_tokens(**generation_args):
+                try:
+                    async for token in self._generate_tokens(**generation_args):
+                        # Handle client disconnect: stop generation and exit
+                        if await request.is_disconnected():
+                            self.stop_event.set()
+                            break
 
-                    # Create an event
-                    delta_event = ResponseTextDeltaEvent(
-                        content_index=0,
-                        delta=("<think>" + token if reasoning_first_token else token),
-                        item_id="0 ",
-                        output_index=0,
-                        type="response.output_text.delta",
-                        sequence_number=0,
-                    )
-                    full_response += token
+                        # Create an event
+                        delta_event = ResponseTextDeltaEvent(
+                            content_index=0,
+                            delta=(
+                                "<think>" + token if reasoning_first_token else token
+                            ),
+                            item_id="0 ",
+                            output_index=0,
+                            type="response.output_text.delta",
+                            sequence_number=0,
+                        )
+                        full_response += token
 
-                    # Format as SSE
-                    reasoning_first_token = False
-                    yield f"data: {delta_event.model_dump_json()}\n\n".encode("utf-8")
+                        # Format as SSE
+                        reasoning_first_token = False
+                        yield f"data: {delta_event.model_dump_json()}\n\n".encode(
+                            "utf-8"
+                        )
 
-                # Send the completed event
-                response_output_message = ResponseOutputMessage(
-                    id="0",
-                    content=[
-                        ResponseOutputText(
-                            annotations=[],
-                            text=full_response,
-                            type="output_text",
+                    # Send the completed event (only if still connected)
+                    if not await request.is_disconnected():
+                        response_output_message = ResponseOutputMessage(
+                            id="0",
+                            content=[
+                                ResponseOutputText(
+                                    annotations=[],
+                                    text=full_response,
+                                    type="output_text",
+                                )
+                            ],
+                            role="assistant",
+                            status="completed",
+                            type="message",
+                        )
+                        response = Response(
+                            id="0",
+                            model=self.llm_loaded.checkpoint,
+                            created_at=int(time.time()),
+                            object="response",
+                            output=[response_output_message],
+                            parallel_tool_calls=True,
+                            tool_choice="auto",
+                            tools=[],
+                        )
+                        completed_event = ResponseCompletedEvent(
+                            response=response,
+                            type="response.completed",
+                            sequence_number=0,
+                        )
+                        yield f"data: {completed_event.model_dump_json()}\n\n".encode(
+                            "utf-8"
                         )
-                    ],
-                    role="assistant",
-                    status="completed",
-                    type="message",
-                )
-                response = Response(
-                    id="0",
-                    model=self.llm_loaded.checkpoint,
-                    created_at=int(time.time()),
-                    object="response",
-                    output=[response_output_message],
-                    parallel_tool_calls=True,
-                    tool_choice="auto",
-                    tools=[],
-                )
-                completed_event = ResponseCompletedEvent(
-                    response=response,
-                    type="response.completed",
-                    sequence_number=0,
-                )
-                yield f"data: {completed_event.model_dump_json()}\n\n".encode("utf-8")
 
-                # Send the [DONE] marker
-                yield b"data: [DONE]\n\n"
+                        # Send the [DONE] marker
+                        yield b"data: [DONE]\n\n"
+                except asyncio.CancelledError:
+                    self.stop_event.set()
+                    return
 
             return StreamingResponse(
                 generate(),
@@ -1312,8 +1420,10 @@ class Server:
         """
         Send performance statistics to the client.
         """
-        # If using llama server, get telemetry from the telemetry instance
-        if self.llm_loaded and self.llm_loaded.recipe == "llamacpp":
+        # If using wrapped server, get telemetry from the telemetry instance
+        if self.llm_loaded and (
+            self.llm_loaded.recipe == "llamacpp" or self.llm_loaded.recipe == "flm"
+        ):
             return self.wrapped_server.telemetry.get_telemetry_data()
 
         # For built-in server, use the existing telemetry
@@ -1414,6 +1524,7 @@ class Server:
             checkpoint=config.checkpoint,
             recipe=config.recipe,
             reasoning=config.reasoning,
+            vision=config.vision,
             mmproj=config.mmproj,
             # The pull endpoint will download an upgraded model if available, even
             # if we already have a local copy of the model
@@ -1493,8 +1604,8 @@ class Server:
             ):
                 if (
                     self.llm_loaded.recipe == "llamacpp"
-                    and self.wrapped_server.process.poll()
-                ):
+                    or self.llm_loaded.recipe == "flm"
+                ) and self.wrapped_server.process.poll():
                     # wrapped server process has gone away for some reason, so we should
                     # proceed with loading to get it back
                     pass
@@ -1518,6 +1629,14 @@ class Server:
                         do_not_upgrade=True,
                     )
 
+                elif config_to_use.recipe == "flm":
+                    self.wrapped_server = FlmServer()
+                    self.wrapped_server.load(
+                        model_config=config_to_use,
+                        ctx_size=self.ctx_size,
+                        do_not_upgrade=True,
+                    )
+
                 else:
                     self.model, self.tokenizer = lemonade_api.from_pretrained(
                         checkpoint=config_to_use.checkpoint, recipe=config_to_use.recipe
@@ -1554,7 +1673,7 @@ class Server:
                 for _ in range(self.max_concurrent_generations):
                     await self._generate_semaphore.acquire()
 
-            if self.llm_loaded.recipe == "llamacpp":
+            if self.llm_loaded.recipe == "llamacpp" or self.llm_loaded.recipe == "flm":
                 self.wrapped_server.process.terminate()
 
             self.llm_loaded = None
@@ -1657,6 +1776,12 @@ class Server:
                     logging.debug(f"Total request time: {request_time:.4f} seconds")
             return response
 
+    async def logs_ws(self, websocket: WebSocket):
+        if not self.log_file or not os.path.exists(self.log_file):
+            await websocket.close(code=4000)
+            return
+        await log_streamer(websocket, self.log_file)
+
 
 # This file was originally licensed under Apache 2.0. It has been modified.
 # Modifications Copyright (c) 2025 AMD