Build streaming microservice with llama.cpp

[ajndkr]

Related #164

While I work on integrating this feature (need to create sync mode callback handlers for all), here's a script to build a microservice with llama.cpp:

import asyncio
from functools import partial
from typing import Any

from fastapi import Depends, status
from lanarky import Lanarky
from lanarky.adapters.langchain.callbacks import (
    TokenStreamingCallbackHandler as _TokenStreamingCallbackHandler,
)
from lanarky.adapters.langchain.callbacks import get_token_data
from lanarky.adapters.langchain.responses import StreamingResponse as _StreamingResponse
from lanarky.events import Events, ServerSentEvent, ensure_bytes
from lanarky.logging import logger
from lanarky.responses import HTTPStatusDetail
from lanarky.responses import StreamingResponse as _StreamingResponse
from langchain.chains import ConversationChain
from langchain.chains.base import Chain
from langchain.llms.llamacpp import LlamaCpp
from pydantic import BaseModel
from starlette.types import Send


class TokenStreamingCallbackHandler(_TokenStreamingCallbackHandler):
    async def on_llm_new_token(self, token: str, **kwargs: dict[str, Any]) -> None:
        """Run on new LLM token. Only available when streaming is enabled."""
        if not self.streaming:
            self.streaming = True

        if self.llm_cache_used:  # cache missed (or was never enabled) if we are here
            self.llm_cache_used = False

        message = self._construct_message(
            data=get_token_data(token, self.mode), event=Events.COMPLETION
        )
        task = asyncio.create_task(self._send(message))
        await task



class StreamingResponse(_StreamingResponse):
    """StreamingResponse class for LangChain resources."""

    def __init__(
        self,
        chain: Chain,
        config: dict[str, Any],
        *args: Any,
        **kwargs: dict[str, Any],
    ) -> None:
        """Constructor method.

        Args:
            chain: A LangChain instance.
            config: A config dict.
            *args: Positional arguments to pass to the parent constructor.
            **kwargs: Keyword arguments to pass to the parent constructor.
        """
        super().__init__(*args, **kwargs)

        self.chain = chain
        self.config = config

    async def stream_response(self, send: Send) -> None:
        """Stream LangChain outputs.

        If an exception occurs while iterating over the LangChain, an
        internal server error is sent to the client.

        Args:
            send: The ASGI send callable.
        """
        await send(
            {
                "type": "http.response.start",
                "status": self.status_code,
                "headers": self.raw_headers,
            }
        )

        if "callbacks" in self.config:
            for callback in self.config["callbacks"]:
                if hasattr(callback, "send"):
                    callback.send = send

        try:
            # TODO: migrate to `.ainvoke` when adding support
            # for LCEL
            loop = asyncio.get_event_loop()
            outputs = await loop.run_in_executor(
                None, partial(self.chain, **self.config)
            )
            if self.background is not None:
                self.background.kwargs.update({"outputs": outputs})
        except Exception as e:
            logger.error(f"chain runtime error: {e}")
            if self.background is not None:
                self.background.kwargs.update({"outputs": {}, "error": e})
            chunk = ServerSentEvent(
                data=dict(
                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                    detail=HTTPStatusDetail.INTERNAL_SERVER_ERROR,
                ),
                event=Events.ERROR,
            )
            await send(
                {
                    "type": "http.response.body",
                    "body": ensure_bytes(chunk, None),
                    "more_body": True,
                }
            )

        await send({"type": "http.response.body", "body": b"", "more_body": False})


app = Lanarky()


class ChatInput(BaseModel):
    input: str


def chain_factory() -> ConversationChain:
    return ConversationChain(
        llm=LlamaCpp(
            model_path="models/mistral-7b-instruct-v0.1.Q2_K.gguf",
            temperature=0.9,
            streaming=True,
        ),
        verbose=True,
    )


@app.post("/chat")
async def chat(request: ChatInput, chain: ConversationChain = Depends(chain_factory)):
    return StreamingResponse(
        chain=chain,
        config={
            "inputs": request.model_dump(),
            "callbacks": [
                TokenStreamingCallbackHandler(output_key=chain.output_key),
            ],
        },
    )


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app)

To start server, run:

python app.py

NOTE: change the model_path before running server.

To test:

curl -X 'POST' 'http://localhost:8000/chat'  \
-H 'accept: text/event-stream'  \
-H 'Content-Type: application/json'  \
-d '{"input": "hi"}'

Hope this helps while I work on updating the library!

[neokd]

Thanks a lot @ajndkr

[ajndkr]

@neokd you can now upgrade to v0.8.2 and use the below code for your llama.cpp application:

from fastapi import Depends
from langchain.chains import ConversationChain
from langchain.llms.llamacpp import LlamaCpp
from pydantic import BaseModel

from lanarky import Lanarky
from lanarky.adapters.langchain.callbacks import TokenStreamingCallbackHandler
from lanarky.adapters.langchain.responses import StreamingResponse

app = Lanarky()


class ChatInput(BaseModel):
    input: str


def chain_factory() -> ConversationChain:
    return ConversationChain(
        llm=LlamaCpp(
            model_path="<your-model-path>",
            temperature=0.9,
            streaming=True,
        ),
        verbose=True,
    )


@app.post("/chat")
async def chat(request: ChatInput, chain: ConversationChain = Depends(chain_factory)):
    return StreamingResponse(
        chain=chain,
        config={
            "inputs": request.model_dump(),
            "callbacks": [
                TokenStreamingCallbackHandler(output_key=chain.output_key),
            ],
        },
        run_mode="sync",
    )


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app)