feat(core): add di and log extra print

Author: gusye1234

src/server/core/acontext_core/__init__.py

@@ -1,4 +1,4 @@
-from .entry import setup, cleanup
+from .di import setup, cleanup
 
 __version__ = "0.0.1.dev"
 __url__ = "https://github.com/memodb-io/Acontext"

src/server/core/acontext_core/entry.py src/server/core/acontext_core/di.pysrc/server/core/acontext_core/entry.py renamed to src/server/core/acontext_core/di.py

@@ -205,7 +205,7 @@ async def _process_message(self, config: ConsumerConfig, message: Message) -> No
                             f"Message processing unknown error - queue: {config.queue_name}, "
                             f"attempt: {retry_count}/{config.max_retries}, "
                             f"retry after {_wait_for}s, "
-                            f"error: {str(e)}. {traceback.format_exc()}",
+                            f"error: {str(e)}.",
                             extra={"traceback": traceback.format_exc()},
                         )
                         await asyncio.sleep(_wait_for)  # Exponential backoff

src/server/core/acontext_core/infra/async_mq.py

@@ -139,8 +139,10 @@ async def get_session_context(self) -> AsyncGenerator[AsyncSession, None]:
             yield session
             await session.commit()
         except Exception as e:
-            logger.error(f"DB Session failed: {str(e)}. Rollback...")
-            logger.error(traceback.format_exc(), extra={"__internal__": True})
+            logger.error(
+                f"DB Session failed: {str(e)}. Rollback...",
+                extra={"traceback": traceback.format_exc()},
+            )
             await session.rollback()
             raise e
         finally:

src/server/core/acontext_core/infra/db.py

@@ -1,3 +1,4 @@
+import re
 from typing import List
 from urllib import response
 from ...env import LOG, CONFIG, bound_logging_vars
@@ -34,7 +35,7 @@ async def task_agent_curd(
     session_id: asUUID,
     previous_messages: List[MessageBlob],
     messages: List[MessageBlob],
-    max_iterations=3,
+    max_iterations=1,  # task curd agent only receive one turn of actions
 ) -> Result[None]:
     async with DB_CLIENT.get_session_context() as db_session:
         r = await TD.fetch_current_tasks(db_session, session_id)
