jihuaqiang 4 месяцев назад
Родитель
Сommit
1ef8349851
2 измененных файлов с 11 добавлено и 5 удалено
  1. 10 4
      src/database/models.py
  2. 1 1
      src/tools/scheduler.py

+ 10 - 4
src/database/models.py

@@ -18,7 +18,7 @@ class QueryTaskStatus:
 class KnowledgeSuggestQuery:
     """知识查询建议模型"""
     
-    def __init__(self, task_id: int, question: str, querys: Optional[List[str]] = None, status: int = QueryTaskStatus.PENDING, knowledgeType: str = "", err_msg: str = ""):
+    def __init__(self, task_id: int, question: str, querys: Optional[List[str]] = None, status: int = QueryTaskStatus.PENDING, knowledgeType: str = "", err_msg: str = "", need_store: int = 1):
         """
         初始化查询任务
         
@@ -27,6 +27,9 @@ class KnowledgeSuggestQuery:
             question: 问题
             querys: 查询词列表
             status: 任务状态
+            knowledgeType: 知识类型
+            err_msg: 错误信息
+            need_store: 是否存储查询词
         """
         self.task_id = task_id
         self.question = question
@@ -34,6 +37,7 @@ class KnowledgeSuggestQuery:
         self.status = status
         self.knowledgeType = knowledgeType
         self.err_msg = err_msg or ""
+        self.need_store = need_store
     
     def to_dict(self) -> Dict[str, Any]:
         """转换为字典"""
@@ -43,7 +47,8 @@ class KnowledgeSuggestQuery:
             'querys': json.dumps(self.querys, ensure_ascii=False) if self.querys else None,
             'status': self.status,
             'knowledgeType': self.knowledgeType,
-            'err_msg': self.err_msg or None
+            'err_msg': self.err_msg or None,
+            'need_store': self.need_store
         }
     
     @classmethod
@@ -62,7 +67,8 @@ class KnowledgeSuggestQuery:
             querys=querys,
             status=data['status'],
             knowledgeType=data.get('knowledgeType', ""),
-            err_msg=data.get('err_msg', "")
+            err_msg=data.get('err_msg', ""),
+            need_store=data.get('need_store', 1)
         )
 
 
@@ -86,7 +92,7 @@ class QueryTaskDAO:
         try:
             with self.db_manager.get_cursor() as cursor:
                 sql = """
-                INSERT INTO knowledge_suggest_query (task_id, question, status, knowledgeType, err_msg, needStore)
+                INSERT INTO knowledge_suggest_query (task_id, question, status, knowledgeType, err_msg, need_store)
                 VALUES (%s, %s, %s, %s, %s, %s)
                 ON DUPLICATE KEY UPDATE
                 question = VALUES(question),

+ 1 - 1
src/tools/scheduler.py

@@ -81,7 +81,7 @@ class TaskScheduler:
             
             try:
                 # 使用Agent生成查询词(Agent内部会在内容类型时直接失败并返回空)
-                queries = await self.agent.generate_queries(task.question, task.needStore, task.task_id)
+                queries = await self.agent.generate_queries(task.question, task.need_store, task.task_id)
                 
                 # 若为空,视为不支持内容类型,标记失败
                 if not queries: