Merge pull request #36 from vaisu-bhut/dd_agent

GenAi error fixed, AgentCore is ready

Author: vaisu-bhut

services/security-agent/app/agents/nodes/security.py

@@ -276,48 +276,22 @@ async def security_check(state: AgentState) -> Dict[str, Any]:
         # Step 4: LLM Security Analysis (if enabled)
         llm_result = {}
         if settings.SECURITY_LLM_CHECK_ENABLED:
-            logger.info("CRASH_DEBUG: Starting LLM Security Check setup")
-
             # Reconstruct the chain manually as before
             try:
-                logger.info("CRASH_DEBUG: About to import ChatPromptTemplate")
                 from langchain_core.prompts import ChatPromptTemplate
-
-                logger.info("CRASH_DEBUG: ChatPromptTemplate imported successfully")
-
-                logger.info("CRASH_DEBUG: About to import JsonOutputParser")
                 from langchain_core.output_parsers import JsonOutputParser
 
-                logger.info("CRASH_DEBUG: JsonOutputParser imported successfully")
-
-                logger.info("CRASH_DEBUG: About to create prompt template")
                 prompt = ChatPromptTemplate.from_template(SECURITY_PROMPT)
-                logger.info("CRASH_DEBUG: Prompt template created successfully")
-
-                logger.info("CRASH_DEBUG: About to call get_llm()")
                 llm = get_llm()
-                logger.info("CRASH_DEBUG: get_llm() returned successfully")
-
-                logger.info("CRASH_DEBUG: About to create output parser")
                 parser = JsonOutputParser()
-                logger.info("CRASH_DEBUG: Output parser created successfully")
 
-                logger.info("CRASH_DEBUG: About to create chain with pipe operator")
                 chain = prompt | llm | parser
-                logger.info("CRASH_DEBUG: LLM chain initialized successfully")
             except Exception as e:
-                logger.error(f"CRASH_DEBUG: Failed to initialize LLM chain: {e}")
+                logger.error(f"Failed to initialize LLM chain: {e}")
                 raise
 
-            logger.info(
-                "CRASH_DEBUG: Preparing LLM input",
-                input_length=len(user_input),
-                threshold=settings.SECURITY_LLM_CHECK_THRESHOLD,
-            )
-
             # Using current time for detailed timing
             llm_start = time.perf_counter()
-            logger.info(f"CRASH_DEBUG: Invoking LLM chain... (Time: {llm_start})")
 
             try:
                 # Add a timeout to the LLM call if possible or just log around it
@@ -327,18 +301,13 @@ async def security_check(state: AgentState) -> Dict[str, Any]:
                         "threshold": settings.SECURITY_LLM_CHECK_THRESHOLD,
                     }
                 )
-                logger.info("CRASH_DEBUG: LLM chain returned successfully")
             except Exception as llm_exc:
-                logger.error(f"CRASH_DEBUG: LLM chain raised exception: {llm_exc}")
+                logger.error(f"LLM chain raised exception: {llm_exc}")
                 # Depending on policy, you might want to block here or set a default score
                 # For now, re-raise to be caught by the outer try/except
                 raise
 
             check_time_ms = (time.perf_counter() - llm_start) * 1000
-
-            logger.info(
-                f"CRASH_DEBUG: LLM Security Check completed in {check_time_ms:.2f}ms"
-            )
             metrics.record_stage_latency("llm_check", check_time_ms)
             metrics_builder.add_latency("llm_check", check_time_ms)
 

services/security-agent/app/main.py

@@ -18,35 +18,26 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(title=settings.PROJECT_NAME, version=settings.VERSION, lifespan=lifespan)
 
-print("DEBUG: Setting up telemetry...", flush=True)
 # Setup telemetry IMMEDIATELY
 setup_telemetry(app)
-print("DEBUG: Telemetry setup complete", flush=True)
 
 # Initialize global logger after telemetry
-print("DEBUG: Initializing structlog...", flush=True)
+# Initialize global logger after telemetry
 logger = structlog.get_logger()
-print("DEBUG: Structlog initialized", flush=True)
 
 # Import modules AFTER logging is configured
-print("DEBUG: Importing app.agents.graph...", flush=True)
+# Import modules AFTER logging is configured
 try:
     from app.agents.graph import agent_graph
-
-    print("DEBUG: Imported app.agents.graph successfully", flush=True)
 except Exception as e:
-    print(f"DEBUG: Failed to import app.agents.graph: {e}", flush=True)
     import traceback
 
     traceback.print_exc()
     raise
 
-print("DEBUG: Importing schemas and metrics...", flush=True)
 from app.schemas.security import ChatRequest, ChatResponse
 from app.core.metrics import get_security_metrics
 
-print("DEBUG: Imports completed successfully", flush=True)
-
 
 @app.get("/health")
 async def health_check():
@@ -65,7 +56,8 @@ async def chat(request: ChatRequest):
     """
     import time
 
-    print(f"CRASH_DEBUG: Chat request endpoint entered at {time.time()}", flush=True)
+    import time
+
     logger.info(
         "Chat request received",
         client_ip=request.client_ip,
@@ -83,12 +75,10 @@ async def chat(request: ChatRequest):
         "metrics_data": None,
     }
 
-    print("CRASH_DEBUG: Invoking agent graph...", flush=True)
     try:
         result = await agent_graph.ainvoke(initial_state)
-        print("CRASH_DEBUG: Agent graph returned successfully", flush=True)
     except Exception as e:
-        print(f"CRASH_DEBUG: Agent graph failed: {e}", flush=True)
+        logger.error(f"Agent graph execution failed: {e}")
         raise
 
     if result.get("is_blocked"):