fix: few needed adjustments

Signed-off-by: Zvi Grinberg <[email protected]>

Author: zvigrinberg

src/vuln_analysis/functions/cve_agent.py

@@ -69,20 +69,28 @@ async def _create_agent(config: CVEAgentExecutorToolConfig, builder: Builder,
     tools = builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
     llm = await builder.get_llm(llm_name=config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
     prompt = PromptTemplate.from_template(get_agent_prompt(config.prompt, config.prompt_examples))
-
+    ecosystem = state.original_input.input.image.ecosystem
+    transitive_search_tool_supports_ecosystem: bool = True
+    # I
+    if ecosystem:
+        try:
+            get_language_function_parser(ecosystem)
+        except NotImplementedError:
+            transitive_search_tool_supports_ecosystem = False
+            logger.warning(f"Transitive code search tool doesn't support programming language {ecosystem},"
+                           f" disabling tool...")
     # Filter tools that are not available
     tools = [
         tool for tool in tools
         if not ((tool.name == "Container Image Code QA System" and state.code_vdb_path is None) or
                 (tool.name == "Container Image Developer Guide QA System" and state.doc_vdb_path is None) or
                 (tool.name == "Lexical Search Container Image Code QA System" and state.code_index_path is None) or
                 (tool.name == "Transitive code search tool" and (not config.transitive_search_tool_enabled or
-                                                                 state.code_index_path is None)) or
+                                                                 state.code_index_path is None or
+                                                                 not transitive_search_tool_supports_ecosystem)) or
                 (tool.name == "Calling Function Name Extractor" and (not config.transitive_search_tool_enabled or
-                                                                     state.code_index_path is None)) or
-                (tool.name == "Package and Function Locator" and (not config.transitive_search_tool_enabled or
-                                                                  state.code_index_path is None))
-
+                                                                     state.code_index_path is None or
+                                                                     not transitive_search_tool_supports_ecosystem))
                 )
     ]
 

src/vuln_analysis/functions/cve_generate_vdbs.py

@@ -28,6 +28,7 @@
 
 from vuln_analysis.data_models.common import AnalysisType
 from vuln_analysis.logging.loggers_factory import LoggingFactory, trace_id
+from vuln_analysis.utils.dep_tree import Ecosystem
 
 logger = LoggingFactory.get_agent_logger(__name__)
 
@@ -98,7 +99,7 @@ async def generate_vdb(config: CVEGenerateVDBsToolConfig, builder: Builder):
                                  pickle_cache_directory=config.base_pickle_dir)
 
     def _create_code_index(source_infos: list[SourceDocumentsInfo], embedder: DocumentEmbedding,
-                           output_path: Path) -> bool :
+                           output_path: Path) -> bool:
         logger.info("Collecting documents from git repos. Source Infos: %s",
                     json.dumps([x.model_dump(mode="json") for x in source_infos]))
 
@@ -135,7 +136,8 @@ def _create_code_index(source_infos: list[SourceDocumentsInfo], embedder: Docume
         logger.info("Completed code indexing in %.2f seconds for '%s'", time.time() - indexing_start_time, output_path)
         return True
 
-    def _build_code_index(source_infos: list[SourceDocumentsInfo], ecosystem, manifest_path) -> Path | None:
+    def _build_code_index(source_infos: list[SourceDocumentsInfo], ecosystem: Ecosystem = None,
+                          manifest_path: str = None) -> Path | None:
         code_index_path: Path | None = None
 
         # Filter to only code sources

src/vuln_analysis/tools/transitive_code_search.py

@@ -37,6 +37,8 @@
 logger = LoggingFactory.get_agent_logger(__name__)
 
 
+
+
 class TransitiveCodeSearchToolConfig(FunctionBaseConfig, name="transitive_code_search"):
     """
     Transitive code search tool used to search source code.

src/vuln_analysis/utils/functions_parsers/lang_functions_parsers_factory.py

@@ -11,8 +11,7 @@ def get_language_function_parser(ecosystem: Ecosystem, tree: DependencyTree | No
     :param ecosystem: the desired programming language parser.
     :param tree: the dependency tree for the ecosystem, can be None.
     :return:
-    The right language functions parser associated to the ecosystem, if not exists, return ABC parent that
-     doesn't do anything
+    The right language functions parser associated to the ecosystem, if not exists, throw an NotImplementedError
     """
     if ecosystem == Ecosystem.GO:
         return GoLanguageFunctionsParser()
@@ -24,4 +23,4 @@ def get_language_function_parser(ecosystem: Ecosystem, tree: DependencyTree | No
             parser_obj.init_CParser(tree)
         return parser_obj
     else:
-        return LanguageFunctionsParser()
+        raise NotImplementedError(f"Language functions parser for {ecosystem} not implemented.")