Streaming LlamaCPP in Langserve or FastAPI

[weissenbacherpwc]

Hi,

I have built a rag app and I am loading a LLM with Llamacpp. However I have problems with making Streaming work for FastAPI or Langserve requests. Streaming is working in my Terminal, but I don't know what I have to change to make it work in FastAPI/Langserve.

Here is my Langserve code:

from langchain_community.vectorstores.pgvector import PGVector
from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
from langchain_core.runnables import RunnableParallel
import os
from langchain_community.embeddings import HuggingFaceBgeEmbeddings, HuggingFaceEmbeddings
import box 
import yaml
from langchain_community.llms import LlamaCpp
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from operator import itemgetter
from typing import TypedDict
from fastapi import FastAPI
from fastapi.responses import RedirectResponse
from langserve import add_routes
from fastapi.middleware.cors import CORSMiddleware
from starlette.staticfiles import StaticFiles
from langchain_core.output_parsers import StrOutputParser

with open('./config/config.yml', 'r', encoding='utf8') as ymlfile:
    cfg = box.Box(yaml.safe_load(ymlfile))
    
def build_llm(model_path, temperature=cfg.RAG_TEMPERATURE, max_tokens=cfg.MAX_TOKENS, callback = StreamingStdOutCallbackHandler()):
        
        callback_manager = CallbackManager([callback])
        
        n_gpu_layers = 1 # Metal set to 1 is enough. # ausprobiert mit mehreren
        n_batch = 512 #1024 Should be between 1 and n_ctx, consider the amount of RAM of your Apple Silicon Chip.

        llm = LlamaCpp(
                max_tokens = max_tokens,
                n_threads = 8,#8, #für performance,
                model_path=model_path,
                temperature=temperature,
                f16_kv=True,
                n_ctx=15000, # 8k aber mann muss Platz lassen für Instruction, History etc. 
                n_gpu_layers=n_gpu_layers,
                n_batch=n_batch,
                callback_manager=callback_manager, 
                verbose=True, # Verbose is required to pass to the callback manager
                top_p=0.75,
                top_k=40,
                repeat_penalty = 1.1,
                streaming=True,
                model_kwargs={
                        #'repetition_penalty': 1.1,
                        #'mirostat': 2,
                },
        )
        
        return llm

embeddings = HuggingFaceEmbeddings(model_name=cfg.EMBEDDING_MODEL_NAME,
                                            model_kwargs={'device': 'mps'})

PG_COLLECTION_NAME =  "PGVECTOR_BKB"
model_path = "./modelle/sauerkrautlm-mixtral-8x7b-instruct.Q4_K_M.gguf"

CONNECTION_STRING = "MY_CONNTECTION_STRING"
vector_store = PGVector(
    collection_name=PG_COLLECTION_NAME,
    connection_string=CONNECTION_STRING,
    embedding_function=embeddings
)

prompt= """
<s> [INST] Du bist RagBot, ein hilfsbereiter Assistent. Antworte nur auf Deutsch. Verwende die folgenden Kontextinformationen, um die Frage am Ende knapp zu beantworten. Wenn du die Antwort nicht kennst, sag einfach, dass du es nicht weisst. Erfinde keine Antwort! Falls der Nutzer allgemeine Fragen stellt, führe Smalltalk mit Ihm.

### Hier der Kontext: ###
{context}

### Hier die Frage: ###
{question}

Antwort: [/INST]
"""

def model_response_prompt():
    return PromptTemplate(template=prompt, input_variables=['input', 'typescript_string'])
prompt_temp = model_response_prompt()

llm = build_llm(model_path, temperature= cfg.NO_RAG_TEMPERATURE, max_tokens = cfg.NO_RAG_MAX_TOKENS)

class RagInput(TypedDict):
    question: str


final_chain = (
        RunnableParallel(
            context=(itemgetter("question") | vector_store.as_retriever()),
            question=itemgetter("question")
        ) |
        RunnableParallel(
            answer=(prompt_temp| llm),
            docs=itemgetter("context")
        )

).with_types(input_type=RagInput)

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "http://localhost:3000"
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

#app.mount("/rag/static", StaticFiles(directory="./source_docs"), name="static")
@app.get("/")
async def redirect_root_to_docs():
    return RedirectResponse("/docs")


# Edit this to add the chain you want to add
add_routes(app, final_chain, path="/rag")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

