task: add short story logic

task/hot_topic/real_estate_story.py

@@ -1,4 +1,4 @@
-from datetime import datetime
+from datetime import datetime, timedelta
 import json
 from task.manager_task import execute_task
 from config.database import SessionLocal
@@ -9,10 +9,10 @@ from llm import LLMThinkingEngine
 
 def story_edit_task():
     with SessionLocal() as db:
-        # 获取今天的所有帖子信息
+        # 获取昨天至今的最新10条帖子信息
         today_contents = db.query(SourceContent).filter(
-            SourceContent.create_time >= datetime.today().replace(hour=0, minute=0, second=0, microsecond=0)
-        ).limit(10).all()
+            SourceContent.create_time >= (datetime.today().replace(hour=0, minute=0, second=0, microsecond=0) - timedelta(days=1))
+        ).order_by(SourceContent.create_time.desc()).limit(10).all()
         if len(today_contents) == 0:
             logger.info("story_edit_task finish, content size 0")
             return
@@ -49,6 +49,38 @@ def story_edit_task():
             db.add(article)
             db.commit()
             # break  # 目前先处理一条内容，后续再改成批量处理
+        
+        llm_engine = LLMThinkingEngine(system_prompt_file="real_estate_story_short_system_prompt.txt")
+        for content in today_contents:
+            logger.info(f"story_edit_task content id: {content.id}, title: {content.link}, platform: {content.platform}")
+            story = llm_engine.think(f"故事素材：{content.content}")
+            logger.info(f"story_edit_task content id: {content.id} story: {story}")
+            # llm生成的结果有时不是json结构，会在前后增加一些文本，需要提取出json部分进行解析
+            try:
+                json_start = story.find("{")
+                json_end = story.rfind("}") + 1
+                if json_start != -1 and json_end != -1:
+                    story = story[json_start:json_end]
+                else:
+                    logger.warning(f"story_edit_task content id: {content.id} llm生成的结果不是有效的json格式，无法提取故事内容")
+                    continue
+            except json.JSONDecodeError:
+                logger.warning(f"story_edit_task content id: {content.id} llm生成的结果不是有效的json格式，无法解析故事内容")
+                continue
+            # 将生成的故事写入Article表
+            json_story = json.loads(story)
+            title = json_story.get("title", "无标题")
+            paragraphs = json_story.get("body", ["无内容"])
+            article_content = "\n".join(paragraphs)
+            article = Article(
+                title=title,
+                keywords=None,
+                content=article_content,
+                used=False
+            )
+            db.add(article)
+            db.commit()
+            # break  # 目前先处理一条内容，后续再改成批量处理
 
 if __name__ == "__main__":
     execute_task(story_edit_task)