| import threading |
| import queue |
| import tiktoken |
|
|
| from langchain.chat_models import ChatOpenAI |
| from langchain.callbacks.manager import CallbackManager |
| from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler |
|
|
| from openai.error import InvalidRequestError |
|
|
| class ThreadedGenerator: |
| def __init__(self): |
| self.queue = queue.Queue() |
|
|
| def __iter__(self): |
| return self |
|
|
| def __next__(self): |
| item = self.queue.get() |
| if item is StopIteration: raise item |
| return item |
|
|
| def send(self, data): |
| self.queue.put(data) |
|
|
| def close(self): |
| self.queue.put(StopIteration) |
|
|
| class ChainStreamHandler(StreamingStdOutCallbackHandler): |
| def __init__(self, gen): |
| super().__init__() |
| self.gen = gen |
|
|
| def on_llm_new_token(self, token: str, **kwargs): |
| self.gen.send(token) |
|
|
| def run(g, prompts, retries = 0): |
| try: |
| chat = ChatOpenAI( |
| verbose=True, |
| streaming=True, |
| callback_manager=CallbackManager([ChainStreamHandler(g)]), |
| ) |
| chat(prompts) |
| except InvalidRequestError as e: |
| if retries >= 10: |
| raise e |
| else: |
| retries += 1 |
| del prompts[-1] |
| run(g, prompts, retries) |
|
|
| def llm_thread(g, prompts): |
| encoding = tiktoken.encoding_for_model("gpt-3.5-turbo") |
| messages = [] |
| tokens = 0 |
|
|
| for prompt in prompts: |
| tokens += len(encoding.encode(prompt.content)) |
| if tokens >= 4096: |
| break |
| |
| messages.append(prompt) |
|
|
| try: |
| run(g, messages) |
| finally: |
| g.close() |
|
|
|
|
| def chat(prompts): |
| g = ThreadedGenerator() |
| threading.Thread(target=llm_thread, args=(g, prompts)).start() |
| return g |