When trying out in the Langserve Playground (http://0.0.0.0:8000/rag/playground/) the response gets streamed in my Terminal but not in the Playground.

So how can I make this work?

[Gitclop]

I've just build something similar, maybe this will help as a framework for you?

# App definition
app = FastAPI(
  title="LangChain Server",
  version="1.0",
  description="A simple API server using LangChain's Runnable interfaces",
)

# Needed for authentication
app.add_middleware(
    CORSMiddleware,
    allow_origins=['*'],
    allow_methods=['*'],
    allow_headers=['*'],
    allow_credentials=True
)

class ChatInput(BaseModel):
    question: str

class ChatOutput(BaseModel):
    response: str

class MultiChatInput(BaseModel):
    prompts: Dict[str, str]

@app.post("/multi_chat")
async def multi_chat_endpoint(chat_inputs: MultiChatInput):
    try:
        system_prompt = "System Prompt"
        template_parameter = "ChatML"

        responses = {}
        model_path = 'TheBloke/DiscoLM_German_7b_v1-GGUF'
        my_engine = InferenceEngine.Engine(selected_model=model_path, network_serving=False)
        llm_chain = my_engine.create_chain(system_prompt, template_parameter)

        for key, prompt in chat_inputs.prompts.items():
            response = str(llm_chain.invoke({"question": prompt})).lstrip()
            responses[key] = response

    except Exception as e:
        return {"error": str(e)}

    return responses

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="localhost", port=8000)
    
class Engine:
    def __init__(self, selected_model,network_serving=False):

        if not network_serving:
            use_proxy.set_proxy()

            if selected_model == 'TheBloke/DiscoLM_German_7b_v1-GGUF':

                (repo_id, model_file_name) = ("TheBloke/DiscoLM_German_7b_v1-GGUF",
                                              "discolm_german_7b_v1.Q6_K.gguf")
                ...
        

            model_path = hf_hub_download(repo_id=repo_id,
                                         filename=model_file_name,
                                         repo_type="model")
        else:
            model_path = selected_model

        # initialize LlamaCpp LLM model
        self.llm = LlamaCpp(
                    model_path=model_path,
                    seed=2,
                    temperature=0.5,  # Default 0.8
                    max_tokens=256,  # The maximum number of tokens to generate. Default 256
                    top_k=20,  # Default 40
                    top_p=0.85, # Default 0.95
                    n_ctx=1024, # Text context window, 0 = from model
                    n_gpu_layers=64,
                    n_batch=1024,  # Prompt processing maximum batch size Default 512
                    stop = ['</s>'],
                    #n_threads=8,  # Number of threads to use for generation --8 Am schnellsten!
                    #callback_manager=callbacks,
                    verbose=True,  # Verbose is required to pass to the callback manager
                    streaming=False  # Whether to stream the results, token by token.
                )

    def create_chain(self, system_prompt, template_parameter):

        if template_parameter == 'ChatML':
            chatML_template = f"""<|im_start|>system
                                    {system_prompt}<|im_end|>
                                    <|im_start|>user
                                    {{question}}<|im_end|>
                                    <|im_start|>assistant
                                    """
            template = chatML_template

        if template_parameter == 'Alpaca':
            alpaca_template = f"""USER: {system_prompt} USER: {{question}} ASSISTANT:"""

            template = alpaca_template

        if template_parameter == 'Instruct':
            instruct_template = f"""<s>[INST]{system_prompt}[/INST]<s>[INST]{{question}}[/INST]"""
            template = instruct_template

        prompt = PromptTemplate(template=template, input_variables=["question"])

        # We create an llm chain with our LLM and prompt
        # llm_chain = LLMChain(prompt=prompt, llm=llm) # Legacy
        llm_chain = prompt | self.llm  # LCEL

        return llm_chain

[Gitclop]

Oh, and i call it like this:

import requests

api_endpoint = "http://localhost:8000/multi_chat"

def send_chat_request(prompts):
    response = requests.post(api_endpoint, json={"prompts": prompts})
    return response.json()

prompts = {"title": title_prompt, "question": question_prompt, "answer": answer_prompt}
responses = send_chat_request(prompts)

title_response = responses.get('title', 'Keine Antwort erhalten')
question_response = responses.get('question', 'Keine Antwort erhalten')
answer_response = responses.get('answer', 'Keine Antwort erhalten')

[weissenbacherpwc]

thanks, so you are calling Class engine after if __name__ == "__main__":? Or is this a new script?

[Gitclop]

