feat: better logging

Author: nicosammito

src/endpoint/generation/generate_endpoint.py

@@ -71,30 +71,6 @@ def Prompt(self, request: pb2.PromptRequest, context) -> pb2.FlowResponse:
                 details=f"The specified model_identifier '{request.model_identifier}' does not exist. Please provide a valid model_identifier for flow generation."
             )
 
-        stored_functions_count = len(self.function_store.get_all(group_identifier=str(request.project_id)))
-        stored_flow_types_count = len(self.flow_type_store.get_all(group_identifier=str(request.project_id)))
-        stored_few_shots_count = len(self.few_shots_store.get_all(group_identifier="global"))
-        log.info(
-            f"[Prompt] Store state — project={request.project_id} "
-            f"functions={stored_functions_count} "
-            f"flow_types={stored_flow_types_count} "
-            f"few_shots={stored_few_shots_count}"
-        )
-
-        if stored_functions_count <= 0 and (len(request.functions) <= 0):
-            log.warning(f"[Prompt] Rejected — no functions for project={request.project_id}")
-            context.abort(
-                code=grpc.StatusCode.ABORTED,
-                details="No functions found for the given project_id. Please add functions before requesting a prompt generation."
-            )
-
-        if stored_flow_types_count <= 0 and (len(request.flow_types) <= 0):
-            log.warning(f"[Prompt] Rejected — no flow types for project={request.project_id}")
-            context.abort(
-                code=grpc.StatusCode.ABORTED,
-                details="No flow types found for the given project_id. Please add flow_types before requesting a prompt generation."
-            )
-
         functions = [
             map_to_function_schema(fn)
             for fn in request.functions
@@ -124,6 +100,30 @@ def Prompt(self, request: pb2.PromptRequest, context) -> pb2.FlowResponse:
                 data_types=data_types
             )
 
+        stored_functions_count = len(self.function_store.get_all(group_identifier=str(request.project_id)))
+        stored_flow_types_count = len(self.flow_type_store.get_all(group_identifier=str(request.project_id)))
+        stored_few_shots_count = len(self.few_shots_store.get_all(group_identifier="global"))
+        log.info(
+            f"[Prompt] Store state — project={request.project_id} "
+            f"functions={stored_functions_count} "
+            f"flow_types={stored_flow_types_count} "
+            f"few_shots={stored_few_shots_count}"
+        )
+
+        if stored_functions_count <= 0:
+            log.warning(f"[Prompt] Rejected — no functions for project={request.project_id}")
+            context.abort(
+                code=grpc.StatusCode.ABORTED,
+                details="No functions found for the given project_id. Please add functions before requesting a prompt generation."
+            )
+
+        if stored_flow_types_count <= 0:
+            log.warning(f"[Prompt] Rejected — no flow types for project={request.project_id}")
+            context.abort(
+                code=grpc.StatusCode.ABORTED,
+                details="No flow types found for the given project_id. Please add flow_types before requesting a prompt generation."
+            )
+
         prompt_functions = self.function_store.search(
             group_identifier=str(request.project_id),
             prompt=request.prompt,
@@ -274,30 +274,6 @@ def Flow(self, request: pb2.FlowRequest, context) -> pb2.FlowResponse:
                 details=f"The specified model_identifier '{request.model_identifier}' does not exist. Please provide a valid model_identifier for flow generation."
             )
 
-        stored_functions_count = len(self.function_store.get_all(group_identifier=str(request.project_id)))
-        stored_flow_types_count = len(self.flow_type_store.get_all(group_identifier=str(request.project_id)))
-        stored_few_shots_count = len(self.few_shots_store.get_all(group_identifier="global"))
-        log.info(
-            f"[Flow] Store state — project={request.project_id} "
-            f"functions={stored_functions_count} "
-            f"flow_types={stored_flow_types_count} "
-            f"few_shots={stored_few_shots_count}"
-        )
-
-        if stored_functions_count <= 0 and (len(request.functions) <= 0):
-            log.warning(f"[Flow] Rejected — no functions for project={request.project_id}")
-            context.abort(
-                code=grpc.StatusCode.ABORTED,
-                details="No functions found for the given project_id. Please add functions before requesting a prompt generation."
-            )
-
-        if stored_flow_types_count <= 0 and (len(request.flow_types) <= 0):
-            log.warning(f"[Flow] Rejected — no flow types for project={request.project_id}")
-            context.abort(
-                code=grpc.StatusCode.ABORTED,
-                details="No flow types found for the given project_id. Please add flow_types before requesting a prompt generation."
-            )
-
         functions = [
             map_to_function_schema(fn)
             for fn in request.functions
@@ -327,6 +303,30 @@ def Flow(self, request: pb2.FlowRequest, context) -> pb2.FlowResponse:
                 data_types=data_types
             )
 
+        stored_functions_count = len(self.function_store.get_all(group_identifier=str(request.project_id)))
+        stored_flow_types_count = len(self.flow_type_store.get_all(group_identifier=str(request.project_id)))
+        stored_few_shots_count = len(self.few_shots_store.get_all(group_identifier="global"))
+        log.info(
+            f"[Flow] Store state — project={request.project_id} "
+            f"functions={stored_functions_count} "
+            f"flow_types={stored_flow_types_count} "
+            f"few_shots={stored_few_shots_count}"
+        )
+
+        if stored_functions_count <= 0:
+            log.warning(f"[Flow] Rejected — no functions for project={request.project_id}")
+            context.abort(
+                code=grpc.StatusCode.ABORTED,
+                details="No functions found for the given project_id. Please add functions before requesting a prompt generation."
+            )
+
+        if stored_flow_types_count <= 0:
+            log.warning(f"[Flow] Rejected — no flow types for project={request.project_id}")
+            context.abort(
+                code=grpc.StatusCode.ABORTED,
+                details="No flow types found for the given project_id. Please add flow_types before requesting a prompt generation."
+            )
+
         prompt_functions = self.function_store.search(
             group_identifier=str(request.project_id),
             prompt=request.prompt,