|
2 | 2 | IChatResponse,
|
3 | 3 | IUserMessage,
|
4 | 4 | )
|
| 5 | +import logging |
5 | 6 | from app.utils.adaptive_cards.cards import create_adaptive_card
|
6 | 7 | from app.utils.callback import (
|
7 | 8 | CustomAsyncCallbackHandler,
|
@@ -84,55 +85,60 @@ async def websocket_endpoint(websocket: WebSocket):
|
84 | 85 | @router.websocket("/tools")
|
85 | 86 | async def websocket_endpoint(websocket: WebSocket):
|
86 | 87 | await websocket.accept()
|
87 |
| - while True: |
88 |
| - data = await websocket.receive_json() |
89 |
| - user_message = data["message"] |
90 |
| - user_message_card = create_adaptive_card(user_message) |
91 | 88 |
|
92 |
| - resp = IChatResponse( |
93 |
| - sender="you", |
94 |
| - message=user_message_card.to_dict(), |
95 |
| - type="start", |
96 |
| - message_id=str(uuid7()), |
97 |
| - id=str(uuid7()), |
98 |
| - ) |
99 |
| - |
100 |
| - await websocket.send_json(resp.dict()) |
101 |
| - message_id: str = str(uuid7()) |
102 |
| - custom_handler = CustomFinalStreamingStdOutCallbackHandler( |
103 |
| - websocket, message_id=message_id |
104 |
| - ) |
105 |
| - |
106 |
| - tools = [ |
107 |
| - GeneralKnowledgeTool(), |
108 |
| - PokemonSearchTool(), |
109 |
| - ImageSearchTool(), |
110 |
| - YoutubeSearchTool(), |
111 |
| - GeneralWeatherTool(), |
112 |
| - ] |
113 |
| - |
114 |
| - llm = ChatOpenAI( |
115 |
| - streaming=True, |
116 |
| - temperature=0, |
117 |
| - ) |
118 |
| - |
119 |
| - agent = ZeroShotAgent.from_llm_and_tools( |
120 |
| - llm=llm, |
121 |
| - tools=tools, |
122 |
| - prefix=zero_agent_prompt.prefix, |
123 |
| - suffix=zero_agent_prompt.suffix, |
124 |
| - format_instructions=zero_agent_prompt.format_instructions, |
125 |
| - input_variables=zero_agent_prompt.input_variables, |
126 |
| - ) |
127 |
| - # TODO: We should use this |
128 |
| - # * max_execution_time=1, |
129 |
| - # early_stopping_method="generate", |
130 |
| - agent_executor = AgentExecutor.from_agent_and_tools( |
131 |
| - agent=agent, |
132 |
| - tools=tools, |
133 |
| - verbose=False, |
134 |
| - handle_parsing_errors=True, |
135 |
| - memory=memory, |
136 |
| - ) |
| 89 | + while True: |
| 90 | + try: |
| 91 | + data = await websocket.receive_json() |
| 92 | + user_message = data["message"] |
| 93 | + user_message_card = create_adaptive_card(user_message) |
| 94 | + |
| 95 | + resp = IChatResponse( |
| 96 | + sender="you", |
| 97 | + message=user_message_card.to_dict(), |
| 98 | + type="start", |
| 99 | + message_id=str(uuid7()), |
| 100 | + id=str(uuid7()), |
| 101 | + ) |
| 102 | + |
| 103 | + await websocket.send_json(resp.dict()) |
| 104 | + message_id: str = str(uuid7()) |
| 105 | + custom_handler = CustomFinalStreamingStdOutCallbackHandler( |
| 106 | + websocket, message_id=message_id |
| 107 | + ) |
| 108 | + |
| 109 | + tools = [ |
| 110 | + GeneralKnowledgeTool(), |
| 111 | + PokemonSearchTool(), |
| 112 | + ImageSearchTool(), |
| 113 | + YoutubeSearchTool(), |
| 114 | + GeneralWeatherTool(), |
| 115 | + ] |
137 | 116 |
|
138 |
| - await agent_executor.arun(input=user_message, callbacks=[custom_handler]) |
| 117 | + llm = ChatOpenAI( |
| 118 | + streaming=True, |
| 119 | + temperature=0, |
| 120 | + ) |
| 121 | + |
| 122 | + agent = ZeroShotAgent.from_llm_and_tools( |
| 123 | + llm=llm, |
| 124 | + tools=tools, |
| 125 | + prefix=zero_agent_prompt.prefix, |
| 126 | + suffix=zero_agent_prompt.suffix, |
| 127 | + format_instructions=zero_agent_prompt.format_instructions, |
| 128 | + input_variables=zero_agent_prompt.input_variables, |
| 129 | + ) |
| 130 | + # TODO: We should use this |
| 131 | + # * max_execution_time=1, |
| 132 | + # early_stopping_method="generate", |
| 133 | + agent_executor = AgentExecutor.from_agent_and_tools( |
| 134 | + agent=agent, |
| 135 | + tools=tools, |
| 136 | + verbose=False, |
| 137 | + handle_parsing_errors=True, |
| 138 | + memory=memory, |
| 139 | + ) |
| 140 | + |
| 141 | + await agent_executor.arun(input=user_message, callbacks=[custom_handler]) |
| 142 | + except WebSocketDisconnect: |
| 143 | + logging.info("websocket disconnect") |
| 144 | + break |
0 commit comments