Sorry for the confusion. Yes they are two seperate scripts:
InferenceEngine.py handles model and chain initialisation.
A seperate Server.py starts up the server.
The code from the edit is called from a third script (from the client)

[weissenbacherpwc]

got it thanks @Gitclop .
I also made my own script running now and the RetrievalQA chain response gets streamed now. But the response format is now not a dict anymore but a blank text. How can I return a dict so that I can give the user information about the retrieved context (sources and page number)?

from fastapi import FastAPI
import box
import yaml
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from langchain_community.llms import LlamaCpp
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.chains import RetrievalQA
from fastapi.middleware.cors import CORSMiddleware
from functools import lru_cache
import uvicorn
from src.llm import build_llm
from src.utils import get_gpu_device
# Import config vars
with open('config/config.yml', 'r', encoding='utf8') as ymlfile:
    cfg = box.Box(yaml.safe_load(ymlfile))
# caching Embeddings
@lru_cache(maxsize=100)
def get_cached_embeddings():
        embeddings = HuggingFaceEmbeddings(model_name=cfg.EMBEDDING_MODEL_NAME,
                                       model_kwargs={'device': device}, encode_kwargs={'device': device}) #'batch_size': 32
    
        return embeddings

embeddings = get_cached_embeddings()

vectorstore = FAISS.load_local(cfg.BKB_FAISS_PATH, embeddings)
retriever = vectorstore.as_retriever(search_kwargs={'k': cfg.VECTOR_COUNT, 'score_threshold': cfg.SCORE_THRESHOLD}, search_type="similarity_score_threshold")

def build_retrieval_qa(llm, prompt, vectordb):
    chain_type_kwargs={ "prompt": prompt,
                        "verbose": True
                        }
    dbqa = RetrievalQA.from_chain_type(llm=llm,
                                        chain_type='stuff',
                                        retriever=vectordb,
                                        return_source_documents=True,
                                        chain_type_kwargs=chain_type_kwargs,
                                        verbose=True
                                        )
    return dbqa
llm = build_llm(model_path)
def stream_with_context(question: str, task: str):
    rag_prompt = f"""
    <s> [INST] Als dein Assistent liegt meine Aufgabe darin, dir basierend auf dem User Input und dem verfügbaren Kontext präzise und verständliche Antworten zu liefern. Bitte befolge folgende Reihenfolge Schritt für Schritt für eine optimale Antwort:

    1. **Kontextanalyse**: Lies und verstehe den Kontext gründlich, um die Informationen zu identifizieren, die für die Beantwortung relevant sein könnten.

    2. **Aufgabenverständnis**: Kläre die Aufgabenstellung. Was genau wird erwartet und welche spezifischen Informationen werden benötigt?

    3. **Analyse des User Inputs**: Untersuche den User Input genau. Identifiziere die Kernfrage oder das Anliegen des Nutzers.

    4. **Antwortformulierung**: Entwickle eine Antwort, die direkt und präzise auf den User Input eingeht, unter Verwendung einer klaren und verständlichen Sprache.

    5. **Integration von Kontextinformationen**: Beziehe relevante Kontextinformationen ein, aber nur, wenn du ihrer Relevanz und Genauigkeit sicher bist. Bei Unsicherheit über die Relevanz oder Genauigkeit, verzichte auf die Einbeziehung, um die Verbreitung von Fehlinformationen zu verhindern.

    6. **Keine Wiederholung von Texten**: Achte darauf, dass bereits geschriebene Texte nicht wiederholt werden! Die Wiederholung eines Absatzes ist schlecht für die User Experience.

    ### Im Folgenden steht dir der Kontext, die Aufgabenstellung und der User Input zur Verfügung: ###

    **Kontext zum User Input:**
    {{context}}


    **Aufgabenstellung:**
    {task}


    **User Input:**
    {{question}}


    **Antwort (wiederhole keine Absätze):**"""
    rag_prompt = rag_prompt.format(typescript_string=task)
    qa_prompt = PromptTemplate(template=rag_prompt, input_variables=['typescript_string', 'context', 'question'])
    model = build_retrieval_qa(llm=llm,prompt=qa_prompt,vectordb=retriever)
    for res in model.stream({"query":question}):
        result = copy.deepcopy(res)
        yield result
app = FastAPI(
    title="Inference API for Mistral and Mixtral",
    description="A simple API that use Mistral or Mixtral",
    version="1.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
@app.get("/stream_with_context")
def get_stream_with_context(question: str, task: str):
    return StreamingResponse(model_streamer(question, task), media_type='text/event-stream')