Refactor AI service to use single endpoint for itinerary and leg enrichment

Copilot · zlendo1 · Copilot · commit d5fcf91abd59 · 2025-10-29T14:44:09.000Z
Co-authored-by: zlendo1 &lt;115471708+zlendo1@users.noreply.github.com&gt;
diff --git a/app/api/v1/endpoints/routes.py b/app/api/v1/endpoints/routes.py
@@ -61,25 +61,13 @@ async def search_routes(request: RouteSearchRequest) -> Any:
             first=request.num_itineraries,
         )
 
-        # Enhance each leg with AI insights (with graceful degradation)
+        # Enhance each itinerary with AI insights (with graceful degradation)
         for itinerary in itineraries:
-            for leg in itinerary.legs:
-                try:
-                    ai_insight = await ai_agents_service.get_leg_insight(leg)
-                    leg.ai_insight = ai_insight
-                except Exception as e:  # pylint: disable=broad-except
-                    # Gracefully degrade - log warning but continue without AI insight
-                    logger.warning("Failed to get AI insight for leg: %s", str(e))
-                    leg.ai_insight = None
-
-            # Get AI description for the overall itinerary
             try:
-                ai_description = await ai_agents_service.get_itinerary_insight(itinerary)
-                itinerary.ai_description = ai_description
+                await ai_agents_service.get_itinerary_insight(itinerary)
             except Exception as e:  # pylint: disable=broad-except
-                # Gracefully degrade - log warning but continue without AI description
-                logger.warning("Failed to get AI description for itinerary: %s", str(e))
-                itinerary.ai_description = None
+                # Gracefully degrade - log warning but continue without AI insights
+                logger.warning("Failed to get AI insights for itinerary: %s", str(e))
 
         logger.info("Route search successful: found %d itineraries", len(itineraries))
 
diff --git a/app/services/ai_agents_service.py b/app/services/ai_agents_service.py
@@ -65,66 +65,26 @@ async def health_check(self) -> ServiceHealth:
                 message=f"AI-agents API check failed: {str(e)}",
             )
 
-    async def get_leg_insight(self, leg: Leg) -> Optional[str]:
+    async def get_itinerary_insight(self, itinerary: Itinerary) -> None:
         """
-        Get AI-generated insight for a specific leg of the journey.
-
-        Args:
-            leg: The leg object containing journey segment information
-
-        Returns:
-            Optional[str]: AI-generated insight text, or None if unavailable
-
-        Raises:
-            No exceptions - gracefully degrades by returning None on errors
-        """
-        try:
-            client = self._get_client()
-
-            # Prepare request payload with leg information
-            payload: dict[str, str | int | float | dict[str, str] | None] = {
-                "mode": leg.mode.value,
-                "duration": leg.duration,
-                "distance": leg.distance,
-                "from_place": leg.from_place.name,
-                "to_place": leg.to_place.name,
-            }
-
-            # Include route information if available (e.g., bus number)
-            if leg.route:
-                payload["route"] = {
-                    "short_name": leg.route.short_name,
-                    "long_name": leg.route.long_name,
-                }
-
-            response = await client.post("/api/v1/insight/leg", json=payload)
-
-            if response.status_code == 200:
-                data = response.json()
-                return data.get("insight")
-
-            logger.warning("AI agents service returned non-200 status: %s", response.status_code)
-            return None
-
-        except httpx.TimeoutException:
-            logger.warning("AI agents service request timed out")
-            return None
-        except Exception as e:  # pylint: disable=broad-except
-            logger.warning("Failed to get AI insight: %s", str(e))
-            return None
-
-    async def get_itinerary_insight(self, itinerary: Itinerary) -> Optional[str]:
-        """
-        Get AI-generated insight for a complete itinerary.
+        Get AI-generated insights for a complete itinerary and enrich it with AI data.
+        
+        This method sends the entire itinerary to the AI service and receives back:
+        - ai_description: Overall description of the itinerary
+        - ai_insights: List of insights for each leg
+        
+        The method directly modifies the itinerary object in place, setting:
+        - itinerary.ai_description
+        - leg.ai_insight for each leg
 
         Args:
             itinerary: The itinerary object containing the complete journey information
 
         Returns:
-            Optional[str]: AI-generated insight text for the itinerary, or None if unavailable
+            None - modifies the itinerary object in place
 
         Raises:
-            No exceptions - gracefully degrades by returning None on errors
+            No exceptions - gracefully degrades by returning without modification on errors
         """
         try:
             client = self._get_client()
@@ -160,20 +120,29 @@ async def get_itinerary_insight(self, itinerary: Itinerary) -> Optional[str]:
 
             if response.status_code == 200:
                 data = response.json()
-                return data.get("insight")
+                
+                # Set the itinerary-level AI description
+                itinerary.ai_description = data.get("ai_description")
+                
+                # Set AI insights for each leg
+                ai_insights = data.get("ai_insights", [])
+                for idx, leg in enumerate(itinerary.legs):
+                    if idx < len(ai_insights):
+                        leg.ai_insight = ai_insights[idx]
+                    else:
+                        leg.ai_insight = None
+                
+                return
 
             logger.warning(
                 "AI agents service returned non-200 status for itinerary insight: %s",
                 response.status_code,
             )
-            return None
 
         except httpx.TimeoutException:
             logger.warning("AI agents service request timed out for itinerary insight")
-            return None
         except Exception as e:  # pylint: disable=broad-except
             logger.warning("Failed to get AI itinerary insight: %s", str(e))
