在原retrieve_sse接口上增加问题分类器

app/api/v1/zhipu_controller.py

@@ -1,4 +1,4 @@
-import logging
+import logging,json
 from flask import Blueprint, request, Response
 from app.services.zhipu_service import ZhipuService
 from app.services.zhipu_alltool_service import ZhipuAlltoolService
@@ -48,12 +48,28 @@ def retrive_stream():
     knowledge_id = data.get('knowledge_id', '')
     prompt_template = data.get('prompt_template', '')
 
-    def event_stream():
-        for chunk in zhipu_service.retrive_sse(message, knowledge_id, None):
-            if chunk:
-                yield chunk
+    classification_result_str = zhipu_alltool_service.func_call_classify(message)
+    print(f'classification_result: {classification_result_str}')
+    classification_result = json.loads(classification_result_str)
+
+    if classification_result.get('category') == 'web_search':
+        def event_stream_websearch_sse():
+            for chunk in zhipu_alltool_service.web_search_sse(message):
+                if chunk:
+                    yield chunk
+        return Response(event_stream_websearch_sse(), content_type='text/event-stream')        
+    
+    elif classification_result.get('category')== 'retrive_knowledge':
+        def event_stream_retrive():
+            for chunk in zhipu_service.retrive_sse(message, knowledge_id, None):
+                if chunk:
+                    yield chunk
+        return Response(event_stream_retrive(), content_type='text/event-stream')
+
+    elif classification_result.get('category')== 'generate_report':
+        #do something
+        pass
 
-    return Response(event_stream(), content_type='text/event-stream')
 
 
 @zhipu_controller.route('/zhipu/analysis/stream', methods=['POST'])