fix (memory): logging for memory server

Author: itsskofficial

src/server/mcp_hub/memory/auth.py

@@ -1,8 +1,8 @@
-import logging
 from fastmcp import Context
 from fastmcp.exceptions import ToolError
+from fastmcp.utilities.logging import get_logger
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 def get_user_id_from_context(ctx: Context) -> str:
     """

src/server/mcp_hub/memory/db.py

@@ -1,10 +1,10 @@
 import os
-import logging
 import asyncpg
 import asyncio
 from dotenv import load_dotenv
 from typing import Dict
 
+from fastmcp.utilities.logging import get_logger
 from .constants import TOPICS
 
 # Load .env file for 'dev-local' environment.
@@ -14,7 +14,7 @@
     if os.path.exists(dotenv_path):
         load_dotenv(dotenv_path=dotenv_path)
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 # PostgreSQL connection details
 POSTGRES_USER = os.getenv("POSTGRES_USER")

src/server/mcp_hub/memory/llm.py

@@ -1,18 +1,19 @@
 import os
-import logging
 from qwen_agent.agents import Assistant
 from typing import Dict, Any
 import json
 
 from . import prompts
+from main.llm import run_agent_with_fallback
+from fastmcp.utilities.logging import get_logger
+
+logger = get_logger(__name__)
 
 # Use the main server's LLM config
 OPENAI_API_BASE_URL = os.getenv("OPENAI_API_BASE_URL", "http://localhost:11434/v1/") # Keep for consistency in config
 OPENAI_MODEL_NAME = os.getenv("OPENAI_MODEL_NAME", "qwen3:4b")
 OPENAI_API_KEYS = [key for key in [os.getenv("OPENAI_API_KEY"), os.getenv("OPENAI_API_KEY_FALLBACK_1"), os.getenv("OPENAI_API_KEY_FALLBACK_2")] if key]
 
-logger = logging.getLogger(__name__)
-
 def _get_base_llm_config() -> Dict[str, Any]:
     config = {
         'model': OPENAI_MODEL_NAME,
@@ -55,9 +56,6 @@ def run_agent_with_prompt(agent_config: Dict[str, Any], user_prompt: str) -> str
     messages = [{'role': 'user', 'content': user_prompt}]
     final_content = ""
 
-    # Import here to avoid circular dependency issues at module load time
-    from main.llm import run_agent_with_fallback
-
     try:
         for chunk in run_agent_with_fallback(system_message=system_message, function_list=[], messages=messages):
             if isinstance(chunk, list) and chunk and chunk[-1].get("role") == "assistant":

src/server/mcp_hub/memory/main.py

@@ -1,12 +1,12 @@
 import os
-import logging
 import asyncio
 from typing import Dict, Any, List
 from contextlib import asynccontextmanager
 from workers.tasks import cud_memory_task
 
 from dotenv import load_dotenv
 from fastmcp import FastMCP, Context
+from fastmcp.utilities.logging import configure_logging, get_logger
 
 from . import auth, utils, db
 
@@ -17,8 +17,9 @@
     if os.path.exists(dotenv_path):
         load_dotenv(dotenv_path=dotenv_path)
 
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+# --- Standardized Logging Setup ---
+configure_logging(level="INFO")
+logger = get_logger(__name__)
 
 # --- MCP Server Initialization ---
 @asynccontextmanager
@@ -116,5 +117,5 @@ async def delete_memory_by_source(ctx: Context, source_name: str) -> Dict[str, A
     host = os.getenv("MCP_SERVER_HOST", "127.0.0.1")
     port = int(os.getenv("MCP_SERVER_PORT", 8001))
 
-    print(f"Starting Memory MCP Server on http://{host}:{port}")
+    logger.info(f"Starting Memory MCP Server on http://{host}:{port}")
     mcp.run(transport="sse", host=host, port=port)

src/server/mcp_hub/memory/utils.py

@@ -1,25 +1,19 @@
 import os
 import json
-import logging
 import re
 from datetime import datetime, timedelta, timezone
 from typing import Dict, Any, List, Optional
 
 import numpy as np
 import google.generativeai as genai
 from pgvector.asyncpg import register_vector
+from fastmcp.utilities.logging import get_logger
 
 from . import db, llm
-from .prompts import (
-    fact_relevance_user_prompt_template,
-    fact_summarization_user_prompt_template,
-    fact_extraction_user_prompt_template,
-    cud_decision_user_prompt_template,
-    fact_analysis_user_prompt_template,
-)
+from .prompts import fact_analysis_user_prompt_template
+from main.config import EMBEDDING_MODEL_NAME, GEMINI_API_KEY
 
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 # --- Module-level state (initialized by lifespan event) ---
 embed_model_name: str = None
@@ -29,13 +23,11 @@
 def initialize_embedding_model():
     global embed_model_name
     if embed_model_name is None:
-        model_name = os.getenv("EMBEDDING_MODEL_NAME", "models/gemini-embedding-001")
-        api_key = os.getenv("GEMINI_API_KEY")
-        if not api_key:
-            raise ValueError("GEMINI_API_KEY environment variable not set.")
-        logger.info(f"Initializing embedding model: {model_name}")
-        genai.configure(api_key=api_key)
-        embed_model_name = model_name
+        if not GEMINI_API_KEY:
+            raise ValueError("GEMINI_API_KEY is not set in main.config.")
+        logger.info(f"Initializing embedding model: {EMBEDDING_MODEL_NAME}")
+        genai.configure(api_key=GEMINI_API_KEY)
+        embed_model_name = EMBEDDING_MODEL_NAME
 
 def initialize_agents():
     global agents
@@ -129,7 +121,7 @@ def _get_normalized_embedding(text: str, task_type: str) -> np.ndarray:
 
 async def search_memory(user_id: str, query: str) -> str: # noqa: E501
     """Searches memory by performing a semantic search, filtering for relevance, and summarizing results."""
-    print(f"Executing search_memory for user_id='{user_id}' with query: '{query}'")
+    logger.info(f"Executing search_memory for user_id='{user_id}' with query: '{query}'")
     pool = await db.get_db_pool()
     async with pool.acquire() as conn:
         await register_vector(conn)
@@ -148,7 +140,7 @@ async def search_memory(user_id: str, query: str) -> str: # noqa: E501
         )
 
         found_facts = [r['content'] for r in records]
-        print(f"Found {found_facts} potentially relevant facts from vector search.")
+        logger.info(f"Found {len(found_facts)} potentially relevant facts from vector search.")
 
     if not found_facts:
         logger.info("No facts found from vector search. Returning.")
@@ -170,10 +162,10 @@ async def search_memory(user_id: str, query: str) -> str: # noqa: E501
         except (json.JSONDecodeError, AttributeError):
             logger.warning(f"Could not parse relevance check response: {relevance_raw}")
 
-    print(f"Step 3/4: Found {relevant_facts} truly relevant facts after filtering.")
+    logger.info(f"Step 3/4: Found {len(relevant_facts)} truly relevant facts after filtering.")
 
     if not relevant_facts:
-        print("No relevant facts remained after filtering. Returning message to user.")
+        logger.info("No relevant facts remained after filtering. Returning message to user.")
         return "No relevant information found in your memory."
 
     # Step 4: Summarization
@@ -182,7 +174,7 @@ async def search_memory(user_id: str, query: str) -> str: # noqa: E501
     summary_raw = llm.run_agent_with_prompt(agents["fact_summarization"], prompt)
     summary = clean_llm_output(summary_raw)
 
-    print(f"Summary generated: {summary}")
+    logger.info(f"Summary generated: {summary[:100]}...")
 
     logger.info("Search complete. Returning summary.")
     return summary if isinstance(summary, str) and summary else "Could not generate a summary from the retrieved information."