implement formated prompt

app/chains.py

@@ -8,6 +8,7 @@ from langchain_core.runnables import RunnablePassthrough
 import schemas
 from prompts import (
     raw_prompt,
+    raw_prompt_formatted,
     format_context,
     tokenizer
 )
@@ -29,8 +30,8 @@ simple_chain = (raw_prompt | llm).with_types(input_type=schemas.UserQuestion)
 
 # data_indexer = DataIndexer()
 
-# # TODO: create formatted_chain by piping raw_prompt_formatted and the LLM endpoint.
-# formatted_chain = None
+# TODO: create formatted_chain by piping raw_prompt_formatted and the LLM endpoint.
+formatted_chain = (raw_prompt_formatted | llm).with_types(input_type=schemas.UserQuestion)
 
 # # TODO: use history_prompt_formatted and HistoryInput to create the history_chain
 # history_chain = None

app/main.py

@@ -10,7 +10,7 @@ from typing import List
 from sqlalchemy.orm import Session
 
 import schemas
-from chains import simple_chain
+from chains import simple_chain, formatted_chain
 import crud, models, schemas
 from database import SessionLocal, engine
 from callbacks import LogResponseCallback
@@ -28,24 +28,47 @@ def get_db():
         db.close()
 
 # ..
+# "async" marks the function as asynchronous, allowing it to pause and resume during operations like streaming or I/O.
 async def generate_stream(input_data: schemas.BaseModel, runnable: Runnable, callbacks: List[BaseCallbackHandler]=[]):
+    """generate_stream is an asynchronous generator that processes input data, 
+    streams output data from a runnable object, serializes each output, and yields 
+    it to the client in real-time as part of a server-sent event (SSE) stream.
+    It uses callbacks to customize the processing, serializes each piece of output 
+    using WellKnownLCSerializer, and indicates the end of the stream with a final “end” event.
+    """    
     for output in runnable.stream(input_data.dict(), config={"callbacks": callbacks}): 
         data = WellKnownLCSerializer().dumps(output).decode("utf-8")
-        yield {'data': data, "event": "data"} 
+        yield {'data': data, "event": "data"}
+    # After all the data has been streamed and the loop is complete, the function yields a final event to signal 
+    # the end of the stream. This sends an {"event": "end"} message to the client, letting them know that no more 
+    # data will be sent.
     yield {"event": "end"}
 
-
+# This registers the function simple_stream as a handler for HTTP POST requests at the URL endpoint /simple/stream. 
+# It means that when a client sends a POST request to this endpoint, this function will be triggered.
 @app.post("/simple/stream")
 async def simple_stream(request: Request):
+    """the function handles a POST request at the /simple/stream endpoint, 
+    extracts the JSON body, unpacks the "input" field, and then uses it to 
+    initialize a UserQuestion schema object (which performs validation 
+    and data transformation) and then initiates a server-sent event response 
+    to stream data back to the client based on the user’s question.
+    """    
+    # await is used because parsing the JSON may involve asynchronous I/O operations, 
+    # especially when handling larger payloads.
     data = await request.json()
     user_question = schemas.UserQuestion(**data['input'])
+    # This line returns an EventSourceResponse, which is typically used to handle server-sent events (SSE). 
+    # It’s a special kind of response that streams data back to the client in real time. 
     return EventSourceResponse(generate_stream(user_question, simple_chain))
 
 
-# @app.post("/formatted/stream")
-# async def formatted_stream(request: Request):
-#     # TODO: use the formatted_chain to implement the "/formatted/stream" endpoint.
-#     raise NotImplemented
+@app.post("/formatted/stream")
+async def formatted_stream(request: Request):
+    # TODO: use the formatted_chain to implement the "/formatted/stream" endpoint.
+    data = await request.json()
+    user_question = schemas.UserQuestion(**data['input'])
+    return EventSourceResponse(generate_stream(user_question, formatted_chain))
 
 
 # @app.post("/history/stream")

app/prompts.py

@@ -16,15 +16,29 @@ login(os.environ['HF_TOKEN'])
 tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
 
 def format_prompt(prompt) -> PromptTemplate:
-    # TODO: format the input prompt by using the model specific instruction template
-    # TODO: return a langchain PromptTemplate
-    raise NotImplemented
+    """this function prepares a well-formatted prompt template for interacting with a 
+    large language model, ensuring that the model has a clear role (AI assistant) 
+    and understands the user’s input.
+    It first 1. format the input prompt by using the model specific instruction template
+    2. return a langchain PromptTemplate
+    """    
+    chat = [
+        {"role": "system", "content": "You are a helpful AI assistant."},
+        {"role": "user", "content": prompt},
+    ]
+    formatted_prompt = tokenizer.apply_chat_template(
+        chat, # chat-specific formatting template to a conversation input
+        tokenize=False, #  indicates that the function should return the formatted chat as a string (rather than tokenizing it into numerical tokens).
+        add_generation_prompt=True # adds a prompt or marker that signals where the model should start generating the response in a dialogue or conversation flow
+    )
+    return PromptTemplate.from_template(formatted_prompt)
 
 def format_chat_history(messages: List[models.Message]):
     # TODO:  implement format_chat_history to format 
     # the list of Message into a text of chat history.
     raise NotImplemented
 
+
 def format_context(docs: List[str]):
     # TODO:  the output of the DataIndexer.search is a list of text, 
     # so we need to concatenate that list into a text that can fit into 
@@ -47,7 +61,7 @@ standalone_prompt: str = None
 rag_prompt: str = None
 
 # TODO: create raw_prompt_formatted by using format_prompt
-raw_prompt_formatted = None
+raw_prompt_formatted = format_prompt(raw_prompt)
 raw_prompt = PromptTemplate.from_template(raw_prompt)
 
 # TODO: use format_prompt to create history_prompt_formatted