Spaces:
Paused
Paused
Daniel Marques
commited on
Commit
·
5068745
1
Parent(s):
b3fd6f4
feat: add backend
Browse files
main.py
CHANGED
@@ -58,8 +58,9 @@ RETRIEVER = DB.as_retriever()
|
|
58 |
class MyCustomSyncHandler(BaseCallbackHandler):
|
59 |
def __init__(self):
|
60 |
self.end = False
|
|
|
61 |
|
62 |
-
|
63 |
self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
|
64 |
) -> None:
|
65 |
self.end = False
|
@@ -68,7 +69,7 @@ class MyCustomSyncHandler(BaseCallbackHandler):
|
|
68 |
self.end = True
|
69 |
|
70 |
def on_llm_new_token(self, token: str, **kwargs) -> Any:
|
71 |
-
|
72 |
|
73 |
|
74 |
handlerToken = MyCustomSyncHandler()
|
@@ -257,7 +258,7 @@ async def websocket_endpoint(websocket: WebSocket):
|
|
257 |
finish = False
|
258 |
while finish == False:
|
259 |
finish = handlerToken.end
|
260 |
-
token = handlerToken.
|
261 |
await websocket.send_text(f"result: {token}")
|
262 |
|
263 |
|
|
|
58 |
class MyCustomSyncHandler(BaseCallbackHandler):
|
59 |
def __init__(self):
|
60 |
self.end = False
|
61 |
+
self.token = ""
|
62 |
|
63 |
+
def on_llm_start(
|
64 |
self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
|
65 |
) -> None:
|
66 |
self.end = False
|
|
|
69 |
self.end = True
|
70 |
|
71 |
def on_llm_new_token(self, token: str, **kwargs) -> Any:
|
72 |
+
self.token += token
|
73 |
|
74 |
|
75 |
handlerToken = MyCustomSyncHandler()
|
|
|
258 |
finish = False
|
259 |
while finish == False:
|
260 |
finish = handlerToken.end
|
261 |
+
token = handlerToken.token
|
262 |
await websocket.send_text(f"result: {token}")
|
263 |
|
264 |
|