codeboarding 0.9.0__py3-none-any.whl

agents/abstraction_agent.py

@@ -0,0 +1,150 @@
+import logging
+from pathlib import Path
+
+from langchain_core.prompts import PromptTemplate
+from langchain_core.language_models import BaseChatModel
+
+from agents.agent import CodeBoardingAgent
+from agents.agent_responses import (
+    AnalysisInsights,
+    ClusterAnalysis,
+    MetaAnalysisInsights,
+    assign_component_ids,
+)
+from agents.prompts import (
+    get_system_message,
+    get_cluster_grouping_message,
+    get_final_analysis_message,
+)
+from agents.cluster_methods_mixin import ClusterMethodsMixin
+from agents.validation import (
+    ValidationContext,
+    validate_cluster_coverage,
+    validate_component_relationships,
+    validate_key_entities,
+    validate_cluster_ids_populated,
+    validate_relation_component_names,
+)
+from monitoring import trace
+from static_analyzer.analysis_result import StaticAnalysisResults
+from static_analyzer.graph import ClusterResult
+from static_analyzer.cluster_helpers import build_all_cluster_results, get_all_cluster_ids
+
+logger = logging.getLogger(__name__)
+
+
+class AbstractionAgent(ClusterMethodsMixin, CodeBoardingAgent):
+    def __init__(
+        self,
+        repo_dir: Path,
+        static_analysis: StaticAnalysisResults,
+        project_name: str,
+        meta_context: MetaAnalysisInsights,
+        agent_llm: BaseChatModel,
+        parsing_llm: BaseChatModel,
+    ):
+        super().__init__(repo_dir, static_analysis, get_system_message(), agent_llm, parsing_llm)
+
+        self.project_name = project_name
+        self.meta_context = meta_context
+
+        self.prompts = {
+            "group_clusters": PromptTemplate(
+                template=get_cluster_grouping_message(),
+                input_variables=["project_name", "cfg_clusters", "meta_context", "project_type"],
+            ),
+            "final_analysis": PromptTemplate(
+                template=get_final_analysis_message(),
+                input_variables=["project_name", "cluster_analysis", "meta_context", "project_type"],
+            ),
+        }
+
+    @trace
+    def step_clusters_grouping(self, cluster_results: dict[str, ClusterResult]) -> ClusterAnalysis:
+        logger.info(f"[AbstractionAgent] Grouping CFG clusters for: {self.project_name}")
+
+        meta_context_str = self.meta_context.llm_str() if self.meta_context else "No project context available."
+        project_type = self.meta_context.project_type if self.meta_context else "unknown"
+
+        programming_langs = self.static_analysis.get_languages()
+
+        # Build cluster string using the pre-computed cluster results
+        cluster_str = self._build_cluster_string(programming_langs, cluster_results)
+
+        prompt = self.prompts["group_clusters"].format(
+            project_name=self.project_name,
+            cfg_clusters=cluster_str,
+            meta_context=meta_context_str,
+            project_type=project_type,
+        )
+
+        cluster_analysis = self._validation_invoke(
+            prompt,
+            ClusterAnalysis,
+            validators=[validate_cluster_coverage],
+            context=ValidationContext(
+                cluster_results=cluster_results,
+                expected_cluster_ids=get_all_cluster_ids(cluster_results),
+            ),
+        )
+        return cluster_analysis
+
+    @trace
+    def step_final_analysis(
+        self, cluster_analysis: ClusterAnalysis, cluster_results: dict[str, ClusterResult]
+    ) -> AnalysisInsights:
+        logger.info(f"[AbstractionAgent] Generating final analysis for: {self.project_name}")
+
+        meta_context_str = self.meta_context.llm_str() if self.meta_context else "No project context available."
+        project_type = self.meta_context.project_type if self.meta_context else "unknown"
+
+        cluster_str = cluster_analysis.llm_str() if cluster_analysis else "No cluster analysis available."
+
+        prompt = self.prompts["final_analysis"].format(
+            project_name=self.project_name,
+            cluster_analysis=cluster_str,
+            meta_context=meta_context_str,
+            project_type=project_type,
+        )
+
+        # Build validation context with CFG graphs for edge checking
+        context = ValidationContext(
+            cluster_results=cluster_results,
+            cfg_graphs={lang: self.static_analysis.get_cfg(lang) for lang in self.static_analysis.get_languages()},
+        )
+
+        return self._validation_invoke(
+            prompt,
+            AnalysisInsights,
+            validators=[
+                validate_relation_component_names,
+                validate_component_relationships,
+                validate_key_entities,
+                validate_cluster_ids_populated,
+            ],
+            context=context,
+        )
+
+    def run(self):
+        # Build full cluster results dict for all languages ONCE
+        cluster_results = build_all_cluster_results(self.static_analysis)
+
+        # Step 1: Group related clusters together into logical components
+        cluster_analysis = self.step_clusters_grouping(cluster_results)
+
+        # Step 2: Generate abstract components from grouped clusters
+        analysis = self.step_final_analysis(cluster_analysis, cluster_results)
+        # Step 3: Sanitize cluster IDs (remove invalid ones)
+        self._sanitize_component_cluster_ids(analysis, cluster_results=cluster_results)
+        # Step 4: Assign files to components (deterministic + LLM-based with validation)
+        self.classify_files(analysis, cluster_results, self.static_analysis.get_all_source_files())
+        # Step 5: Fix source code reference lines (resolves reference_file paths for key_entities)
+        analysis = self.fix_source_code_reference_lines(analysis)
+        # Step 6: Ensure unique key entities across components
+        self._ensure_unique_key_entities(analysis)
+        # Step 7: Ensure unique file assignments across components
+        self._ensure_unique_file_assignments(analysis)
+        # Step 8: Assign deterministic component IDs
+        assign_component_ids(analysis)
+
+        return analysis, cluster_results

