update to use async

This commit is contained in:
James Briggs
2023-09-28 11:48:59 +02:00
parent 56f08768e8
commit 11f60624ef
2 changed files with 122 additions and 105 deletions

View File

@@ -72,7 +72,7 @@
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": 20,
"metadata": {},
"outputs": [],
"source": [
@@ -101,7 +101,7 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": 21,
"metadata": {},
"outputs": [
{
@@ -120,26 +120,28 @@
"\n",
"Days turned into weeks, and weeks turned into months, but Lily never lost faith. She faced countless challenges along the way, from treacherous paths to dark and eerie caves. Yet, she pressed on, fueled by her unwavering determination to find the Moonlight Unicorn.\n",
"\n",
"One fateful night, as the full moon illuminated the forest, Lily stumbled upon a clearing bathed in a soft, silvery light. There, standing before her, was the Moonlight Unicorn. Its coat shimmered like stardust, and its horn emitted a gentle glow that illuminated the entire clearing.\n",
"One fateful night, as the full moon illuminated the forest, Lily stumbled upon a clearing bathed in a soft, silver light. There, standing before her, was the Moonlight Unicorn. Its coat shimmered like stardust, and its horn emitted a gentle glow that illuminated the entire clearing.\n",
"\n",
"Overwhelmed with joy, Lily approached the unicorn cautiously, her heart pounding with excitement. The unicorn, sensing her pure intentions, greeted her with a gentle nudge. Lily reached out and touched the unicorn's horn, feeling a surge of warmth and magic course through her veins.\n",
"\n",
"In that moment, the Moonlight Unicorn spoke to Lily, its voice as soft as a whisper. It thanked her for her unwavering determination and rewarded her with a single wish. Lily, without hesitation, wished for the forest to be protected and preserved for generations to come, ensuring its magic would never fade away.\n",
"\n",
"As her wish was granted, the Moonlight Unicorn vanished into the night, leaving Lily standing in the clearing, forever changed by her encounter. She returned to her village, sharing her incredible journey and the importance of preserving the forest's magic.\n",
"As her wish was granted, the forest came alive with renewed energy. The trees stood taller, the animals thrived, and the enchantment of the forest grew stronger than ever before. Lily knew that her journey had not only brought her closer to the Moonlight Unicorn but had also helped preserve the magic of the forest for all to enjoy.\n",
"\n",
"From that day forward, Lily became the guardian of the enchanted forest, dedicating her life to its preservation. She taught others about the wonders of the forest, spreading the message of harmony between humans and nature.\n",
"With a heart full of gratitude, Lily bid farewell to the Moonlight Unicorn, knowing that she had fulfilled her purpose. She returned to her village, where she shared her incredible adventure with her grandmother and the villagers, inspiring them to cherish and protect the forest.\n",
"\n",
"And so, the legend of Lily and the Moonlight Unicorn lived on, inspiring generations to cherish the magic of the natural world and reminding them that with determination and a pure heart, dreams can come true."
"From that day forward, Lily became known as the Guardian of the Forest, and her story was passed down through generations. The enchanted forest thrived, and the Moonlight Unicorn continued to bless those who sought its presence.\n",
"\n",
"And so, the tale of Lily and the Moonlight Unicorn became a legend, reminding people of the power of determination, the magic of nature, and the importance of preserving the wonders that surround us."
]
},
{
"data": {
"text/plain": [
"AIMessageChunk(content=\"Once upon a time, in a small village nestled deep within a lush forest, there lived a young girl named Lily. She was known for her vibrant red hair, sparkling green eyes, and a heart full of curiosity. Lily had always been fascinated by the stories her grandmother would tell her about a magical creature called the Moonlight Unicorn.\\n\\nLegend had it that the Moonlight Unicorn was a rare and majestic creature that only appeared during the full moon. Its coat shimmered like silver, and its horn glowed with a soft, ethereal light. The unicorn was said to possess incredible powers, capable of granting wishes and bringing good fortune to those who encountered it.\\n\\nDriven by her desire to see the Moonlight Unicorn, Lily embarked on a journey through the enchanted forest. Armed with her grandmother's stories and a sense of adventure, she ventured deeper into the woods, following the faint whispers of the wind.\\n\\nAs she wandered through the forest, Lily encountered various magical creatures. She met a mischievous sprite who led her to a hidden waterfall, where she discovered a family of playful water nymphs. They taught her the secrets of the forest and shared their wisdom about the Moonlight Unicorn.\\n\\nWith newfound knowledge, Lily continued her quest, guided by the moon's gentle glow. She encountered talking animals, wise old trees, and even a friendly dragon who offered her protection on her journey. Each encounter brought her closer to her goal, and her heart filled with hope.\\n\\nDays turned into weeks, and weeks turned into months, but Lily never lost faith. She faced countless challenges along the way, from treacherous paths to dark and eerie caves. Yet, she pressed on, fueled by her unwavering determination to find the Moonlight Unicorn.\\n\\nOne fateful night, as the full moon illuminated the forest, Lily stumbled upon a clearing bathed in a soft, silvery light. There, standing before her, was the Moonlight Unicorn. Its coat shimmered like stardust, and its horn emitted a gentle glow that illuminated the entire clearing.\\n\\nOverwhelmed with joy, Lily approached the unicorn cautiously, her heart pounding with excitement. The unicorn, sensing her pure intentions, greeted her with a gentle nudge. Lily reached out and touched the unicorn's horn, feeling a surge of warmth and magic course through her veins.\\n\\nIn that moment, the Moonlight Unicorn spoke to Lily, its voice as soft as a whisper. It thanked her for her unwavering determination and rewarded her with a single wish. Lily, without hesitation, wished for the forest to be protected and preserved for generations to come, ensuring its magic would never fade away.\\n\\nAs her wish was granted, the Moonlight Unicorn vanished into the night, leaving Lily standing in the clearing, forever changed by her encounter. She returned to her village, sharing her incredible journey and the importance of preserving the forest's magic.\\n\\nFrom that day forward, Lily became the guardian of the enchanted forest, dedicating her life to its preservation. She taught others about the wonders of the forest, spreading the message of harmony between humans and nature.\\n\\nAnd so, the legend of Lily and the Moonlight Unicorn lived on, inspiring generations to cherish the magic of the natural world and reminding them that with determination and a pure heart, dreams can come true.\", additional_kwargs={}, example=False)"
"AIMessageChunk(content=\"Once upon a time, in a small village nestled deep within a lush forest, there lived a young girl named Lily. She was known for her vibrant red hair, sparkling green eyes, and a heart full of curiosity. Lily had always been fascinated by the stories her grandmother would tell her about a magical creature called the Moonlight Unicorn.\\n\\nLegend had it that the Moonlight Unicorn was a rare and majestic creature that only appeared during the full moon. Its coat shimmered like silver, and its horn glowed with a soft, ethereal light. The unicorn was said to possess incredible powers, capable of granting wishes and bringing good fortune to those who encountered it.\\n\\nDriven by her desire to see the Moonlight Unicorn, Lily embarked on a journey through the enchanted forest. Armed with her grandmother's stories and a sense of adventure, she ventured deeper into the woods, following the faint whispers of the wind.\\n\\nAs she wandered through the forest, Lily encountered various magical creatures. She met a mischievous sprite who led her to a hidden waterfall, where she discovered a family of playful water nymphs. They taught her the secrets of the forest and shared their wisdom about the Moonlight Unicorn.\\n\\nWith newfound knowledge, Lily continued her quest, guided by the moon's gentle glow. She encountered talking animals, wise old trees, and even a friendly dragon who offered her protection on her journey. Each encounter brought her closer to her goal, and her heart filled with hope.\\n\\nDays turned into weeks, and weeks turned into months, but Lily never lost faith. She faced countless challenges along the way, from treacherous paths to dark and eerie caves. Yet, she pressed on, fueled by her unwavering determination to find the Moonlight Unicorn.\\n\\nOne fateful night, as the full moon illuminated the forest, Lily stumbled upon a clearing bathed in a soft, silver light. There, standing before her, was the Moonlight Unicorn. Its coat shimmered like stardust, and its horn emitted a gentle glow that illuminated the entire clearing.\\n\\nOverwhelmed with joy, Lily approached the unicorn cautiously, her heart pounding with excitement. The unicorn, sensing her pure intentions, greeted her with a gentle nudge. Lily reached out and touched the unicorn's horn, feeling a surge of warmth and magic course through her veins.\\n\\nIn that moment, the Moonlight Unicorn spoke to Lily, its voice as soft as a whisper. It thanked her for her unwavering determination and rewarded her with a single wish. Lily, without hesitation, wished for the forest to be protected and preserved for generations to come, ensuring its magic would never fade away.\\n\\nAs her wish was granted, the forest came alive with renewed energy. The trees stood taller, the animals thrived, and the enchantment of the forest grew stronger than ever before. Lily knew that her journey had not only brought her closer to the Moonlight Unicorn but had also helped preserve the magic of the forest for all to enjoy.\\n\\nWith a heart full of gratitude, Lily bid farewell to the Moonlight Unicorn, knowing that she had fulfilled her purpose. She returned to her village, where she shared her incredible adventure with her grandmother and the villagers, inspiring them to cherish and protect the forest.\\n\\nFrom that day forward, Lily became known as the Guardian of the Forest, and her story was passed down through generations. The enchanted forest thrived, and the Moonlight Unicorn continued to bless those who sought its presence.\\n\\nAnd so, the tale of Lily and the Moonlight Unicorn became a legend, reminding people of the power of determination, the magic of nature, and the importance of preserving the wonders that surround us.\", additional_kwargs={}, example=False)"
]
},
"execution_count": 5,
"execution_count": 21,
"metadata": {},
"output_type": "execute_result"
}
@@ -163,7 +165,7 @@
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": 22,
"metadata": {},
"outputs": [],
"source": [
@@ -314,12 +316,31 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"It's interesting to see during development but we'll want to clean this streaming up a little in any actual use-case. For that we can go with two approaches — either we build a custom callback handler, or use a purpose built callback handler from LangChain (as usual, LangChain has something for everything). Let's first try LangChain's purpose-built `FinalStreamingStdOutCallbackHandler`."
"It's interesting to see during development but we'll want to clean this streaming up a little in any actual use-case. For that we can go with two approaches — either we build a custom callback handler, or use a purpose built callback handler from LangChain (as usual, LangChain has something for everything). Let's first try LangChain's purpose-built `FinalStreamingStdOutCallbackHandler`.\n",
"\n",
"We will overwrite the existing `callbacks` attribute found here:"
]
},
{
"cell_type": "code",
"execution_count": 13,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"agent.agent.llm_chain.llm"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"With the new callback handler:"
]
},
{
"cell_type": "code",
"execution_count": 28,
"metadata": {},
"outputs": [],
"source": [
@@ -327,29 +348,11 @@
" FinalStreamingStdOutCallbackHandler,\n",
")\n",
"\n",
"llm = ChatOpenAI(\n",
" openai_api_key=os.getenv(\"OPENAI_API_KEY\"),\n",
" temperature=0.0,\n",
" model_name=\"gpt-3.5-turbo\",\n",
" streaming=True,\n",
" callbacks=[\n",
" FinalStreamingStdOutCallbackHandler(\n",
" answer_prefix_tokens=[\"Final\", \"Answer\"]\n",
" )\n",
" ] # ! important\n",
")\n",
"\n",
"# reinitialize the agent\n",
"agent = initialize_agent(\n",
" agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,\n",
" tools=tools,\n",
" llm=llm,\n",
" memory=memory,\n",
" verbose=True,\n",
" max_iterations=3,\n",
" early_stopping_method=\"generate\",\n",
" return_intermediate_steps=False\n",
")"
"agent.agent.llm_chain.llm.callbacks = [\n",
" FinalStreamingStdOutCallbackHandler(\n",
" answer_prefix_tokens=[\"Final\", \"Answer\"]\n",
" )\n",
"]"
]
},
{
@@ -450,25 +453,7 @@
" sys.stdout.write(token) # equal to `print(token, end=\"\")`\n",
" sys.stdout.flush()\n",
"\n",
"llm = ChatOpenAI(\n",
" openai_api_key=os.getenv(\"OPENAI_API_KEY\"),\n",
" temperature=0.0,\n",
" model_name=\"gpt-3.5-turbo\",\n",
" streaming=True,\n",
" callbacks=[CallbackHandler()] # ! important\n",
")\n",
"\n",
"# reinitialize the agent\n",
"agent = initialize_agent(\n",
" agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,\n",
" tools=tools,\n",
" llm=llm,\n",
" memory=memory,\n",
" verbose=True,\n",
" max_iterations=3,\n",
" early_stopping_method=\"generate\",\n",
" return_intermediate_steps=False\n",
")"
"agent.agent.llm_chain.llm.callbacks = [CallbackHandler()]"
]
},
{
@@ -535,6 +520,26 @@
"agent(\"what is the square root of 71?\")"
]
},
{
"cell_type": "code",
"execution_count": 27,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[<langchain.callbacks.streaming_stdout.StreamingStdOutCallbackHandler at 0x11fd01640>]"
]
},
"execution_count": 27,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"agent.agent.llm_chain.llm"
]
},
{
"attachments": {},
"cell_type": "markdown",
@@ -565,7 +570,7 @@
},
{
"cell_type": "code",
"execution_count": 22,
"execution_count": 1,
"metadata": {},
"outputs": [
{
@@ -574,7 +579,7 @@
"{'status': '🤙'}"
]
},
"execution_count": 22,
"execution_count": 1,
"metadata": {},
"output_type": "execute_result"
}
@@ -591,16 +596,16 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Unlike with our StdOut streaming, to stream via FastAPI we need to run two parallel chunks of code. One is what we're already running, the agent. The other is the code that feeds those tokens to FastAPI via a `StreamingResponse` object. To handle this we need to use multiple threads to handle the two running chunks of logic, and we need to create a way of communicating between these two threads — to do that we use a `Queue`.\n",
"Unlike with our StdOut streaming, we now need to send our tokens to a generator function that feeds those tokens to FastAPI via a `StreamingResponse` object. To handle this we need to use async code, otherwise our generator will not begin emitting anything until _after_ generation is already complete.\n",
"\n",
"The `Queue` is accessed by our callback handler, as as each token is generated, it puts the token into the queue. In parallel, the main event loop (our main thread) is constantly checking for new tokens being added to the queue. As soon as the main event loop sees a token as been added, it gets the token and yields it to our `StreamingResponse`.\n",
"The `Queue` is accessed by our callback handler, as as each token is generated, it puts the token into the queue. Our generator function asyncronously checks for new tokens being added to the queue. As soon as the generator sees a token has been added, it gets the token and yields it to our `StreamingResponse`.\n",
"\n",
"To see it in action, we'll define a stream requests function called `get_stream`:"
]
},
{
"cell_type": "code",
"execution_count": 23,
"execution_count": 14,
"metadata": {},
"outputs": [],
"source": [
@@ -617,10 +622,20 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 16,
"metadata": {},
"outputs": [],
"source": []
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
" \"Once upon a time, in a faraway land, there was a young prince named Edward. He lived in a magnificent castle surrounded by lush gardens and sparkling fountains. Edward was known for his kindness and generosity, and he was loved by all who knew him. One day, a mysterious old woman appeared at the castle gates. She was hunched over and carried a staff, and her eyes sparkled with wisdom. The old woman told Edward that she was a powerful sorceress, and she had come to test his character. She presented him with a challenge: if he could prove himself worthy, she would grant him a single wish. Edward eagerly accepted the challenge, and the old woman led him on a journey through enchanted forests, treacherous mountains, and dark caves. Along the way, Edward encountered many obstacles and faced numerous trials. He showed bravery in the face of danger, compassion towards those in need, and wisdom in his decision-making. Finally, after what felt like an eternity, Edward reached the end of his journey. The old woman stood before him, her eyes filled with pride. She granted him his wish, and Edward chose to use it to bring prosperity and happiness to his kingdom. From that day forward, Edward ruled with fairness and justice, and his kingdom flourished. The story of Prince Edward's journey and his ultimate triumph spread far and wide, inspiring others to be brave, kind, and wise. And so, the legend of Prince Edward lived on, reminding us all of the power of character and the importance of staying true to ourselves.\"\n"
]
}
],
"source": [
"get_stream(\"tell me a long story\")"
]
}
],
"metadata": {

View File

@@ -1,6 +1,6 @@
import os
from threading import Thread
from typing import Generator
import asyncio
from typing import Any
import uvicorn
from fastapi import FastAPI, Body
@@ -11,39 +11,19 @@ from pydantic import BaseModel
from langchain.agents import AgentType, initialize_agent
from langchain.chat_models import ChatOpenAI
from langchain.memory import ConversationBufferWindowMemory
from langchain.callbacks.base import BaseCallbackHandler
from langchain.callbacks.streaming_aiter import AsyncIteratorCallbackHandler
from langchain.callbacks.streaming_stdout_final_only import FinalStreamingStdOutCallbackHandler
from langchain.schema import LLMResult
app = FastAPI()
class QueueCallbackHandler(BaseCallbackHandler):
def __init__(self, queue: Queue):
self.queue = queue
self.content: str = ""
self.final_answer: bool = False
def on_llm_new_token(self, token: str, **kwargs: any) -> None:
self.content += token
# if we passed the final answer, we put tokens in queue
if self.final_answer:
if '"action_input": "' in self.content:
if token not in ['"', "}"]:
self.queue.put(token)
elif "Final Answer" in self.content:
self.final_answer = True
self.content = ""
# queue is our communication between threads
queue = Queue()
# we use this for signaling the end of the queue
job_done = "<END-STREAM>"
# initialize the agent
# initialize the agent (we need to do this for the callbacks)
llm = ChatOpenAI(
openai_api_key=os.getenv("OPENAI_API_KEY"),
temperature=0.0,
model_name="gpt-3.5-turbo",
streaming=True, # ! important
callbacks=[QueueCallbackHandler(queue)] # ! important
callbacks=[] # ! important (but we will add them later)
)
memory = ConversationBufferWindowMemory(
memory_key="chat_history",
@@ -62,30 +42,52 @@ agent = initialize_agent(
return_intermediate_steps=False
)
# this handles the call to the agent, followed by job_done token
def call(query: str, queue: Queue) -> None:
agent.run(query)
queue.put(job_done)
class AsyncCallbackHandler(AsyncIteratorCallbackHandler):
content: str = ""
final_answer: bool = False
async def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
self.content += token
# if we passed the final answer, we put tokens in queue
if self.final_answer:
if '"action_input": "' in self.content:
if token not in ['"', "}"]:
self.queue.put_nowait(token)
elif "Final Answer" in self.content:
self.final_answer = True
self.content = ""
async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
if self.final_answer:
self.content = ""
self.final_answer = False
self.done.set()
else:
self.content = ""
async def run_call(query: str, stream_it: AsyncCallbackHandler):
# assign callback handler
agent.agent.llm_chain.llm.callbacks = [stream_it]
# now query
await agent.acall(inputs={"input": query})
# request input format
class Query(BaseModel):
text: str
def stream(query: Query) -> Generator:
"""This function is what allows us to stream"""
# move the agent task to another thread
thread = Thread(target=call, args=(query.text, queue))
thread.start()
# on main thread, yield the results from queue to StreamingResponse
while True:
result = queue.get()
if result is job_done:
break
yield result
async def create_gen(query: str, stream_it: AsyncCallbackHandler):
task = asyncio.create_task(run_call(query, stream_it))
async for token in stream_it.aiter():
yield token
await task
@app.post("/chat")
async def chat(query: Query = Body(...)):
return StreamingResponse(stream(query), media_type="text/event-stream")
async def chat(
query: Query = Body(...),
):
stream_it = AsyncCallbackHandler()
gen = create_gen(query.text, stream_it)
return StreamingResponse(gen, media_type="text/event-stream")
@app.get("/health")
async def health():