LLM Inference
=============

If you want to perform LLM inference through a REST API, you can send a POST request to ``/api/generate``.
This is great when you're building your own application.

.. note::
    LLM inference is supported for `Ollama` and `Hugging Face`\ [#]_ models. In the case of `Ollama`, please ensure that it is installed on your
    OS and that the desired model is pulled before starting a conversation.

    - Install Ollama: https://ollama.com/download
    - Pull model: For e.g., ``ollama pull llama3.1:8b`` on your terminal

.. admonition:: Running the code
    :class: warning

    The code below must be run as a Python script or executed in a terminal using ``python3``. It will not work inside notebooks.

.. code-block:: python

    # Import langformers
    from langformers import tasks

    # Create a generator
    generator = tasks.create_generator(provider="ollama", model_name="llama3.1:8b")

    # Run the generator
    generator.run(host="0.0.0.0", port=8000)


.. tabs::

    .. tab:: create_generator()

        .. autofunction:: langformers.tasks.create_generator
           :no-index:

    .. tab:: run()

        ``run()`` takes the following parameters:

        - ``host`` (str, default="0.0.0.0"): The IP address to bind the server to.
        - ``port`` (int), default=8000: The port number to listen on.



API Request Format
--------------------

The ``/api/generate`` endpoint accepts the following:

.. code-block:: json

    {
        "system_prompt": "You are an Aussie AI assistant, reply in an Aussie way.",
        "memory_k": 10,
        "temperature": 0.5,
        "top_p": 1,
        "max_length": 5000,
        "prompt": "Hi"
    }

Note: The endpoint expects at least ``prompt``. The keys ``system_prompt``, ``memory_k``, ``temperature``, ``top_p``, ``max_length`` are optional.

Here's an example of making a request to the API using Python:

.. code-block:: python

    # Imports
    import requests
    import json

    # Endpoint URL
    url = "http://0.0.0.0:8000/api/generate"

    # Define payload
    payload = json.dumps({
        "system_prompt": "You are an Aussie AI assistant, reply in an Aussie way.",
        "memory_k": 10,
        "temperature": 0.5,
        "top_p": 1,
        "max_length": 5000,
        "prompt": "Hi"
    })

    # Headers
    headers = {
        "Content-Type": "application/json",
    }

    # Send request
    response = requests.post(url, headers=headers, data=payload)

    # Print response
    print(response.text)

This streams the tokens generated by the LLM using SSE streams (e.g., data: {"chunk": "Hello"} …). You need to parse these SSE streams. Langformers can handle this natively.

StreamProcessor
----------------

Here's how you parse the SSE streams with StreamProcessor.

.. code-block:: python

    # Import StreamProcessor
    from langformers.generators import StreamProcessor

    # Define headers
    headers = {
        "Content-Type": "application/json",
    }

    # Create an object
    client = StreamProcessor(headers=headers)

    # Define payload
    payload = {
        "system_prompt": "You are an Aussie AI assistant, reply in an Aussie way.",
        "memory_k": 10,
        "temperature": 0.5,
        "top_p": 1,
        "max_length": 5000,
        "prompt": "Hi, how are you today",
    }

    # Send request
    response = client.process(endpoint_url="http://0.0.0.0:8000/api/generate", payload=payload)

    # Print response
    for chunk in response:
        print(chunk, end="", flush=True)


.. tabs::

    .. tab:: /api/generate

        The ``/api/generate`` endpoint takes the following parameters:

        - ``system_prompt`` (str, default=<Langformers.commons.prompts default_chat_prompt_system>): The system-level instruction for the LLM.
        - ``memory_k`` (int, default=10): The number of previous messages to retain in memory.
        - ``temperature`` (float, default=0.5): Controls randomness of responses (higher = more random).
        - ``top_p`` (float, default=1): Nucleus sampling parameter (lower = more focused).
        - ``max_length`` (int, default=5000): Maximum number of tokens to generate.
        - ``prompt`` (str, required): User query.

        .. admonition:: System prompt and Memory
            :class: warning

            Note that any change in ``system_prompt`` clears the previous conversations stored in memory.

    .. tab:: StreamProcessor()

        .. autofunction:: langformers.generators.StreamProcessor.__init__
            :no-index:

    .. tab:: process()

        .. autofunction:: langformers.generators.StreamProcessor.process
            :no-index:

Authentication
-----------------
Securing the ``/api/generate`` endpoint is straightforward. You can pass a dependency function to ``dependency`` when creating the generator.

.. code-block:: python

    async def auth_dependency():
        """Authorization dependency for request validation.

        - Implement your own logic here (e.g., API key check, authentication).
        - If the function returns a value, access is granted.
        - Raising an HTTPException will block access.
        """
        return True  # Modify this logic as needed

    generator = tasks.create_generator(provider="ollama", model_name="llama3.1:8b", dependency=auth_dependency)


**Example: Using API Key Authentication**

You can implement a simple authentication dependency like this:

.. code-block::

    # Imports
    from langformers import tasks
    from fastapi import Request, HTTPException

    # Define a set of valid API keys
    API_KEYS = {"12345", "67890"}

    async def auth_dependency(request: Request):
        """
        Extracts the Bearer token and verifies it against a list of valid API keys.
        """
        auth_header = request.headers.get("Authorization")

        if not auth_header or not auth_header.startswith("Bearer "):
            raise HTTPException(status_code=401, detail="Invalid authorization header format.")

        token = auth_header.split("Bearer ")[1]
        if token not in API_KEYS:
            raise HTTPException(status_code=401, detail="Unauthorized.")

        return True  # Allow access

    # Create a generator with authentication
    generator = tasks.create_generator(provider="ollama", model_name="llama3.1:8b", dependency=auth_dependency)

    # Run the generator
    generator.run(host="0.0.0.0", port=8000)

With this setup, only requests that include a valid API key in the headers will be authorized. All you need to do is include an ``Authorization: Bearer <token>`` header with one of the API keys as the token and make a POST request.

.. code-block:: python

    headers = {
        'Authorization': 'Bearer 12345',
        'Content-Type': 'application/json'
    }

.. warning::

    For industry-standard authentication in FastAPI, you can use OAuth2 with JWT (JSON Web Token), which is widely adopted for securing APIs.

**Footnotes**

.. [#] Hugging Face support is limited to chat-tuned models (instruct) that include a ``chat_template`` in their ``tokenizer_config.json`` and are compatible with the `transformers` library and your system's hardware.