fix: graph

sahayak/agent.py

@@ -4,6 +4,7 @@
 from sahayak import prompt
 from sahayak.subagents.planner.agent import planner_agent
 from sahayak.tools.memory import _load_initial_state
+from sahayak.subagents.graph.agent import graph_agent
 
 print(prompt.ROOT_INSTRUCTION)
 root_agent = Agent(
@@ -14,7 +15,7 @@
         "including lesson planning, content creation, and classroom management."
     ),
     instruction=prompt.ROOT_INSTRUCTION,
-    sub_agents=[planner_agent],
+    sub_agents=[planner_agent, graph_agent],
     # before_agent_callback=_load_initial_state,
 )
 

sahayak/prompt.py

@@ -29,10 +29,11 @@
 # """
 
 ROOT_INSTRUCTION = """
-        "You are an educational assistant coordinator. Your role is to orchestrate tasks between three main agents:\n\n"
+        "You are an educational assistant coordinator. Your role is to orchestrate tasks between four main agents:\n\n"
         "1. Planner Agent: Handles curriculum planning, lesson design, and content creation\n"
         "2. Task and Reminder Agent: Manages reminders, tasks, and scheduling\n"
         "3. Academic Helper: Processes documents, assessments, and performance analysis\n\n"
+        "4. Graph Visualizer: Queries student performance in Neo4j graph database\n\n"
         "Your responsibilities include:\n"
         "- Understanding user requests and routing them to appropriate agents\n"
         "- Coordinating complex tasks that require multiple agents\n"

sahayak/subagents/graph/agent.py

@@ -0,0 +1,14 @@
+from google.adk.agents import Agent
+from sahayak.tools.graph import GraphVisualizer
+from sahayak.subagents.graph.prompt import GRAPH_AGENT_INSTR
+
+# Create the agent with the proper tool instance
+graph_agent = Agent(
+    name="graph_visualizer",
+    model="gemini-2.5-flash",
+    description="Queries student performance in Neo4j.",
+    tools=[GraphVisualizer()],
+    instruction=GRAPH_AGENT_INSTR,
+    disallow_transfer_to_parent=True,
+    disallow_transfer_to_peers=True,
+)

sahayak/subagents/graph/prompt.py

@@ -0,0 +1,31 @@
+GRAPH_AGENT_INSTR = """
+You are a student performance analysis assistant that queries Neo4j graph database.
+
+Workflow:
+1. When asked about student performance, use the `graph_visualizer` tool
+2. After the tool executes, you will receive the results  
+3. Based on the query type, format your response accordingly:
+
+**For Team Formation Queries** (detect keywords: "form team", "create team", "make team", "team up"):
+- Return ONLY compact JSON without line breaks or extra formatting
+- Example format: {"teams": [{"name": "Study Team 1", "type": "study_buddy", "members": [{"name": "Student1", "strengths": ["Topic1"], "needs_help": ["Topic2"]}, {"name": "Student2", "strengths": ["Topic2"], "needs_help": ["Topic1"]}], "pairing_logic": "Explanation of how they help each other."}]}
+
+**For All Other Queries** (statistics, individual performance, etc.):
+- Return natural language text
+- Be clear and concise
+
+Examples:
+User: "Form teams based on performance in topics 'light' and 'plants' for grade 6"
+Assistant: {"teams": [{"name": "Study Team 1", "type": "study_buddy", "members": [{"name": "Tanya Patel", "strengths": ["Light"], "needs_help": ["Human Body"]}, {"name": "Arjun Kumar", "strengths": ["Human Body"], "needs_help": ["Light"]}], "pairing_logic": "Tanya Patel excels in Light and can help Arjun Kumar, who needs improvement in Light. Arjun Kumar is strong in Human Body and can assist Tanya Patel, who needs help in Human Body."}, {"name": "Study Team 2", "type": "study_buddy", "members": [{"name": "Priya Sharma", "strengths": ["Human Body"], "needs_help": ["Light"]}, {"name": "Ananya Naik", "strengths": ["Light"], "needs_help": ["Human Body"]}], "pairing_logic": "Priya Sharma is strong in Human Body and can help Ananya Naik. Ananya Naik is strong in Light and can help Priya Sharma, creating a balanced study partnership."}]}
+
+User: "Find the top student in topic 'light' for grade 6"  
+Assistant: The top student in Light (Grade 6) is Tanya Patel with a score of 10.
+
+User: "Show me statistics for Light topic in grade 6"
+Assistant: Here are the statistics for the 'Light' topic in Grade 6: Highest Score: 10 (Tanya Patel), Lowest Score: 3 (Neha Gupta), Average Score: 6.5
+
+CRITICAL RULES:
+- For team queries: Return ONLY compact JSON on a single line, nothing else
+- For other queries: Return natural language
+- Never return tool call syntax or markdown
+"""

sahayak/tools/graph.py

@@ -0,0 +1,314 @@
+from google.adk.tools import BaseTool
+from google.genai import types
+from sahayak.tools.graph_state import GraphQueryState
+from neo4j import AsyncGraphDatabase
+import logging
+import os
+from typing import Optional, List, Dict, Any
+import json
+
+logger = logging.getLogger(__name__)
+
+class GraphVisualizer(BaseTool):
+    def __init__(self):
+        super().__init__(
+            name="graph_visualizer",
+            description="Executes Neo4j graph queries to analyze student performance."
+        )
+
+    def run(self, user_intent: str, topic_a: str, grade: Optional[str] = None, topic_b: Optional[str] = None):
+        """
+        Sync wrapper for async Neo4j query execution.
+        """
+        logger.info(f"🔧 GraphVisualizer.run() called with:")
+        logger.info(f"   user_intent={user_intent}")
+        logger.info(f"   topic_a={topic_a}")
+        logger.info(f"   grade={grade}")
+        logger.info(f"   topic_b={topic_b}")
+
+        try:
+            # Build state dict
+            state = {
+                "user_intent": user_intent,
+                "topic_a": topic_a,
+                "grade": grade,
+                "topic_b": topic_b
+            }
+
+            # Validate state
+            parsed_state = GraphQueryState(**{k: v for k, v in state.items() if v is not None})
+
+        except Exception as e:
+            error_msg = f"Invalid parameters: {e}"
+            logger.error(error_msg)
+            return types.Content(
+                role="tool",
+                parts=[types.Part(text=error_msg)]
+            )
+
+        # Since we're in sync context, run the async method in a new event loop
+        import asyncio
+        import nest_asyncio  # ✅ Handle nested event loops
+
+        try:
+            # Allow nested event loops
+            nest_asyncio.apply()
+            result = asyncio.run(self._run_query(parsed_state))
+            if isinstance(result, dict):
+                json_string = json.dumps(result, indent=2) 
+                logger.info(f"🔧 GraphVisualizer.run() returning JSON string: {json_string[:200]}...") 
+                return types.Content(
+                    role="tool",
+                    parts=[types.Part(text=json_string)] 
+                )
+            else:
+                return types.Content(
+                    role="tool",
+                    parts=[types.Part(text=result)]
+                )
+        except Exception as e:
+            error_msg = f"Query failed: {e}"
+            logger.exception("Neo4j query failed")
+            return types.Content(
+                role="tool",
+                parts=[types.Part(text=error_msg)]
+            )
+
+    async def _run_query(self, parsed_state: GraphQueryState) -> str:
+        """Async method to run Neo4j queries."""
+        logger.info("🚀 Executing Neo4j query...")
+
+        driver = AsyncGraphDatabase.driver(
+            os.getenv("NEO4J_URI", "bolt://localhost:7687"),
+            auth=(os.getenv("NEO4J_USERNAME"), os.getenv("NEO4J_PASSWORD")),
+        )
+
+        try:
+            async with driver.session(database="neo4j") as session:
+                if parsed_state.user_intent == "find_highest":
+                    return await self._find_highest_student(session, parsed_state)
+
+                elif parsed_state.user_intent == "find_top_students":
+                    return await self._find_top_students(session, parsed_state)
+
+                elif parsed_state.user_intent == "form_teams":
+                    return await self._form_teams(session, parsed_state)
+
+                elif parsed_state.user_intent == "get_statistics":
+                    return await self._get_topic_statistics(session, parsed_state)
+
+                elif parsed_state.user_intent == "compare_topics":
+                    return await self._compare_topics(session, parsed_state)
+
+                else:
+                    return "Unsupported or incomplete intent"
+        except Exception as e:
+            logger.exception("Neo4j query failed")
+            return f"Query execution failed: {e}"
+        finally:
+            await driver.close()
+
+    async def _find_highest_student(self, session, parsed_state: GraphQueryState) -> str:
+        """Find the highest scoring student."""
+        result = await session.run("""
+            MATCH (s:Student)-[r:SCORED_IN]->(t:Topic {name: $topic_a, grade: $grade})
+            RETURN s.name AS student_name, r.score AS score
+            ORDER BY r.score DESC
+            LIMIT 1
+        """, {
+            "topic_a": parsed_state.topic_a.title(),
+            "grade": parsed_state.grade
+        })
+        record = await result.single()
+        if record:
+            return f"🏆 Top student in {parsed_state.topic_a} (Grade {parsed_state.grade}): {record['student_name']} (Score: {record['score']})"
+        return f"No student found for {parsed_state.topic_a} topic in Grade {parsed_state.grade}"
+
+    async def _find_top_students(self, session, parsed_state: GraphQueryState) -> str:
+        """Find top 5 students for a topic."""
+        result = await session.run("""
+            MATCH (s:Student)-[r:SCORED_IN]->(t:Topic {name: $topic_a, grade: $grade})
+            RETURN s.name AS student_name, r.score AS score
+            ORDER BY r.score DESC
+            LIMIT 5
+        """, {
+            "topic_a": parsed_state.topic_a.title(),
+            "grade": parsed_state.grade
+        })
+        students = await result.data()
+        if students:
+            rankings = []
+            for i, student in enumerate(students, 1):
+                rankings.append(f"{i}. {student['student_name']}: {student['score']}")
+            return f"🏅 Top 5 students in {parsed_state.topic_a} (Grade {parsed_state.grade}):\n" + "\n".join(rankings)
+        return f"No students found for {parsed_state.topic_a} topic in Grade {parsed_state.grade}"
+
+    async def _form_teams(self, session, parsed_state: GraphQueryState) -> Dict[str, Any]:
+        """Enhanced team formation returning structured JSON."""
+        if parsed_state.topic_b:
+            # Dual topic team formation
+            logger.info(f"Forming teams based on {parsed_state.topic_a} + {parsed_state.topic_b}")
+
+            topic_check = await session.run("""
+                MATCH (t:Topic {grade: $grade})
+                WHERE t.name IN [$topic_a, $topic_b]
+                RETURN t.name AS topic_name
+            """, {
+                "topic_a": parsed_state.topic_a.title(),
+                "topic_b": parsed_state.topic_b.title(),
+                "grade": parsed_state.grade
+            })
+            existing_topics = [record["topic_name"] async for record in topic_check]
+
+            if len(existing_topics) < 2:
+                missing = set([parsed_state.topic_a.title(), parsed_state.topic_b.title()]) - set(existing_topics)
+                return {"error": f"Cannot form teams: Topics not found: {', '.join(missing)}"}
+
+            result = await session.run("""
+                MATCH (s:Student)-[r1:SCORED_IN]->(t1:Topic {name: $topic_a, grade: $grade}),
+                    (s)-[r2:SCORED_IN]->(t2:Topic {name: $topic_b, grade: $grade})
+                RETURN s.name AS student_name, r1.score AS score_a, r2.score AS score_b
+                ORDER BY (r1.score + r2.score) DESC
+                LIMIT 8
+            """, {
+                "topic_a": parsed_state.topic_a.title(),
+                "topic_b": parsed_state.topic_b.title(),
+                "grade": parsed_state.grade
+            })
+            students = await result.data()
+
+            if not students:
+                return {"error": f"No students found who have scores for both '{parsed_state.topic_a}' and '{parsed_state.topic_b}'"}
+        else:
+            # Single topic team formation
+            logger.info(f"Forming teams based on {parsed_state.topic_a} only")
+
+            topic_check = await session.run("""
+                MATCH (t:Topic {name: $topic_a, grade: $grade})
+                RETURN t.name AS topic_name
+            """, {
+                "topic_a": parsed_state.topic_a.title(),
+                "grade": parsed_state.grade
+            })
+            topic_exists = await topic_check.single()
+            if not topic_exists:
+                return {"error": f"Cannot form teams: Topic '{parsed_state.topic_a}' not found for grade {parsed_state.grade}"}
+
+            result = await session.run("""
+                MATCH (s:Student)-[r:SCORED_IN]->(t:Topic {name: $topic_a, grade: $grade})
+                RETURN s.name AS student_name, r.score AS score
+                ORDER BY r.score DESC
+                LIMIT 8
+            """, {
+                "topic_a": parsed_state.topic_a.title(),
+                "grade": parsed_state.grade
+            })
+            students = await result.data()
+
+            if not students:
+                return {"error": f"No students found with scores for '{parsed_state.topic_a}' topic"}
+
+        # ✅ Build structured team JSON
+        teams = await self._create_balanced_teams(students, parsed_state.topic_b is not None)
+        return {
+            "type": "study_buddy",
+            "teams": teams
+        }
+
+
+    async def _create_balanced_teams(self, students: List[Dict], is_dual_topic: bool) -> str:
+        """Create balanced teams with better distribution."""
+        if len(students) == 1:
+            student = students[0]
+            if is_dual_topic:
+                return f"Team 1: {student['student_name']} (Scores: {student['score_a']}, {student['score_b']}) (solo)"
+            else:
+                score = student.get('score', student.get('score_a', 'N/A'))
+                return f"Team 1: {student['student_name']} (Score: {score}) (solo)"
+
+        # Create teams of 2-3 students for better collaboration
+        teams = []
+        team_size = 3 if len(students) >= 6 else 2
+
+        for i in range(0, len(students), team_size):
+            team_students = students[i:i+team_size]
+            team_members = []
+
+            for student in team_students:
+                if is_dual_topic and 'score_a' in student:
+                    team_members.append(f"{student['student_name']} ({student['score_a']}, {student['score_b']})")
+                else:
+                    score = student.get('score', student.get('score_a', 'N/A'))
+                    team_members.append(f"{student['student_name']} ({score})")
+
+            team_str = f"Team {(i//team_size) + 1}: {', '.join(team_members)}"
+            teams.append(team_str)
+
+        return "\n".join(teams)
+
+    async def _get_topic_statistics(self, session, parsed_state: GraphQueryState) -> str:
+        """Get statistics for a topic."""
+        result = await session.run("""
+            MATCH (s:Student)-[r:SCORED_IN]->(t:Topic {name: $topic_a, grade: $grade})
+            RETURN 
+                count(r) AS total_scores,
+                avg(r.score) AS average_score,
+                min(r.score) AS min_score,
+                max(r.score) AS max_score
+        """, {
+            "topic_a": parsed_state.topic_a.title(),
+            "grade": parsed_state.grade
+        })
+        record = await result.single()
+        if record:
+            return (f"📊 Statistics for {parsed_state.topic_a} (Grade {parsed_state.grade}):\n"
+                   f"• Total students: {record['total_scores']}\n"
+                   f"• Average score: {record['average_score']:.1f}\n"
+                   f"• Highest score: {record['max_score']}\n"
+                   f"• Lowest score: {record['min_score']}")
+        return f"No data found for {parsed_state.topic_a} topic in Grade {parsed_state.grade}"
+
+    async def _compare_topics(self, session, parsed_state: GraphQueryState) -> str:
+        """Compare performance between two topics."""
+        if not parsed_state.topic_b:
+            return "Please specify two topics to compare"
+
+        # Get average scores for both topics
+        result_a = await session.run("""
+            MATCH (s:Student)-[r:SCORED_IN]->(t:Topic {name: $topic_a, grade: $grade})
+            RETURN avg(r.score) AS avg_score, count(r) AS student_count
+        """, {
+            "topic_a": parsed_state.topic_a.title(),
+            "grade": parsed_state.grade
+        })
+        record_a = await result_a.single()
+
+        result_b = await session.run("""
+            MATCH (s:Student)-[r:SCORED_IN]->(t:Topic {name: $topic_b, grade: $grade})
+            RETURN avg(r.score) AS avg_score, count(r) AS student_count
+        """, {
+            "topic_b": parsed_state.topic_b.title(),
+            "grade": parsed_state.grade
+        })
+        record_b = await result_b.single()
+
+        if not record_a or not record_b:
+            return "Insufficient data for comparison"
+
+        avg_a = record_a['avg_score'] or 0
+        avg_b = record_b['avg_score'] or 0
+
+        comparison = "📈 Topic Comparison:\n"
+        comparison += f"• {parsed_state.topic_a}: {avg_a:.1f} avg (from {record_a['student_count']} students)\n"
+        comparison += f"• {parsed_state.topic_b}: {avg_b:.1f} avg (from {record_b['student_count']} students)\n"
+
+        if avg_a > avg_b:
+            diff = avg_a - avg_b
+            comparison += f"• {parsed_state.topic_a} performed better by {diff:.1f} points"
+        elif avg_b > avg_a:
+            diff = avg_b - avg_a
+            comparison += f"• {parsed_state.topic_b} performed better by {diff:.1f} points"
+        else:
+            comparison += "• Both topics showed similar performance"
+
+        return comparison