agents/agent.py

@@ -0,0 +1,467 @@
+import json
+import logging
+import os
+import time
+from pathlib import Path
+
+from google.api_core.exceptions import ResourceExhausted
+from langchain_core.exceptions import OutputParserException
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
+from langchain_core.output_parsers import PydanticOutputParser
+from langchain_core.prompts import PromptTemplate
+from langchain.agents import create_agent
+from langgraph.graph.state import CompiledStateGraph
+from pydantic import ValidationError
+from trustcall import create_extractor
+
+from agents.agent_responses import AnalysisInsights, ComponentFiles, FileClassification
+from agents.prompts import (
+    get_unassigned_files_classification_message,
+    get_validation_feedback_message,
+)
+from agents.tools.base import RepoContext
+from agents.tools.toolkit import CodeBoardingToolkit
+from agents.validation import ValidationContext, validate_file_classifications
+from monitoring.mixin import MonitoringMixin
+from repo_utils.ignore import RepoIgnoreManager
+from agents.llm_config import MONITORING_CALLBACK
+from static_analyzer.analysis_result import StaticAnalysisResults
+from static_analyzer.reference_resolve_mixin import ReferenceResolverMixin
+
+logger = logging.getLogger(__name__)
+
+
+class CodeBoardingAgent(ReferenceResolverMixin, MonitoringMixin):
+    def __init__(
+        self,
+        repo_dir: Path,
+        static_analysis: StaticAnalysisResults,
+        system_message: str,
+        agent_llm: BaseChatModel,
+        parsing_llm: BaseChatModel,
+    ):
+        ReferenceResolverMixin.__init__(self, repo_dir, static_analysis)
+        MonitoringMixin.__init__(self)
+        self.parsing_llm = parsing_llm
+        self.repo_dir = repo_dir
+        self.ignore_manager = RepoIgnoreManager(repo_dir)
+
+        # Initialize the professional toolkit
+        context = RepoContext(repo_dir=repo_dir, ignore_manager=self.ignore_manager, static_analysis=static_analysis)
+        self.toolkit = CodeBoardingToolkit(context=context)
+
+        self.agent: CompiledStateGraph = create_agent(
+            model=agent_llm,
+            tools=self.toolkit.get_agent_tools(),
+        )
+        self.static_analysis = static_analysis
+        self.system_message = SystemMessage(content=system_message)
+
+    @property
+    def read_source_reference(self):
+        return self.toolkit.read_source_reference
+
+    @property
+    def read_packages_tool(self):
+        return self.toolkit.read_packages
+
+    @property
+    def read_structure_tool(self):
+        return self.toolkit.read_structure
+
+    @property
+    def read_file_structure(self):
+        return self.toolkit.read_file_structure
+
+    @property
+    def read_cfg_tool(self):
+        return self.toolkit.read_cfg
+
+    @property
+    def read_method_invocations_tool(self):
+        return self.toolkit.read_method_invocations
+
+    @property
+    def read_file_tool(self):
+        return self.toolkit.read_file
+
+    @property
+    def read_docs(self):
+        return self.toolkit.read_docs
+
+    @property
+    def external_deps_tool(self):
+        return self.toolkit.external_deps
+
+    def _invoke(self, prompt, callbacks: list | None = None) -> str:
+        """Unified agent invocation method with timeout and exponential backoff.
+
+        Uses exponential backoff based on total attempts, with different multipliers
+        for different error types. This ensures backoff increases appropriately even
+        when errors alternate between types.
+        """
+        max_retries = 5
+
+        for attempt in range(max_retries):
+            timeout_seconds = 300 if attempt == 0 else 600
+            try:
+                callback_list = callbacks or []
+                # Always append monitoring callback - logging config controls output
+                callback_list.append(MONITORING_CALLBACK)
+                callback_list.append(self.agent_monitoring_callback)
+
+                logger.info(
+                    f"Starting agent.invoke() [attempt {attempt + 1}/{max_retries}] with prompt length: {len(prompt)}, timeout: {timeout_seconds}s"
+                )
+
+                response = self._invoke_with_timeout(
+                    timeout_seconds=timeout_seconds, callback_list=callback_list, prompt=prompt
+                )
+
+                logger.info(
+                    f"Completed agent.invoke() - message count: {len(response['messages'])}, last message type: {type(response['messages'][-1])}"
+                )
+
+                agent_response = response["messages"][-1]
+                assert isinstance(agent_response, AIMessage), f"Expected AIMessage, but got {type(agent_response)}"
+                if isinstance(agent_response.content, str):
+                    return agent_response.content
+                if isinstance(agent_response.content, list):
+                    return "".join(
+                        [
+                            str(message) if not isinstance(message, str) else message
+                            for message in agent_response.content
+                        ]
+                    )
+
+            except TimeoutError as e:
+                if attempt < max_retries - 1:
+                    # Exponential backoff: 10s * 2^attempt (10s, 20s, 40s, 80s)
+                    delay = min(10 * (2**attempt), 120)
+                    logger.warning(
+                        f"Agent invocation timed out after {timeout_seconds}s, retrying in {delay}s... (attempt {attempt + 1}/{max_retries})"
+                    )
+                    time.sleep(delay)
+                else:
+                    logger.error(f"Agent invocation timed out after {timeout_seconds}s on final attempt")
+                    raise
+
+            except ResourceExhausted as e:
+                if attempt < max_retries - 1:
+                    # Longer backoff for rate limits: 30s * 2^attempt (30s, 60s, 120s, 240s)
+                    delay = min(30 * (2**attempt), 300)
+                    logger.warning(
+                        f"ResourceExhausted (rate limit): {e}\n"
+                        f"Retrying in {delay}s... (attempt {attempt + 1}/{max_retries})"
+                    )
+                    time.sleep(delay)
+                else:
+                    logger.error(f"Max retries ({max_retries}) reached. ResourceExhausted: {e}")
+                    raise
+
+            except Exception as e:
+                # Other errors (network, parsing, etc.) get standard exponential backoff
+                if attempt < max_retries - 1:
+                    delay = min(10 * (2**attempt), 120)
+                    logger.warning(
+                        f"Agent error: {type(e).__name__}: {e}, retrying in {delay}s... (attempt {attempt + 1}/{max_retries})"
+                    )
+                    time.sleep(delay)
+                # On final attempt, fall through to return error message below
+
+        logger.error("Max retries reached. Failed to get response from the agent.")
+        return "Could not get response from the agent."
+
+    def _invoke_with_timeout(self, timeout_seconds: int, callback_list: list, prompt: str):
+        """Invoke agent with a timeout using threading."""
+        import threading
+        from queue import Queue, Empty
+
+        result_queue: Queue = Queue()
+        exception_queue: Queue = Queue()
+
+        def invoke_target():
+            try:
+                response = self.agent.invoke(
+                    {"messages": [self.system_message, HumanMessage(content=prompt)]},
+                    config={"callbacks": callback_list, "recursion_limit": 40},
+                )
+                result_queue.put(response)
+            except Exception as e:
+                exception_queue.put(e)
+
+        thread = threading.Thread(target=invoke_target, daemon=True)
+        thread.start()
+        thread.join(timeout=timeout_seconds)
+
+        if thread.is_alive():
+            # Thread is still running - timeout occurred
+            logger.error(f"Agent invoke thread still running after {timeout_seconds}s timeout")
+            raise TimeoutError(f"Agent invocation exceeded {timeout_seconds}s timeout")
+
+        # Check for exceptions
+        try:
+            exception = exception_queue.get_nowait()
+            raise exception
+        except Empty:
+            pass
+
+        # Get result
+        try:
+            return result_queue.get_nowait()
+        except Empty:
+            raise RuntimeError("Agent invocation completed but no result was returned")
+
+    def _parse_invoke(self, prompt, type):
+        response = self._invoke(prompt)
+        assert isinstance(response, str), f"Expected a string as response type got {response}"
+        return self._parse_response(prompt, response, type)
+
+    def _validation_invoke(
+        self, prompt: str, return_type: type, validators: list, context, max_validation_retries: int = 1
+    ):
+        """
+        Invoke LLM with validation and feedback loop.
+
+        Args:
+            prompt: The original prompt
+            return_type: Pydantic type to parse into
+            validators: List of validation functions to run
+            context: ValidationContext with data needed for validation
+            max_validation_retries: Maximum retry attempts with feedback (default: 1)
+
+        Returns:
+            Validated result of return_type
+        """
+        result = self._parse_invoke(prompt, return_type)
+
+        for attempt in range(max_validation_retries):
+            # Run all validators
+            all_feedback = []
+            for validator in validators:
+                validation_result = validator(result, context)
+                if not validation_result.is_valid:
+                    all_feedback.extend(validation_result.feedback_messages)
+
+            if not all_feedback:
+                logger.info(f"[Validation] All validations passed on attempt {attempt + 1}")
+                return result  # All validations passed
+
+            # Build feedback prompt using the prompt factory
+            feedback_template = get_validation_feedback_message()
+            feedback_prompt = feedback_template.format(
+                original_output=result.llm_str(),
+                feedback_list="\n".join(f"- {msg}" for msg in all_feedback),
+                original_prompt=prompt,
+            )
+
+            logger.info(
+                f"[Validation] Retry {attempt + 1}/{max_validation_retries} with {len(all_feedback)} feedback items"
+            )
+            result = self._parse_invoke(feedback_prompt, return_type)
+
+        return result
+
+    def _parse_response(self, prompt, response, return_type, max_retries=5, attempt=0):
+        if attempt >= max_retries:
+            logger.error(f"Max retries ({max_retries}) reached for parsing response: {response}")
+            raise Exception(f"Max retries reached for parsing response: {response}")
+
+        extractor = create_extractor(self.parsing_llm, tools=[return_type], tool_choice=return_type.__name__)
+        if response is None or response.strip() == "":
+            logger.error(f"Empty response for prompt: {prompt}")
+        try:
+            result = extractor.invoke(
+                return_type.extractor_str() + response,
+                config={"callbacks": [MONITORING_CALLBACK, self.agent_monitoring_callback]},
+            )
+            if "responses" in result and len(result["responses"]) != 0:
+                return return_type.model_validate(result["responses"][0])
+            if "messages" in result and len(result["messages"]) != 0:
+                message = result["messages"][0].content
+                parser = PydanticOutputParser(pydantic_object=return_type)
+                return self._try_parse(message, parser)
+            parser = PydanticOutputParser(pydantic_object=return_type)
+            return self._try_parse(response, parser)
+        except AttributeError as e:
+            # Workaround for trustcall bug: https://github.com/hinthornw/trustcall/issues/47
+            # 'ExtractionState' object has no attribute 'tool_call_id' occurs during validation retry
+            if "tool_call_id" in str(e):
+                logger.warning(f"Trustcall bug encountered, falling back to Pydantic parser: {e}")
+                parser = PydanticOutputParser(pydantic_object=return_type)
+                return self._try_parse(response, parser)
+            raise
+        except IndexError as e:
+            # try to parse with the json parser if possible
+            logger.warning(f"IndexError while parsing response (attempt {attempt + 1}/{max_retries}): {e}")
+            return self._parse_response(prompt, response, return_type, max_retries, attempt + 1)
+        except ResourceExhausted as e:
+            # Parsing uses exponential backoff for rate limits
+            if attempt < max_retries - 1:
+                # Exponential backoff: 30s * 2^attempt, capped at 300s
+                delay = min(30 * (2**attempt), 300)
+                logger.warning(
+                    f"ResourceExhausted during parsing (rate limit): {e}\n"
+                    f"Retrying in {delay}s... (attempt {attempt + 1}/{max_retries})"
+                )
+                time.sleep(delay)
+                return self._parse_response(prompt, response, return_type, max_retries, attempt + 1)
+            else:
+                logger.error(f"Resource exhausted on final parsing attempt: {e}")
+                raise
+
+    def _try_parse(self, message_content, parser):
+        try:
+            prompt_template = """You are an JSON expert. Here you need to extract information in the following json format: {format_instructions}
+
+            Here is the content to parse and fix: {adjective}
+
+            Please provide only the JSON output without any additional text."""
+            prompt = PromptTemplate(
+                template=prompt_template,
+                input_variables=["adjective"],
+                partial_variables={"format_instructions": parser.get_format_instructions()},
+            )
+            chain = prompt | self.parsing_llm | parser
+            return chain.invoke(
+                {"adjective": message_content},
+                config={"callbacks": [MONITORING_CALLBACK, self.agent_monitoring_callback]},
+            )
+        except (ValidationError, OutputParserException):
+            for _, v in json.loads(message_content).items():
+                try:
+                    return self._try_parse(json.dumps(v), parser)
+                except:
+                    pass
+        raise ValueError(f"Couldn't parse {message_content}")
+
+    def classify_files(self, analysis: AnalysisInsights, cluster_results: dict, scope_files: list[str]) -> None:
+        """
+        Two-pass file assignment for AnalysisInsights:
+        1. Deterministic: assign files from cluster_ids and key_entities
+        2. LLM-based: classify remaining unassigned files
+
+        Args:
+            analysis: AnalysisInsights object to classify files for
+            cluster_results: Dict mapping language -> ClusterResult (for the relevant scope)
+            scope_files: List of file paths to limit classification scope.
+
+        Requires self to be a mixin with ClusterMethodsMixin for helper methods.
+        """
+        for comp in analysis.components:
+            # Deterministic assignment (uses mixin methods)
+            self._assign_files_to_component(comp, cluster_results)  # type: ignore[attr-defined]
+        self._classify_unassigned_files_llm(analysis, scope_files)
+        self._log_unclassified_files_count(analysis, scope_files)
+
+    def _classify_unassigned_files_llm(self, analysis: AnalysisInsights, scope_files: list[str]) -> None:
+        """
+        Classify files from the scope files that weren't assigned to any component.
+        Uses a single LLM call to classify all unassigned files.
+        Args:
+            analysis: AnalysisInsights object
+            scope_files: List of file paths to limit classification scope.
+        """
+        # Get unassigned files using the helper method
+        unassigned_files = self._get_unassigned_files(analysis, scope_files)
+
+        if not unassigned_files:
+            logger.info("[Agent] All files already assigned, skipping LLM classification")
+            return
+
+        logger.info(f"[Agent] Found {len(unassigned_files)} unassigned files, using LLM classification")
+
+        # 4. Build component summary for LLM using llm_str()
+        valid_components = [comp for comp in analysis.components if comp.name != "Unclassified"]
+        components_summary = "\n\n".join(comp.llm_str() for comp in valid_components)
+        component_map = {comp.name: comp for comp in valid_components}
+
+        # 5. Classify all unassigned files with LLM
+        classifications: list[FileClassification] = self._classify_unassigned_files_with_llm(
+            unassigned_files, components_summary, analysis
+        )
+
+        # 6. Append successfully classified files to components
+        for fc in classifications:
+            if fc.component_name in component_map:
+                comp = component_map[fc.component_name]
+                if fc.file_path not in comp.assigned_files:
+                    comp.assigned_files.append(fc.file_path)
+                    logger.debug(f"[Agent] Assigned {fc.file_path} to {fc.component_name}")
+            else:
+                logger.warning(
+                    f"[Agent] Invalid component name '{fc.component_name}' for file {fc.file_path}, skipping"
+                )
+
+        logger.info(f"[Agent] File classification complete: {len(classifications)} files classified")
+
+    def _get_unassigned_files(self, analysis: AnalysisInsights, scope_files: list[str]) -> list[str]:
+        """
+        Check which files remain unassigned after classification.
+        Args:
+            analysis: AnalysisInsights object with classified components
+            scope_files: List of file paths to limit the scope.
+        Returns:
+            List of file paths that are still unassigned
+        """
+        # 1. Gather all assigned files
+        assigned_files = set()
+        for comp in analysis.components:
+            for f in comp.assigned_files:
+                abs_path = os.path.join(self.repo_dir, f) if not os.path.isabs(f) else f
+                assigned_files.add(os.path.relpath(abs_path, self.repo_dir))
+
+        # 2. Get files to consider for classification
+        # If scope_files is provided (e.g., DetailsAgent), use those
+        # Otherwise use all source files from static_analysis (e.g., AbstractionAgent)
+        all_files = set()
+        for file_path in scope_files:
+            file_path_str = str(file_path)
+            rel_path = os.path.relpath(file_path_str, self.repo_dir) if os.path.isabs(file_path_str) else file_path_str
+            all_files.add(rel_path)
+
+        # 3. Return unassigned files
+        return sorted(all_files - assigned_files)
+
+    def _log_unclassified_files_count(self, analysis: AnalysisInsights, scope_files: list[str]) -> None:
+        """
+        Log how many files remain unclassified within the analysis.
+
+        Args:
+            analysis: AnalysisInsights object with classified components
+            scope_files: List of file paths which are expected to be within the analysis.
+        """
+        unassigned = self._get_unassigned_files(analysis, scope_files)
+        if unassigned:
+            logger.warning(f"[Agent] {len(unassigned)} files have not been classified successfully: {unassigned}")
+        else:
+            logger.info("[Agent] All files have been classified successfully")
+
+    def _classify_unassigned_files_with_llm(
+        self, unassigned_files: list[str], components_summary: str, analysis: AnalysisInsights
+    ) -> list[FileClassification]:
+        """
+        Classify unassigned files using LLM with validation.
+        Returns list of FileClassification objects.
+        """
+
+        prompt = PromptTemplate(
+            template=get_unassigned_files_classification_message(), input_variables=["unassigned_files", "components"]
+        ).format(unassigned_files="\n".join(unassigned_files), components=components_summary)
+
+        # Get valid component names from the components_summary
+        # Parse component names from the summary (components have format "**Component:** `ComponentName`")
+        valid_component_names = set([comp.name for comp in analysis.components])
+
+        # Build validation context
+        context = ValidationContext(
+            expected_files=set(unassigned_files),
+            valid_component_names=valid_component_names,
+            repo_dir=str(self.repo_dir),
+        )
+
+        file_classifications = self._validation_invoke(
+            prompt, ComponentFiles, validators=[validate_file_classifications], context=context
+        )
+        return file_classifications.file_paths