changes from review

Author: yrobla

src/codegate/api/v1.py

@@ -449,6 +449,7 @@ async def get_workspace_alerts_summary(workspace_name: str) -> v1_models.AlertSu
             malicious_packages=summary.total_packages_count,
             pii=summary.total_pii_count,
             secrets=summary.total_secrets_count,
+            total_alerts=summary.total_alerts,
         )
     except Exception:
         logger.exception("Error while getting alerts summary")
@@ -477,60 +478,75 @@ async def get_workspace_messages(
         raise HTTPException(status_code=500, detail="Internal server error")
 
     offset = (page - 1) * page_size
+    valid_conversations: List[v1_models.ConversationSummary] = []
+    fetched_prompts = 0
+
+    while len(valid_conversations) < page_size:
+        batch_size = page_size * 2  # Fetch more prompts to compensate for potential skips
+
+        prompts = await dbreader.get_prompts(
+            ws.id,
+            offset + fetched_prompts,
+            batch_size,
+            filter_by_ids,
+            list([AlertSeverity.CRITICAL.value]),  # TODO: Configurable severity
+            filter_by_alert_trigger_types,
+        )
+
+        # iterate for all prompts to compose the conversation summary
+        conversation_summaries: List[v1_models.ConversationSummary] = []
+        for prompt in prompts:
+            fetched_prompts += 1
+            if not prompt.request:
+                logger.warning(f"Skipping prompt {prompt.id}. Empty request field")
+                continue
+
+            messages, _ = await v1_processing.parse_request(prompt.request)
+            if not messages or len(messages) == 0:
+                logger.warning(f"Skipping prompt {prompt.id}. No messages found")
+                continue
+
+            # message is just the first entry in the request
+            message_obj = v1_models.ChatMessage(
+                message=messages[0], timestamp=prompt.timestamp, message_id=prompt.id
+            )
+
+            # count total alerts for the prompt
+            total_alerts_row = await dbreader.get_alerts_summary(prompt_id=prompt.id)
+
+            # get token usage for the prompt
+            prompts_outputs = await dbreader.get_prompts_with_output(prompt_id=prompt.id)
+            ws_token_usage = await v1_processing.parse_workspace_token_usage(prompts_outputs)
+
+            conversation_summary = v1_models.ConversationSummary(
+                chat_id=prompt.id,
+                prompt=message_obj,
+                provider=prompt.provider,
+                type=prompt.type,
+                conversation_timestamp=prompt.timestamp,
+                alerts_summary=v1_models.AlertSummary(
+                    malicious_packages=total_alerts_row.total_packages_count,
+                    pii=total_alerts_row.total_pii_count,
+                    secrets=total_alerts_row.total_secrets_count,
+                    total_alerts=total_alerts_row.total_alerts,
+                ),
+                total_alerts=total_alerts_row.total_alerts,
+                token_usage_agg=ws_token_usage,
+            )
+
+            valid_conversations.append(conversation_summary)
+            if len(valid_conversations) >= page_size:
+                break
 
-    prompts = await dbreader.get_prompts(
-        ws.id,
-        offset,
-        page_size,
-        filter_by_ids,
-        list([AlertSeverity.CRITICAL.value]),  # TODO: Configurable severity
-        filter_by_alert_trigger_types,
-    )
     # Fetch total message count
     total_count = await dbreader.get_total_messages_count_by_workspace_id(
         ws.id, AlertSeverity.CRITICAL.value
     )
 
-    # iterate for all prompts to compose the conversation summary
-    conversation_summaries: List[v1_models.ConversationSummary] = []
-    for prompt in prompts:
-        if not prompt.request:
-            logger.warning(f"Skipping prompt {prompt.id}. Empty request field")
-            continue
-
-        messages, _ = await v1_processing.parse_request(prompt.request)
-        if not messages or len(messages) == 0:
-            logger.warning(f"Skipping prompt {prompt.id}. No messages found")
-            continue
-
-        # message is just the first entry in the request
-        message_obj = v1_models.ChatMessage(
-            message=messages[0], timestamp=prompt.timestamp, message_id=prompt.id
-        )
-
-        # count total alerts for the prompt
-        total_alerts_row = await dbreader.get_alerts_summary(prompt_id=prompt.id)
-
-        # get token usage for the prompt
-        prompts_outputs = await dbreader.get_prompts_with_output(prompt_id=prompt.id)
-        ws_token_usage = await v1_processing.parse_workspace_token_usage(prompts_outputs)
-
-        conversation_summary = v1_models.ConversationSummary(
-            chat_id=prompt.id,
-            prompt=message_obj,
-            provider=prompt.provider,
-            type=prompt.type,
-            conversation_timestamp=prompt.timestamp,
-            total_alerts=total_alerts_row.total_alerts,
-            token_usage_agg=ws_token_usage,
-        )
-
-        conversation_summaries.append(conversation_summary)
-
     return v1_models.PaginatedMessagesResponse(
-        data=conversation_summaries,
+        data=valid_conversations,
         limit=page_size,
-        offset=(page - 1) * page_size,
+        offset=offset,
         total=total_count,
     )
 
@@ -543,7 +559,7 @@ async def get_workspace_messages(
 async def get_messages_by_prompt_id(
     workspace_name: str,
     prompt_id: str,
-) -> List[v1_models.Conversation]:
+) -> v1_models.Conversation:
     """Get messages for a workspace."""
     try:
         ws = await wscrud.get_workspace_by_name(workspace_name)
@@ -552,12 +568,13 @@ async def get_messages_by_prompt_id(
     except Exception:
         logger.exception("Error while getting workspace")
         raise HTTPException(status_code=500, detail="Internal server error")
-
     prompts_outputs = await dbreader.get_prompts_with_output(
         workspace_id=ws.id, prompt_id=prompt_id
     )
     conversations, _ = await v1_processing.parse_messages_in_conversations(prompts_outputs)
-    return conversations
+    if not conversations:
+        raise HTTPException(status_code=404, detail="Conversation not found")
+    return conversations[0]
 
 
 @v1.get(

src/codegate/api/v1_models.py

@@ -191,6 +191,7 @@ class AlertSummary(pydantic.BaseModel):
     malicious_packages: int
     pii: int
     secrets: int
+    total_alerts: int
 
 
 class PartialQuestionAnswer(pydantic.BaseModel):
@@ -225,7 +226,7 @@ class ConversationSummary(pydantic.BaseModel):
 
     chat_id: str
     prompt: ChatMessage
-    total_alerts: int
+    alerts_summary: AlertSummary
     token_usage_agg: Optional[TokenUsageAggregate]
     provider: Optional[str]
     type: QuestionType