@@ -87,7 +88,7 @@ async def task_agent_curd(
             use_ctx = TaskCtx(
                 db_session=db_session,
                 session_id=session_id,
-                task_ids_index=[t.id for t in tasks],
+                task_ids_index=[t.id for t in current_tasks],
                 message_ids_index=[m.message_id for m in messages],
             )
             tool_response = []

src/server/core/acontext_core/llm/agent/task.py

@@ -28,20 +28,20 @@ def system_prompt(cls) -> str:
 
 ## Analysis Guidelines
 ### Planning Detection
-- Look for explicit task planning language ("I need to...", "My goal is...", "I will follow ... steps")
-- Read out the planning, and separate the tasks from it.
-- Link those planning messages to the planning section, since they aren't related to any specific task execution.
-- Collect all current tasks without missing future ones
-- 
+- Look for explicit task planning language ("My plan is to...")
+- Look for user requirements and preferences.
+- General plannings from user/agent.
+- The messages that cause you to create/update tasks.
 
 ### New Task Detection
 - Avoid creating tasks for simple questions answerable directly
 - Only collect tasks stated by agents/users, don't invent them
 - User's requirement should be confimed by the agent's response, then it becomes a valid task, and append those requirements to planning section.
-- [think] The degree of task splitting should follow the agent's plan in the conversation; do not arbitrarily split into finer or coarser granularity.
-- [think] Notice any task modification from agent.
-- [think] Infer execution order and insert tasks sequentially, make sure you arrange the tasks in logical execution order, no the mentioned order.
-- [think] Ensure no task overlap, make sure the tasks are MECE(mutually exclusive, collectively exhaustive).
+- The degree of task splitting should follow the agent's plan in the conversation; do not arbitrarily split into finer or coarser granularity.
+- Notice any task modification from agent.
+- Infer execution order and insert tasks sequentially, make sure you arrange the tasks in logical execution order, no the mentioned order.
+- Ensure no task overlap, make sure the tasks are MECE(mutually exclusive, collectively exhaustive).
+- When valid new tasks mentioned, always try to capture them all, not only the first one.
 
 ### Task Assignment  
 - Match agent responses/actions to existing task descriptions and contexts
@@ -66,15 +66,16 @@ def system_prompt(cls) -> str:
   - `## Current Message with IDs`: the current messages that you need to analyze [with message ids]
 - Message with ID format: <message id=N> ... </message>, inside the tag is the message content, the id field indicates the message id.
 
-## Think before calling tools
+## Report your thinking before calling tools
 - Use extremely brief sentences to state the plans & tasks conversation mentioned, if any.
 - Use one-two sentences to briefly describe your plan.
+- At the end, confirm you can call finish tool and call it at the end of your actions.
 
 ## Action Guidelines
 - Be precise, context-aware, and conservative. 
 - Focus on meaningful task management that organizes conversation objectives effectively. 
-- Use parallel tool calls when possible, and make sure you call the tools in the correct order.
-- After completing all task management actions, call the `finish` tool.
+- Use parallel tool calls, and make sure you call the tools in the correct order.
+- Make sure you called every tool that you need to call based on your report.
 """
 
     @classmethod

src/server/core/acontext_core/llm/prompt/task.py

@@ -6,6 +6,10 @@
 T = TypeVar("T")
 
 
+class ResultError(Exception):
+    pass
+
+
 class Error(BaseModel):
     status: Code = Code.SUCCESS
     errmsg: str = ""
@@ -41,3 +45,7 @@ def ok(self) -> bool:
         if self.error.status != Code.SUCCESS:
             return False
         return True
+
+    def raise_error(self):
+        if self.error.status != Code.SUCCESS:
+            raise ResultError(str(self.error))

src/server/core/acontext_core/schema/result.py

@@ -3,6 +3,7 @@
 from ...schema.session.task import TaskStatus
 from ...schema.session.message import MessageBlob
 from ...schema.utils import asUUID
+from ...schema.result import ResultError
 from ...env import LOG, CONFIG
 from ...llm.agent import task as AT
 
@@ -41,15 +42,19 @@ async def process_session_pending_message(session_id: asUUID):
             ]
 
         r = await AT.task_agent_curd(session_id, previous_messages_data, messages_data)
+
+        after_status = TaskStatus.SUCCESS
+        if not r.ok():
+            after_status = TaskStatus.FAILED
         async with DB_CLIENT.get_session_context() as session:
             await MD.update_message_status_to(
-                session, pending_message_ids, TaskStatus.SUCCESS
+                session, pending_message_ids, after_status
             )
     except Exception as e:
         if pending_message_ids is None:
             raise e
         LOG.error(
-            f"Exception while processing session pending message: {e}, rollback {len(pending_message_ids)} message status to pending"
+            f"Exception while processing session pending message: {e}, rollback {len(pending_message_ids)} message status to failed"
         )
         async with DB_CLIENT.get_session_context() as session:
             await MD.update_message_status_to(

src/server/core/acontext_core/service/controller/message.py

@@ -55,11 +55,34 @@ def __get_json_logger():
     return logger
 
 
+class CustomTextFormatter(logging.Formatter):
+    """Custom formatter that supports extra fields for text logging"""
+
+    def format(self, record):
+        # First, format the standard message
+        base_msg = super().format(record)
+
+        # Check if there are extra fields (beyond the standard ones)
+        printout_attrs = {"traceback"}
+
+        all_fields = record.__dict__
+        extra_fields = []
+        for attr in printout_attrs:
+            if attr in all_fields:
+                extra_fields.append(f"{attr}={all_fields[attr]}")
+
+        # If there are extra fields, append them to the message
+        if extra_fields:
+            extra_text = "\n- " + "\n- ".join(extra_fields)
+            return base_msg + extra_text
+        return base_msg
+
+
 def __get_text_logger():
     logger = logging.getLogger("acontext-core")
     logger.setLevel(logging.INFO)
 
-    formatter = logging.Formatter(
+    formatter = CustomTextFormatter(
         f"{TerminalColorMarks.BOLD}{TerminalColorMarks.BLUE}%(name)s |{TerminalColorMarks.END} %(asctime)s - %(levelname)s - %(message)s"
     )
     handler = logging.StreamHandler()

src/server/core/acontext_core/telemetry/log.py

@@ -1,6 +1,6 @@
 import asyncio
 import traceback
-from acontext_core.entry import MQ_CLIENT, LOG, setup, cleanup
+from acontext_core.di import MQ_CLIENT, LOG, setup, cleanup
 
 
 async def app(scope, receive, send):