Update agents/routing_agent.py

agents/routing_agent.py

@@ -2,29 +2,168 @@ from prompts import ROUTER_PROMPT
 
 
 class RouterAgent:
+    """
+    Routes the user's request to the appropriate agent/tool.
+    """
 
-    def __init__(self, llm):
+    def __init__(self, llm=None):
         self.llm = llm
 
     def route(self, user_input: str) -> str:
-        """
-        Returns one of:
-        - generate
-        - edit
-        - analyze
-        - ocr
-        - pdf
-        """
 
-        prompt = ROUTER_PROMPT.format(text=user_input)
+        if not user_input:
+            return "chat"
 
-        response = self.llm.invoke(prompt)
+        prompt = user_input.lower().strip()
 
-        decision = response.content.strip().lower()
+        # -----------------------------
+        # ANALYZE IMAGE
+        # -----------------------------
+        analyze_keywords = [
+            "analyze",
+            "analyse",
+            "describe",
+            "explain",
+            "what is in",
+            "what's in",
+            "identify",
+            "detect",
+            "caption",
+            "inspect",
+            "look at",
+            "tell me about",
+            "analyze this image",
+            "describe this image"
+        ]
 
-        valid = ["generate", "edit", "analyze", "ocr", "pdf"]
+        # -----------------------------
+        # EDIT IMAGE
+        # -----------------------------
+        edit_keywords = [
+            "edit",
+            "change",
+            "modify",
+            "remove",
+            "replace",
+            "erase",
+            "add",
+            "swap",
+            "enhance",
+            "improve",
+            "background",
+            "blur",
+            "crop",
+            "resize",
+            "rotate",
+            "anime",
+            "cartoon",
+            "ghibli",
+            "pixar",
+            "watercolor",
+            "oil painting",
+            "black and white",
+            "make it",
+            "turn it into",
+            "convert"
+        ]
 
-        if decision not in valid:
+        # -----------------------------
+        # GENERATE IMAGE
+        # -----------------------------
+        generate_keywords = [
+            "generate",
+            "create",
+            "draw",
+            "design",
+            "imagine",
+            "make a",
+            "new image",
+            "text to image",
+            "logo",
+            "poster",
+            "banner"
+        ]
+
+        # -----------------------------
+        # OCR
+        # -----------------------------
+        ocr_keywords = [
+            "ocr",
+            "extract text",
+            "read text",
+            "scan text",
+            "recognize text"
+        ]
+
+        # -----------------------------
+        # PDF
+        # -----------------------------
+        pdf_keywords = [
+            "pdf",
+            "convert pdf",
+            "save pdf",
+            "export pdf"
+        ]
+
+        # -----------------------------
+        # ROUTING
+        # -----------------------------
+
+        # Analyze
+        if any(word in prompt for word in analyze_keywords):
+            return "analyze"
+
+        # OCR
+        if any(word in prompt for word in ocr_keywords):
+            return "ocr"
+
+        # PDF
+        if any(word in prompt for word in pdf_keywords):
+            return "pdf"
+
+        # Edit
+        if any(word in prompt for word in edit_keywords):
+            return "edit"
+
+        # Generate
+        if any(word in prompt for word in generate_keywords):
             return "generate"
 
-        return decision
+        # -----------------------------
+        # LLM Fallback (optional)
+        # -----------------------------
+        if self.llm:
+            try:
+                response = self.llm.invoke(f"""
+You are an AI routing assistant.
+
+Return ONLY ONE WORD from:
+
+analyze
+edit
+generate
+ocr
+pdf
+chat
+
+User Request:
+{user_input}
+""")
+
+                task = response.content.strip().lower()
+
+                if task in [
+                    "analyze",
+                    "edit",
+                    "generate",
+                    "ocr",
+                    "pdf",
+                    "chat"
+                ]:
+                    return task
+
+            except Exception:
+                pass
+
+        # Default
+        return "chat"