Refactor Telegram agent initialization and shared state management; add error handling and model path resolution

.dockerignore

@@ -0,0 +1,10 @@
+.env
+*.env
+.venv/
+__pycache__/
+*.pyc
+models/
+*.gguf
+.qwen
+.git/
+.gitignore

main.py

@@ -95,31 +95,35 @@ def start_jira_server():
 def start_telegram_and_mcp():
     """Run the Telegram agent + MCP server in a dedicated event loop."""
     import shared_state
-    from telegram_agent import (
-        start_telegram_with_mcp, _init_telegram_bot, _init_llm,
-        _init_embedder, _init_pinecone, bot, llm,
-    )
+    import telegram_agent as tg
+    shared_state.startup_stage = "initializing"
 
     # Initialize components to detect LLM
     try:
-        _init_embedder()
-        _init_llm()
-        _init_pinecone()
-        _init_telegram_bot()
+        tg._init_embedder()
+        tg._init_llm()
+        tg._init_pinecone()
+        tg._init_telegram_bot()
     except Exception as e:
         state.error = f"Failed to init components: {str(e)}"
+        shared_state.telegram_error = state.error
+        shared_state.startup_stage = "failed"
         print(f"[orchestrator] ERROR: {state.error}")
         return
 
-    # Set refs
-    shared_state.llm_instance = llm
-    shared_state.telegram_bot = bot
+    # Set refs from live module globals (not from stale imported names)
+    shared_state.llm_instance = tg.llm
+    shared_state.telegram_bot = tg.bot
+    shared_state.model_path = tg.resolved_model_path or tg.MODEL_PATH
+    shared_state.telegram_error = None
 
     # Run the MCP + Telegram loop
     try:
-        asyncio.run(start_telegram_with_mcp())
+        asyncio.run(tg.start_telegram_with_mcp())
     except Exception as e:
         state.error = f"Telegram/MCP failed: {str(e)}"
+        shared_state.telegram_error = state.error
+        shared_state.startup_stage = "failed"
         print(f"[orchestrator] ERROR: {state.error}")
 
 
@@ -185,7 +189,7 @@ def status_mcp():
 
     if session is None or tools is None:
         return {"service": "mcp_server", "status": "down", "tools_count": 0,
-                "tools": [], "error": "MCP server not initialized"}
+                "tools": [], "error": shared_state.telegram_error or "MCP server not initialized"}
     return {"service": "mcp_server", "status": "live", "tools_count": len(tools),
             "tools": [t["function"]["name"] for t in tools]}
 
@@ -200,11 +204,15 @@ def status_telegram():
     bot_status = "live" if b is not None else "down"
     llm_status = "loaded" if l is not None else "not_loaded"
 
-    model_path = os.environ.get("MODEL_PATH", os.path.join(BASE_DIR, "models", "microsoft_Phi-4-mini-instruct-Q4_K_M.gguf"))
+    model_path = shared_state.model_path or os.environ.get(
+        "MODEL_PATH", os.path.join(BASE_DIR, "models", "microsoft_Phi-4-mini-instruct-Q4_K_M.gguf")
+    )
     model_exists = os.path.exists(model_path)
 
     return {"service": "telegram_bot", "bot_status": bot_status, "llm_status": llm_status,
-            "model_path": model_path, "model_downloaded": model_exists}
+            "model_path": model_path, "model_downloaded": model_exists,
+            "startup_stage": shared_state.startup_stage,
+            "error": shared_state.telegram_error}
 
 
 @app.get("/status")

shared_state.py

@@ -6,3 +6,6 @@ mcp_session = None
 openai_tools = None
 telegram_bot = None
 llm_instance = None
+model_path = None
+telegram_error = None
+startup_stage = "not_started"

telegram_agent.py

@@ -41,6 +41,8 @@ PINECONE_INDEX   = os.environ.get("PINECONE_INDEX", "jira-agent-memory")
 MCP_SERVER_SCRIPT = os.path.join(os.path.dirname(__file__), "jira_mcp_server.py")
 EMBED_DIM = 768
 MODEL_PATH = os.path.join(os.path.dirname(__file__), "models", "microsoft_Phi-4-mini-instruct-Q4_K_M.gguf")
