Selaa lähdekoodia

Update agent_service: create sql engine for agent db

StrayWarrior 3 viikkoa sitten
vanhempi
commit
4e346dbf27
1 muutettua tiedostoa jossa 6 lisäystä ja 1 poistoa
  1. 6 1
      pqai_agent/agent_service.py

+ 6 - 1
pqai_agent/agent_service.py

@@ -15,6 +15,7 @@ import traceback
 import apscheduler.triggers.cron
 import rocketmq
 from apscheduler.schedulers.background import BackgroundScheduler
+from sqlalchemy.orm import sessionmaker
 
 from pqai_agent import configs
 from pqai_agent.configs import apollo_config
@@ -28,9 +29,11 @@ from pqai_agent.push_service import PushScanThread, PushTaskWorkerPool
 from pqai_agent.rate_limiter import MessageSenderRateLimiter
 from pqai_agent.response_type_detector import ResponseTypeDetector
 from pqai_agent.user_manager import UserManager, UserRelationManager
-from pqai_agent.message_queue_backend import MessageQueueBackend, MemoryQueueBackend, AliyunRocketMQQueueBackend
+from pqai_agent.message_queue_backend import MessageQueueBackend, AliyunRocketMQQueueBackend
 from pqai_agent.user_profile_extractor import UserProfileExtractor
 from pqai_agent.message import MessageType, Message, MessageChannel
+from pqai_agent.utils.db_utils import create_sql_engine
+
 
 class AgentService:
     def __init__(
@@ -56,6 +59,8 @@ class AgentService:
         self.response_type_detector = ResponseTypeDetector()
         self.agent_registry: Dict[str, DialogueManager] = {}
         self.history_dialogue_db = HistoryDialogueDatabase(self.config['storage']['user']['mysql'])
+        self.agent_db_engine = create_sql_engine(self.config['storage']['agent_state']['mysql'])
+        self.AgentDBSession = sessionmaker(bind=self.agent_db_engine)
 
         chat_config = self.config['chat_api']['openai_compatible']
         self.text_model_name = chat_config['text_model']