-            return None
 
 
 # Singleton instance for dependency injection
diff --git a/tests/endpoints/test_routes.py b/tests/endpoints/test_routes.py
@@ -440,18 +440,13 @@ def test_search_routes_with_ai_insights_success(client: TestClient, sample_itine
         with patch("app.api.v1.endpoints.routes.ai_agents_service") as mock_ai_service:
             mock_routing_service.get_itinaries = AsyncMock(return_value=sample_itineraries)
 
-            # Mock AI service to return insights for legs
-            async def mock_get_insight(leg):
-                if leg.mode == TransportMode.WALK:
-                    return "Short walk to the bus stop."
-                return "Express bus with comfortable seats."
+            # Mock AI service to enrich itinerary in place
+            async def mock_get_itinerary_insight(itinerary):
+                itinerary.ai_description = "This route offers a good balance of walking and public transport."
+                itinerary.legs[0].ai_insight = "Short walk to the bus stop."
+                itinerary.legs[1].ai_insight = "Express bus with comfortable seats."
 
-            mock_ai_service.get_leg_insight = AsyncMock(side_effect=mock_get_insight)
-
-            # Mock AI service to return itinerary description
-            mock_ai_service.get_itinerary_insight = AsyncMock(
-                return_value="This route offers a good balance of walking and public transport."
-            )
+            mock_ai_service.get_itinerary_insight = AsyncMock(side_effect=mock_get_itinerary_insight)
 
             response = client.post(
                 "/api/v1/routes/search",
@@ -475,9 +470,6 @@ async def mock_get_insight(leg):
             # Check second leg (BUS)
             assert itinerary["legs"][1]["ai_insight"] == "Express bus with comfortable seats."
 
-            # Verify AI service was called for each leg
-            assert mock_ai_service.get_leg_insight.call_count == 2
-
             # Verify AI description for the itinerary
             assert (
                 itinerary["ai_description"]
@@ -493,9 +485,12 @@ def test_search_routes_with_ai_service_unavailable(client: TestClient, sample_it
         with patch("app.api.v1.endpoints.routes.ai_agents_service") as mock_ai_service:
             mock_routing_service.get_itinaries = AsyncMock(return_value=sample_itineraries)
 
-            # Mock AI service to return None (service unavailable)
-            mock_ai_service.get_leg_insight = AsyncMock(return_value=None)
-            mock_ai_service.get_itinerary_insight = AsyncMock(return_value=None)
+            # Mock AI service to do nothing (service unavailable)
+            async def mock_get_itinerary_insight_unavailable(itinerary):
+                # Service unavailable - does not modify itinerary
+                pass
+
+            mock_ai_service.get_itinerary_insight = AsyncMock(side_effect=mock_get_itinerary_insight_unavailable)
 
             response = client.post(
                 "/api/v1/routes/search",
@@ -522,22 +517,17 @@ def test_search_routes_with_ai_service_unavailable(client: TestClient, sample_it
 
 
 def test_search_routes_with_ai_service_partial_failure(client: TestClient, sample_itineraries):
-    """Test graceful degradation when AI service fails for some legs."""
+    """Test graceful degradation when AI service provides partial data."""
     with patch("app.api.v1.endpoints.routes.routing_service") as mock_routing_service:
         with patch("app.api.v1.endpoints.routes.ai_agents_service") as mock_ai_service:
             mock_routing_service.get_itinaries = AsyncMock(return_value=sample_itineraries)
 
-            # Mock AI service to succeed for first leg, fail for second
-            call_count = 0
-
-            async def mock_get_insight_partial(leg):
-                nonlocal call_count
-                call_count += 1
-                if call_count == 1:
-                    return "Short walk to the bus stop."
-                return None  # Simulate failure for second leg
+            # Mock AI service to provide partial data
+            async def mock_get_itinerary_insight_partial(itinerary):
+                # Only set description, not leg insights
+                itinerary.ai_description = "This is a good route."
 
-            mock_ai_service.get_leg_insight = AsyncMock(side_effect=mock_get_insight_partial)
+            mock_ai_service.get_itinerary_insight = AsyncMock(side_effect=mock_get_itinerary_insight_partial)
 
             response = client.post(
                 "/api/v1/routes/search",
@@ -550,10 +540,9 @@ async def mock_get_insight_partial(leg):
             assert response.status_code == 200
             data = response.json()
 
-            # Verify first leg has insight
-            assert data["itineraries"][0]["legs"][0]["ai_insight"] == "Short walk to the bus stop."
-
-            # Verify second leg has no insight (graceful degradation)
+            # Verify itinerary has description but legs don't have insights
+            assert data["itineraries"][0]["ai_description"] == "This is a good route."
+            assert data["itineraries"][0]["legs"][0]["ai_insight"] is None
             assert data["itineraries"][0]["legs"][1]["ai_insight"] is None
 
 
@@ -564,7 +553,6 @@ def test_search_routes_with_ai_service_exception(client: TestClient, sample_itin
             mock_routing_service.get_itinaries = AsyncMock(return_value=sample_itineraries)
 
             # Mock AI service to raise an exception
-            mock_ai_service.get_leg_insight = AsyncMock(side_effect=Exception("AI service error"))
             mock_ai_service.get_itinerary_insight = AsyncMock(
                 side_effect=Exception("AI service error")
             )
diff --git a/tests/services/test_ai_agents_service.py b/tests/services/test_ai_agents_service.py