+MODEL_REPO_ID = os.environ.get("MODEL_REPO_ID", "bartowski/microsoft_Phi-4-mini-instruct-GGUF")
+MODEL_FILENAME = os.environ.get("MODEL_FILENAME", "microsoft_Phi-4-mini-instruct-Q4_K_M.gguf")
 
 # ── Globals (lazy-initialized) ────────────────────────────────────────────────
 bot: Optional[telebot.TeleBot] = None
@@ -51,6 +53,7 @@ user_sessions = {}
 loop_ref = None
 mcp_session_ref = None
 openai_tools_ref = None
+resolved_model_path = None
 
 
 # ── Lazy initializers ────────────────────────────────────────────────────────
@@ -80,7 +83,7 @@ def embed(text: str) -> list[float]:
 
 
 def _init_llm() -> Llama:
-    global llm
+    global llm, resolved_model_path
     if llm is not None:
         return llm
 
@@ -89,9 +92,10 @@ def _init_llm() -> Llama:
     from huggingface_hub import hf_hub_download
     print("  Resolving model from HuggingFace cache...\n")
     cached_path = hf_hub_download(
-        repo_id="bartowski/microsoft_Phi-4-mini-instruct-GGUF",
-        filename="microsoft_Phi-4-mini-instruct-Q4_K_M.gguf",
+        repo_id=MODEL_REPO_ID,
+        filename=MODEL_FILENAME,
     )
+    resolved_model_path = cached_path
     print(f"  Model cached at: {cached_path}")
     print(f"  Loading model...\n")
 
@@ -108,35 +112,55 @@ def _init_pinecone():
     if pine_index is not None:
         return pine_index
     if not PINECONE_API_KEY:
-        raise ValueError("Missing PINECONE_API_KEY")
-    pc = Pinecone(api_key=PINECONE_API_KEY)
-    existing = [i.name for i in pc.list_indexes()]
-    if PINECONE_INDEX not in existing:
-        print(f"  Creating Pinecone index '{PINECONE_INDEX}'...")
-        pc.create_index(name=PINECONE_INDEX, dimension=EMBED_DIM, metric="cosine",
-                        spec=ServerlessSpec(cloud="aws", region="us-east-1"))
-        print("  Index created.\n")
-    pine_index = pc.Index(PINECONE_INDEX)
+        print("  Pinecone disabled: missing PINECONE_API_KEY.\n")
+        return None
+    try:
+        pc = Pinecone(api_key=PINECONE_API_KEY)
+        existing = [i.name for i in pc.list_indexes()]
+        if PINECONE_INDEX not in existing:
+            print(f"  Creating Pinecone index '{PINECONE_INDEX}'...")
+            pc.create_index(name=PINECONE_INDEX, dimension=EMBED_DIM, metric="cosine",
+                            spec=ServerlessSpec(cloud="aws", region="us-east-1"))
+            print("  Index created.\n")
+        pine_index = pc.Index(PINECONE_INDEX)
+    except Exception as e:
+        print(f"  Pinecone disabled due to init error: {str(e)}\n")
+        pine_index = None
     return pine_index
 
 
 def store(user_input, agent_reply, tool_name, tool_args, tool_result):
+    if pine_index is None or _embedder is None:
+        return
     doc = f"User: {user_input}\nAgent: {agent_reply}\nTool: {tool_name}\nArgs: {json.dumps(tool_args)}\nResult: {json.dumps(tool_result)}"
