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

Improve and add more tests (#807)

Browse Source
This commit is contained in:
Sidharth Mohanty
2023-10-18 02:36:47 +05:30
committed by GitHub
parent d065cbf934
commit e8a2846449
11 changed files with 403 additions and 260 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
embedchain/llm/openai.py
View File

@@ -13,13 +13,9 @@ class OpenAILlm(BaseLlm):
def __init__(self, config: Optional[BaseLlmConfig] = None):
super().__init__(config=config)
def get_llm_model_answer(self, prompt):
def get_llm_model_answer(self, prompt) -> str:
response = OpenAILlm._get_answer(prompt, self.config)
if self.config.stream:
return response
else:
return response.content
return response
def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
messages = []
@@ -35,10 +31,9 @@ class OpenAILlm(BaseLlm):
if config.top_p:
kwargs["model_kwargs"]["top_p"] = config.top_p
if config.stream:
from langchain.callbacks.streaming_stdout import \
StreamingStdOutCallbackHandler
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
chat = ChatOpenAI(**kwargs, streaming=config.stream, callbacks=[StreamingStdOutCallbackHandler()])
else:
chat = ChatOpenAI(**kwargs)
return chat(messages)
return chat(messages).content