klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Show details for query tokens (#1392)

Browse Source
This commit is contained in:
Dev Khant
2024-07-05 00:10:56 +05:30
committed by GitHub
parent ea09b5f7f0
commit 4880557d51
25 changed files with 1825 additions and 517 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
embedchain/llm/base.py
View File

@@ -164,7 +164,7 @@ class BaseLlm(JSONSerializable):
return search.run(input_query)
@staticmethod
def _stream_response(answer: Any) -> Generator[Any, Any, None]:
def _stream_response(answer: Any, token_info: Optional[dict[str, Any]] = None) -> Generator[Any, Any, None]:
"""Generator to be used as streaming response
:param answer: Answer chunk from llm
@@ -177,6 +177,8 @@ class BaseLlm(JSONSerializable):
streamed_answer = streamed_answer + chunk
yield chunk
logger.info(f"Answer: {streamed_answer}")
if token_info:
logger.info(f"Token Info: {token_info}")
def query(self, input_query: str, contexts: list[str], config: BaseLlmConfig = None, dry_run=False):
"""
@@ -219,11 +221,18 @@ class BaseLlm(JSONSerializable):
if dry_run:
return prompt
answer = self.get_answer_from_llm(prompt)
if self.config.token_usage:
answer, token_info = self.get_answer_from_llm(prompt)
else:
answer = self.get_answer_from_llm(prompt)
if isinstance(answer, str):
logger.info(f"Answer: {answer}")
if self.config.token_usage:
return answer, token_info
return answer
else:
if self.config.token_usage:
return self._stream_response(answer, token_info)
return self._stream_response(answer)
finally:
if config:
@@ -276,13 +285,13 @@ class BaseLlm(JSONSerializable):
if dry_run:
return prompt
answer = self.get_answer_from_llm(prompt)
answer, token_info = self.get_answer_from_llm(prompt)
if isinstance(answer, str):
logger.info(f"Answer: {answer}")
return answer
return answer, token_info
else:
# this is a streamed response and needs to be handled differently.
return self._stream_response(answer)
return self._stream_response(answer, token_info)
finally:
if config:
# Restore previous config