added return only one query context

Author: tarnopolskyi

uncoder-core/app/routers/translate.py

@@ -1,6 +1,7 @@
 from fastapi import APIRouter, Body
 
 from app.models.translation import InfoMessage, OneTranslationData, Platform, TranslatorPlatforms
+from app.translator.core.context_vars import return_only_first_query_ctx_var
 from app.translator.cti_translator import CTITranslator
 from app.translator.translator import Translator
 
@@ -14,8 +15,9 @@
 def translate_one(
     source_platform_id: str = Body(..., embed=True),
     target_platform_id: str = Body(..., embed=True),
-    text: str = Body(..., embed=True),
+    return_only_first_query: bool = False,
 ) -> OneTranslationData:
+    return_only_first_query_ctx_var.set(return_only_first_query)
     status, data = translator.translate_one(text=text, source=source_platform_id, target=target_platform_id)
     if status:
         return OneTranslationData(status=status, translation=data, target_platform_id=target_platform_id)
@@ -27,8 +29,11 @@ def translate_one(
 @st_router.post("/translate/all", tags=["translator"], description="Generate all translations")
 @st_router.post("/translate/all/", include_in_schema=False)
 def translate_all(
-    source_platform_id: str = Body(..., embed=True), text: str = Body(..., embed=True)
+    source_platform_id: str = Body(..., embed=True),
+    text: str = Body(..., embed=True),
+    return_only_first_query: bool = False,
 ) -> list[OneTranslationData]:
+    return_only_first_query_ctx_var.set(return_only_first_query)
     result = translator.translate_all(text=text, source=source_platform_id)
     translations = []
     for platform_result in result:

uncoder-core/app/translator/core/context_vars.py

@@ -0,0 +1,4 @@
+from contextvars import ContextVar
+
+return_only_first_query_ctx_var: ContextVar[bool] = ContextVar("return_only_first_query_ctx_var", default=False)
+"""Set to True to return ony first query if rendered multiple options"""

uncoder-core/app/translator/core/render.py

@@ -35,6 +35,7 @@
 from app.translator.core.models.query_container import MetaInfoContainer, RawQueryContainer, TokenizedQueryContainer
 from app.translator.core.str_value_manager import StrValue, StrValueManager
 from app.translator.core.tokenizer import TOKEN_TYPE
+from app.translator.core.context_vars import return_only_first_query_ctx_var
 
 
 class BaseQueryFieldValue(ABC):
@@ -283,6 +284,7 @@ def finalize_query(
         **kwargs,  # noqa: ARG002
     ) -> str:
         query = self.query_pattern.format(prefix=prefix, query=query, functions=functions).strip()
+
         query = self.wrap_query_with_meta_info(meta_info=meta_info, query=query)
         if not_supported_functions:
             rendered_not_supported = self.render_not_supported_functions(not_supported_functions)
@@ -368,6 +370,8 @@ def _generate_from_tokenized_query_container(self, query_container: TokenizedQue
                     meta_info=query_container.meta_info,
                     source_mapping=source_mapping,
                 )
+                if return_only_first_query_ctx_var.get() is True:
+                    return finalized_query
                 queries_map[source_mapping.source_id] = finalized_query
         if not queries_map and errors:
             raise errors[0]

uncoder-core/app/translator/platforms/forti_siem/renders/forti_siem_rule.py

@@ -18,6 +18,7 @@
 from typing import Optional, Union
 
 from app.translator.const import DEFAULT_VALUE_TYPE
+from app.translator.core.context_vars import return_only_first_query_ctx_var
 from app.translator.core.custom_types.meta_info import SeverityType
 from app.translator.core.custom_types.tokens import GroupType, LogicalOperatorType, OperatorType
 from app.translator.core.custom_types.values import ValueType
@@ -273,6 +274,8 @@ def _generate_from_tokenized_query_container(self, query_container: TokenizedQue
                 source_mapping=source_mapping,
                 fields=mapped_fields_set,
             )
+            if return_only_first_query_ctx_var.get() is True:
+                return finalized_query
             queries_map[source_mapping.source_id] = finalized_query
 
         return self.finalize(queries_map)

uncoder-core/app/translator/platforms/logrhythm_axon/renders/logrhythm_axon_query.py

@@ -19,6 +19,7 @@
 from typing import Union
 
 from app.translator.const import DEFAULT_VALUE_TYPE
+from app.translator.core.context_vars import return_only_first_query_ctx_var
 from app.translator.core.custom_types.tokens import LogicalOperatorType
 from app.translator.core.custom_types.values import ValueType
 from app.translator.core.exceptions.core import StrictPlatformException
@@ -262,6 +263,8 @@ def _generate_from_tokenized_query_container(self, query_container: TokenizedQue
                 meta_info=query_container.meta_info,
                 source_mapping=source_mapping,
             )
+            if return_only_first_query_ctx_var.get() is True:
+                return finalized_query
             queries_map[source_mapping.source_id] = finalized_query
 
         return self.finalize(queries_map)