-    pine_index.upsert(vectors=[{
-        "id": str(uuid.uuid4()), "values": embed(doc),
-        "metadata": {"doc": doc, "user_input": user_input, "agent_reply": agent_reply,
-                      "tool_name": tool_name, "tool_args": json.dumps(tool_args),
-                      "tool_result": json.dumps(tool_result),
-                      "timestamp": datetime.datetime.utcnow().isoformat()},
-    }])
+    try:
+        pine_index.upsert(vectors=[{
+            "id": str(uuid.uuid4()), "values": embed(doc),
+            "metadata": {"doc": doc, "user_input": user_input, "agent_reply": agent_reply,
+                          "tool_name": tool_name, "tool_args": json.dumps(tool_args),
+                          "tool_result": json.dumps(tool_result),
+                          "timestamp": datetime.datetime.utcnow().isoformat()},
+        }])
+    except Exception as e:
+        print(f"  Warning: failed to store memory: {str(e)}")
 
 
 def recall(user_input: str, top_k: int = 2) -> str:
-    stats = pine_index.describe_index_stats()
+    if pine_index is None or _embedder is None:
+        return ""
+    try:
+        stats = pine_index.describe_index_stats()
+    except Exception as e:
+        print(f"  Warning: failed to query memory stats: {str(e)}")
+        return ""
     total = stats.get("total_vector_count", 0)
     if total == 0:
         return ""
-    results = pine_index.query(vector=embed(user_input), top_k=min(top_k, total), include_metadata=True)
+    try:
+        results = pine_index.query(vector=embed(user_input), top_k=min(top_k, total), include_metadata=True)
+    except Exception as e:
+        print(f"  Warning: failed to query memory: {str(e)}")
+        return ""
     chunks = [f"[past]\n{m['metadata'].get('user_input','')} → {m['metadata'].get('tool_name','')}"
               for m in results.get("matches", []) if m.get('score', 0) > 0.5]
     return "\nRECENT SIMILAR (context only):\n" + "\n---\n".join(chunks) + "\n" if chunks else ""
@@ -291,17 +315,28 @@ async def start_telegram_with_mcp():
     global mcp_session_ref, openai_tools_ref, loop_ref
     loop_ref = asyncio.get_event_loop()
 
+    import shared_state
+    shared_state.startup_stage = "init_telegram"
     # Lazy init all components
     _init_telegram_bot()
+
+    shared_state.startup_stage = "init_embedder"
     _init_embedder()
+
+    shared_state.startup_stage = "init_llm"
     _init_llm()
+
+    shared_state.startup_stage = "init_pinecone"
     _init_pinecone()
+
+    shared_state.startup_stage = "register_handlers"
     _register_handlers()
 
     server_params = StdioServerParameters(command=sys.executable, args=[MCP_SERVER_SCRIPT], env={**os.environ})
 
     async with stdio_client(server_params) as (read, write):
         async with ClientSession(read, write) as session:
+            shared_state.startup_stage = "init_mcp_session"
             await session.initialize()
             tools_resp = await session.list_tools()
             mcp_tools = tools_resp.tools
@@ -313,10 +348,17 @@ async def start_telegram_with_mcp():
             shared_state.openai_tools = openai_tools
             shared_state.telegram_bot = bot
             shared_state.llm_instance = llm
+            shared_state.model_path = resolved_model_path or MODEL_PATH
+            shared_state.telegram_error = None
             mcp_session_ref = session
             openai_tools_ref = openai_tools
 
-            stats = pine_index.describe_index_stats()
+            stats = {"total_vector_count": 0}
+            if pine_index is not None:
+                try:
+                    stats = pine_index.describe_index_stats()
+                except Exception as e:
+                    print(f"  Warning: failed to read Pinecone stats: {str(e)}")
             print("=" * 60)
             print(f"  Telegram + MCP | Tools: {[t.name for t in mcp_tools]}")
             print(f"  Memory: {stats.get('total_vector_count', 0)} vectors")
@@ -328,6 +370,7 @@ async def start_telegram_with_mcp():
 
             t = threading.Thread(target=run_telegram, daemon=True)
             t.start()
+            shared_state.startup_stage = "running"
 
             try:
                 while True:
@@ -335,6 +378,7 @@ async def start_telegram_with_mcp():
             except (KeyboardInterrupt, SystemExit):
                 print("\nShutting down Telegram...")
                 bot.stop_polling()
+                shared_state.startup_stage = "stopped"
 
 
 # ── Standalone run ────────────────────────────────────────────────────────────