github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-07-09 06:24:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Track ChatFireworks time to first_token (#11672)

Browse Source
This commit is contained in:
Erick Friis 2023-10-11 13:37:03 -07:00 committed by GitHub
parent 2c1e735403
commit 28ee6a7c12
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 16 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
libs/langchain/langchain/chat_models/fireworks.py
View File

@ -96,7 +96,10 @@ class ChatFireworks(BaseChatModel):
try:
import fireworks.client
except ImportError as e:
raise ImportError("") from e
raise ImportError(
"Could not import fireworks-ai python package. "
"Please install it with `pip install fireworks-ai`."
) from e
fireworks_api_key = get_from_dict_or_env(
values, "fireworks_api_key", "FIREWORKS_API_KEY"
)
@ -194,6 +197,8 @@ class ChatFireworks(BaseChatModel):
)
default_chunk_class = chunk.__class__
yield ChatGenerationChunk(message=chunk, generation_info=generation_info)
if run_manager:
run_manager.on_llm_new_token(chunk.content, chunk=chunk)
async def _astream(
self,
@ -221,6 +226,8 @@ class ChatFireworks(BaseChatModel):
)
default_chunk_class = chunk.__class__
yield ChatGenerationChunk(message=chunk, generation_info=generation_info)
if run_manager:
await run_manager.on_llm_new_token(token=chunk.content, chunk=chunk)
def completion_with_retry(

9
libs/langchain/langchain/llms/fireworks.py
View File

@ -45,7 +45,10 @@ class Fireworks(LLM):
try:
import fireworks.client
except ImportError as e:
raise ImportError("") from e
raise ImportError(
"Could not import fireworks-ai python package. "
"Please install it with `pip install fireworks-ai`."
) from e
fireworks_api_key = get_from_dict_or_env(
values, "fireworks_api_key", "FIREWORKS_API_KEY"
)
@ -113,6 +116,8 @@ class Fireworks(LLM):
):
chunk = _stream_response_to_generation_chunk(stream_resp)
yield chunk
if run_manager:
run_manager.on_llm_new_token(chunk.text, chunk=chunk)
async def _astream(
self,
@ -132,6 +137,8 @@ class Fireworks(LLM):
):
chunk = _stream_response_to_generation_chunk(stream_resp)
yield chunk
if run_manager:
await run_manager.on_llm_new_token(chunk.text, chunk=chunk)
def stream(
self,